kimi-cli 0.35__py3-none-any.whl → 0.52__py3-none-any.whl

kimi_cli/tools/file/read.py

@@ -1,11 +1,11 @@
 from pathlib import Path
-from typing import override
+from typing import Any, override
 
 import aiofiles
 from kosong.tooling import CallableTool2, ToolError, ToolOk, ToolReturnType
 from pydantic import BaseModel, Field
 
-from kimi_cli.agent import BuiltinSystemPromptArgs
+from kimi_cli.soul.runtime import BuiltinSystemPromptArgs
 from kimi_cli.tools.utils import load_desc, truncate_line
 
 MAX_LINES = 1000
@@ -47,8 +47,9 @@ class ReadFile(CallableTool2[Params]):
     )
     params: type[Params] = Params
 
-    def __init__(self, builtin_args: BuiltinSystemPromptArgs, **kwargs):
+    def __init__(self, builtin_args: BuiltinSystemPromptArgs, **kwargs: Any) -> None:
         super().__init__(**kwargs)
+
         self._work_dir = builtin_args.KIMI_WORK_DIR
 
     @override
@@ -84,7 +85,7 @@ class ReadFile(CallableTool2[Params]):
 
             lines: list[str] = []
             n_bytes = 0
-            truncated_line_numbers = []
+            truncated_line_numbers: list[int] = []
             max_lines_reached = False
             max_bytes_reached = False
             async with aiofiles.open(p, encoding="utf-8", errors="replace") as f:
@@ -108,7 +109,7 @@ class ReadFile(CallableTool2[Params]):
                         break
 
             # Format output with line numbers like `cat -n`
-            lines_with_no = []
+            lines_with_no: list[str] = []
             for line_num, line in zip(
                 range(params.line_offset, params.line_offset + len(lines)), lines, strict=True
             ):

kimi_cli/tools/file/replace.py

@@ -1,11 +1,14 @@
 from pathlib import Path
-from typing import override
+from typing import Any, override
 
 import aiofiles
 from kosong.tooling import CallableTool2, ToolError, ToolOk, ToolReturnType
 from pydantic import BaseModel, Field
 
-from kimi_cli.agent import BuiltinSystemPromptArgs
+from kimi_cli.soul.approval import Approval
+from kimi_cli.soul.runtime import BuiltinSystemPromptArgs
+from kimi_cli.tools.file import FileActions
+from kimi_cli.tools.utils import ToolRejectedError, load_desc
 
 
 class Edit(BaseModel):
@@ -26,12 +29,13 @@ class Params(BaseModel):
 
 class StrReplaceFile(CallableTool2[Params]):
     name: str = "StrReplaceFile"
-    description: str = (Path(__file__).parent / "replace.md").read_text()
+    description: str = load_desc(Path(__file__).parent / "replace.md")
     params: type[Params] = Params
 
-    def __init__(self, builtin_args: BuiltinSystemPromptArgs, **kwargs):
+    def __init__(self, builtin_args: BuiltinSystemPromptArgs, approval: Approval, **kwargs: Any):
         super().__init__(**kwargs)
         self._work_dir = builtin_args.KIMI_WORK_DIR
+        self._approval = approval
 
     def _validate_path(self, path: Path) -> ToolError | None:
         """Validate that the path is safe to edit."""
@@ -87,6 +91,14 @@ class StrReplaceFile(CallableTool2[Params]):
                     brief="Invalid path",
                 )
 
+            # Request approval
+            if not await self._approval.request(
+                self.name,
+                FileActions.EDIT,
+                f"Edit file `{params.path}`",
+            ):
+                return ToolRejectedError()
+
             # Read the file content
             async with aiofiles.open(p, encoding="utf-8", errors="replace") as f:
                 content = await f.read()

kimi_cli/tools/file/write.py

@@ -1,11 +1,14 @@
 from pathlib import Path
-from typing import Literal, override
+from typing import Any, Literal, override
 
 import aiofiles
 from kosong.tooling import CallableTool2, ToolError, ToolOk, ToolReturnType
 from pydantic import BaseModel, Field
 
-from kimi_cli.agent import BuiltinSystemPromptArgs
+from kimi_cli.soul.approval import Approval
+from kimi_cli.soul.runtime import BuiltinSystemPromptArgs
+from kimi_cli.tools.file import FileActions
+from kimi_cli.tools.utils import ToolRejectedError, load_desc
 
 
 class Params(BaseModel):
@@ -23,12 +26,13 @@ class Params(BaseModel):
 
 class WriteFile(CallableTool2[Params]):
     name: str = "WriteFile"
-    description: str = (Path(__file__).parent / "write.md").read_text()
+    description: str = load_desc(Path(__file__).parent / "write.md")
     params: type[Params] = Params
 
-    def __init__(self, builtin_args: BuiltinSystemPromptArgs, **kwargs):
+    def __init__(self, builtin_args: BuiltinSystemPromptArgs, approval: Approval, **kwargs: Any):
         super().__init__(**kwargs)
         self._work_dir = builtin_args.KIMI_WORK_DIR
+        self._approval = approval
 
     def _validate_path(self, path: Path) -> ToolError | None:
         """Validate that the path is safe to write."""
@@ -85,6 +89,14 @@ class WriteFile(CallableTool2[Params]):
                     brief="Invalid write mode",
                 )
 
+            # Request approval
+            if not await self._approval.request(
+                self.name,
+                FileActions.EDIT,
+                f"Write file `{params.path}`",
+            ):
+                return ToolRejectedError()
+
             # Determine file mode for aiofiles
             file_mode = "w" if params.mode == "overwrite" else "a"
 

kimi_cli/tools/mcp.py

@@ -1,12 +1,15 @@
+from typing import Any
+
 import fastmcp
 import mcp
 from fastmcp.client.client import CallToolResult
-from kosong.base.message import AudioURLPart, ContentPart, ImageURLPart, TextPart
+from fastmcp.client.transports import ClientTransport
+from kosong.message import AudioURLPart, ContentPart, ImageURLPart, TextPart
 from kosong.tooling import CallableTool, ToolOk, ToolReturnType
 
 
-class MCPTool(CallableTool):
-    def __init__(self, mcp_tool: mcp.Tool, client: fastmcp.Client, **kwargs):
+class MCPTool[T: ClientTransport](CallableTool):
+    def __init__(self, mcp_tool: mcp.Tool, client: fastmcp.Client[T], **kwargs: Any):
         super().__init__(
             name=mcp_tool.name,
             description=mcp_tool.description or "",
@@ -16,7 +19,7 @@ class MCPTool(CallableTool):
         self._mcp_tool = mcp_tool
         self._client = client
 
-    async def __call__(self, *args, **kwargs) -> ToolReturnType:
+    async def __call__(self, *args: Any, **kwargs: Any) -> ToolReturnType:
         async with self._client as client:
             result = await client.call_tool(self._mcp_tool.name, kwargs, timeout=20)
             return convert_tool_result(result)

kimi_cli/tools/task/__init__.py

@@ -1,17 +1,22 @@
+import asyncio
 from pathlib import Path
-from typing import override
+from typing import Any, override
 
 from kosong.tooling import CallableTool2, ToolError, ToolOk, ToolReturnType
 from pydantic import BaseModel, Field
 
-from kimi_cli.agent import Agent, AgentGlobals, AgentSpec, load_agent
-from kimi_cli.soul import MaxStepsReached
+from kimi_cli.agentspec import ResolvedAgentSpec, SubagentSpec
+from kimi_cli.soul import MaxStepsReached, get_wire_or_none, run_soul
+from kimi_cli.soul.agent import Agent, load_agent
 from kimi_cli.soul.context import Context
 from kimi_cli.soul.kimisoul import KimiSoul
-from kimi_cli.soul.wire import ApprovalRequest, Wire, WireMessage, get_wire_or_none
+from kimi_cli.soul.runtime import Runtime
+from kimi_cli.soul.toolset import get_current_tool_call_or_none
 from kimi_cli.tools.utils import load_desc
 from kimi_cli.utils.message import message_extract_text
 from kimi_cli.utils.path import next_available_rotation
+from kimi_cli.wire import WireUISide
+from kimi_cli.wire.message import ApprovalRequest, SubagentEvent, WireMessage
 
 # Maximum continuation attempts for task summary
 MAX_CONTINUE_ATTEMPTS = 1
@@ -45,29 +50,37 @@ class Task(CallableTool2[Params]):
     name: str = "Task"
     params: type[Params] = Params
 
-    def __init__(self, agent_spec: AgentSpec, agent_globals: AgentGlobals, **kwargs):
-        subagents: dict[str, Agent] = {}
-        descs = []
-
-        # load all subagents
-        assert agent_spec.subagents is not None, "Task tool expects subagents"
-        for name, spec in agent_spec.subagents.items():
-            subagents[name] = load_agent(spec.path, agent_globals)
-            descs.append(f"- `{name}`: {spec.description}")
-
+    def __init__(self, agent_spec: ResolvedAgentSpec, runtime: Runtime, **kwargs: Any):
         super().__init__(
             description=load_desc(
                 Path(__file__).parent / "task.md",
                 {
-                    "SUBAGENTS_MD": "\n".join(descs),
+                    "SUBAGENTS_MD": "\n".join(
+                        f"- `{name}`: {spec.description}"
+                        for name, spec in agent_spec.subagents.items()
+                    ),
                 },
             ),
             **kwargs,
         )
 
-        self._agent_globals = agent_globals
-        self._session = agent_globals.session
-        self._subagents = subagents
+        self._runtime = runtime
+        self._session = runtime.session
+        self._subagents: dict[str, Agent] = {}
+
+        try:
+            loop = asyncio.get_running_loop()
+            self._load_task = loop.create_task(self._load_subagents(agent_spec.subagents))
+        except RuntimeError:
+            # In case there's no running event loop, e.g., during synchronous tests
+            self._load_task = None
+            asyncio.run(self._load_subagents(agent_spec.subagents))
+
+    async def _load_subagents(self, subagent_specs: dict[str, SubagentSpec]) -> None:
+        """Load all subagents specified in the agent spec."""
+        for name, spec in subagent_specs.items():
+            agent = await load_agent(spec.path, self._runtime, mcp_configs=[])
+            self._subagents[name] = agent
 
     async def _get_subagent_history_file(self) -> Path:
         """Generate a unique history file path for subagent."""
@@ -82,6 +95,10 @@ class Task(CallableTool2[Params]):
 
     @override
     async def __call__(self, params: Params) -> ToolReturnType:
+        if self._load_task is not None:
+            await self._load_task
+            self._load_task = None
+
         if params.subagent_name not in self._subagents:
             return ToolError(
                 message=f"Subagent not found: {params.subagent_name}",
@@ -99,20 +116,34 @@ class Task(CallableTool2[Params]):
 
     async def _run_subagent(self, agent: Agent, prompt: str) -> ToolReturnType:
         """Run subagent with optional continuation for task summary."""
+        super_wire = get_wire_or_none()
+        assert super_wire is not None
+        current_tool_call = get_current_tool_call_or_none()
+        assert current_tool_call is not None
+        current_tool_call_id = current_tool_call.id
+
+        def _super_wire_send(msg: WireMessage) -> None:
+            if isinstance(msg, ApprovalRequest):
+                super_wire.soul_side.send(msg)
+                return
+
+            event = SubagentEvent(
+                task_tool_call_id=current_tool_call_id,
+                event=msg,
+            )
+            super_wire.soul_side.send(event)
+
+        async def _ui_loop_fn(wire: WireUISide) -> None:
+            while True:
+                msg = await wire.receive()
+                _super_wire_send(msg)
+
         subagent_history_file = await self._get_subagent_history_file()
         context = Context(file_backend=subagent_history_file)
-        soul = KimiSoul(
-            agent,
-            agent_globals=self._agent_globals,
-            context=context,
-            loop_control=self._agent_globals.config.loop_control,
-        )
-        wire = get_wire_or_none()
-        assert wire is not None, "Wire is expected to be set"
-        sub_wire = _SubWire(wire)
+        soul = KimiSoul(agent, runtime=self._runtime, context=context)
 
         try:
-            await soul.run(prompt, sub_wire)
+            await run_soul(soul, prompt, _ui_loop_fn, asyncio.Event())
         except MaxStepsReached as e:
             return ToolError(
                 message=(
@@ -135,22 +166,10 @@ class Task(CallableTool2[Params]):
         # Check if response is too brief, if so, run again with continuation prompt
         n_attempts_remaining = MAX_CONTINUE_ATTEMPTS
         if len(final_response) < 200 and n_attempts_remaining > 0:
-            await soul.run(CONTINUE_PROMPT, sub_wire)
+            await run_soul(soul, CONTINUE_PROMPT, _ui_loop_fn, asyncio.Event())
 
             if len(context.history) == 0 or context.history[-1].role != "assistant":
                 return ToolError(message=_error_msg, brief="Failed to run subagent")
             final_response = message_extract_text(context.history[-1])
 
         return ToolOk(output=final_response)
-
-
-class _SubWire(Wire):
-    def __init__(self, super_wire: Wire):
-        super().__init__()
-        self._super_wire = super_wire
-
-    @override
-    def send(self, msg: WireMessage):
-        if isinstance(msg, ApprovalRequest):
-            self._super_wire.send(msg)
-        # TODO: visualize subagent behavior by sending other messages in some way

kimi_cli/tools/task/task.md

@@ -4,7 +4,7 @@ Spawn a subagent to perform a specific task. Subagent will be spawned with a fre
 
 Context isolation is one of the key benefits of using subagents. By delegating tasks to subagents, you can keep your main context clean and focused on the main goal requested by the user.
 
-Here are some scenerios you may want this tool for context isolation:
+Here are some scenarios you may want this tool for context isolation:
 
 - You wrote some code and it did not work as expected. In this case you can spawn a subagent to fix the code, asking the subagent to return how it is fixed. This can potentially benefit because the detailed process of fixing the code may not be relevant to your main goal, and may clutter your context.
 - When you need some latest knowledge of a specific library, framework or technology to proceed with your task, you can spawn a subagent to search on the internet for the needed information and return to you the gathered relevant information, for example code examples, API references, etc. This can avoid ton of irrelevant search results in your own context.

kimi_cli/tools/todo/__init__.py

@@ -4,6 +4,8 @@ from typing import Literal, override
 from kosong.tooling import CallableTool2, ToolOk, ToolReturnType
 from pydantic import BaseModel, Field
 
+from kimi_cli.tools.utils import load_desc
+
 
 class Todo(BaseModel):
     title: str = Field(description="The title of the todo", min_length=1)
@@ -16,7 +18,7 @@ class Params(BaseModel):
 
 class SetTodoList(CallableTool2[Params]):
     name: str = "SetTodoList"
-    description: str = (Path(__file__).parent / "set_todo_list.md").read_text()
+    description: str = load_desc(Path(__file__).parent / "set_todo_list.md")
     params: type[Params] = Params
 
     @override
@@ -24,4 +26,4 @@ class SetTodoList(CallableTool2[Params]):
         rendered = ""
         for todo in params.todos:
             rendered += f"- {todo.title} [{todo.status}]\n"
-        return ToolOk(output=rendered)
+        return ToolOk(output="", message="Todo list updated", brief=rendered)

kimi_cli/tools/utils.py

@@ -7,7 +7,7 @@ from kosong.tooling import ToolError, ToolOk
 
 def load_desc(path: Path, substitutions: dict[str, str] | None = None) -> str:
     """Load a tool description from a file, with optional substitutions."""
-    description = path.read_text()
+    description = path.read_text(encoding="utf-8")
     if substitutions:
         description = string.Template(description).substitute(substitutions)
     return description

kimi_cli/tools/web/fetch.py

@@ -7,6 +7,7 @@ from kosong.tooling import CallableTool2, ToolReturnType
 from pydantic import BaseModel, Field
 
 from kimi_cli.tools.utils import ToolResultBuilder, load_desc
+from kimi_cli.utils.aiohttp import new_client_session
 
 
 class Params(BaseModel):
@@ -24,7 +25,7 @@ class FetchURL(CallableTool2[Params]):
 
         try:
             async with (
-                aiohttp.ClientSession() as session,
+                new_client_session() as session,
                 session.get(
                     params.url,
                     headers={

kimi_cli/tools/web/search.py

@@ -1,14 +1,15 @@
 from pathlib import Path
-from typing import override
+from typing import Any, override
 
-import aiohttp
 from kosong.tooling import CallableTool2, ToolReturnType
 from pydantic import BaseModel, Field, ValidationError
 
-import kimi_cli
 from kimi_cli.config import Config
+from kimi_cli.constant import USER_AGENT
 from kimi_cli.soul.toolset import get_current_tool_call_or_none
+from kimi_cli.tools import SkipThisTool
 from kimi_cli.tools.utils import ToolResultBuilder, load_desc
+from kimi_cli.utils.aiohttp import new_client_session
 
 
 class Params(BaseModel):
@@ -39,14 +40,13 @@ class SearchWeb(CallableTool2[Params]):
     description: str = load_desc(Path(__file__).parent / "search.md", {})
     params: type[Params] = Params
 
-    def __init__(self, config: Config, **kwargs):
+    def __init__(self, config: Config, **kwargs: Any):
         super().__init__(**kwargs)
-        if config.services.moonshot_search is not None:
-            self._base_url = config.services.moonshot_search.base_url
-            self._api_key = config.services.moonshot_search.api_key.get_secret_value()
-        else:
-            self._base_url = ""
-            self._api_key = ""
+        if config.services.moonshot_search is None:
+            raise SkipThisTool()
+        self._base_url = config.services.moonshot_search.base_url
+        self._api_key = config.services.moonshot_search.api_key.get_secret_value()
+        self._custom_headers = config.services.moonshot_search.custom_headers or {}
 
     @override
     async def __call__(self, params: Params) -> ToolReturnType:
@@ -62,13 +62,14 @@ class SearchWeb(CallableTool2[Params]):
         assert tool_call is not None, "Tool call is expected to be set"
 
         async with (
-            aiohttp.ClientSession() as session,
+            new_client_session() as session,
             session.post(
                 self._base_url,
                 headers={
-                    "User-Agent": kimi_cli.USER_AGENT,
+                    "User-Agent": USER_AGENT,
                     "Authorization": f"Bearer {self._api_key}",
                     "X-Msh-Tool-Call-Id": tool_call.id,
+                    **self._custom_headers,
                 },
                 json={
                     "text_query": params.query,

kimi_cli/ui/__init__.py

@@ -1,68 +0,0 @@
-import asyncio
-import contextlib
-from collections.abc import Callable, Coroutine
-from typing import Any
-
-from kimi_cli.soul import Soul
-from kimi_cli.soul.wire import Wire
-from kimi_cli.utils.logging import logger
-
-type UILoopFn = Callable[[Wire], Coroutine[Any, Any, None]]
-"""A long-running async function to visualize the agent behavior."""
-
-
-class RunCancelled(Exception):
-    """The run was cancelled by the cancel event."""
-
-
-async def run_soul(
-    soul: Soul,
-    user_input: str,
-    ui_loop_fn: UILoopFn,
-    cancel_event: asyncio.Event,
-):
-    """
-    Run the soul with the given user input.
-
-    `cancel_event` is a outside handle that can be used to cancel the run. When the event is set,
-    the run will be gracefully stopped and a `RunCancelled` will be raised.
-
-    Raises:
-        ChatProviderError: When the LLM provider returns an error.
-        MaxStepsReached: When the maximum number of steps is reached.
-        RunCancelled: When the run is cancelled by the cancel event.
-    """
-    wire = Wire()
-    logger.debug("Starting UI loop with function: {ui_loop_fn}", ui_loop_fn=ui_loop_fn)
-
-    ui_task = asyncio.create_task(ui_loop_fn(wire))
-    soul_task = asyncio.create_task(soul.run(user_input, wire))
-
-    cancel_event_task = asyncio.create_task(cancel_event.wait())
-    await asyncio.wait(
-        [soul_task, cancel_event_task],
-        return_when=asyncio.FIRST_COMPLETED,
-    )
-
-    try:
-        if cancel_event.is_set():
-            logger.debug("Cancelling the run task")
-            soul_task.cancel()
-            try:
-                await soul_task
-            except asyncio.CancelledError:
-                raise RunCancelled from None
-        else:
-            assert soul_task.done()  # either stop event is set or the run task is done
-            cancel_event_task.cancel()
-            with contextlib.suppress(asyncio.CancelledError):
-                await cancel_event_task
-            soul_task.result()  # this will raise if any exception was raised in the run task
-    finally:
-        logger.debug("Shutting down the visualization loop")
-        # shutting down the event queue should break the visualization loop
-        wire.shutdown()
-        try:
-            await asyncio.wait_for(ui_task, timeout=0.5)
-        except TimeoutError:
-            logger.warning("Visualization loop timed out")