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

kimi_cli/soul/kimisoul.py

@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING
 import kosong
 import tenacity
 from kosong import StepResult
-from kosong.base.message import Message
+from kosong.base.message import ContentPart, ImageURLPart, Message
 from kosong.chat_provider import (
     APIConnectionError,
     APIStatusError,
@@ -16,13 +16,19 @@ from kosong.chat_provider import (
 from kosong.tooling import ToolResult
 from tenacity import RetryCallState, retry_if_exception, stop_after_attempt, wait_exponential_jitter
 
-from kimi_cli.config import LoopControl
-from kimi_cli.soul import LLMNotSet, MaxStepsReached, Soul, StatusSnapshot, wire_send
+from kimi_cli.soul import (
+    LLMNotSet,
+    LLMNotSupported,
+    MaxStepsReached,
+    Soul,
+    StatusSnapshot,
+    wire_send,
+)
 from kimi_cli.soul.agent import Agent
 from kimi_cli.soul.compaction import SimpleCompaction
 from kimi_cli.soul.context import Context
-from kimi_cli.soul.globals import AgentGlobals
 from kimi_cli.soul.message import system, tool_result_to_messages
+from kimi_cli.soul.runtime import Runtime
 from kimi_cli.tools.dmail import NAME as SendDMail_NAME
 from kimi_cli.tools.utils import ToolRejectedError
 from kimi_cli.utils.logging import logger
@@ -43,30 +49,28 @@ class KimiSoul(Soul):
     def __init__(
         self,
         agent: Agent,
-        agent_globals: AgentGlobals,
+        runtime: Runtime,
         *,
         context: Context,
-        loop_control: LoopControl,
     ):
         """
         Initialize the soul.
 
         Args:
             agent (Agent): The agent to run.
-            agent_globals (AgentGlobals): Global states and parameters.
+            runtime (Runtime): Runtime parameters and states.
             context (Context): The context of the agent.
-            loop_control (LoopControl): The control parameters for the agent loop.
         """
         self._agent = agent
-        self._agent_globals = agent_globals
-        self._denwa_renji = agent_globals.denwa_renji
-        self._approval = agent_globals.approval
+        self._runtime = runtime
+        self._denwa_renji = runtime.denwa_renji
+        self._approval = runtime.approval
         self._context = context
-        self._loop_control = loop_control
+        self._loop_control = runtime.config.loop_control
         self._compaction = SimpleCompaction()  # TODO: maybe configurable and composable
         self._reserved_tokens = RESERVED_TOKENS
-        if self._agent_globals.llm is not None:
-            assert self._reserved_tokens <= self._agent_globals.llm.max_context_size
+        if self._runtime.llm is not None:
+            assert self._reserved_tokens <= self._runtime.llm.max_context_size
 
         for tool in agent.toolset.tools:
             if tool.name == SendDMail_NAME:
@@ -81,25 +85,36 @@ class KimiSoul(Soul):
 
     @property
     def model(self) -> str:
-        return self._agent_globals.llm.chat_provider.model_name if self._agent_globals.llm else ""
+        return self._runtime.llm.chat_provider.model_name if self._runtime.llm else ""
 
     @property
     def status(self) -> StatusSnapshot:
         return StatusSnapshot(context_usage=self._context_usage)
 
+    @property
+    def context(self) -> Context:
+        return self._context
+
     @property
     def _context_usage(self) -> float:
-        if self._agent_globals.llm is not None:
-            return self._context.token_count / self._agent_globals.llm.max_context_size
+        if self._runtime.llm is not None:
+            return self._context.token_count / self._runtime.llm.max_context_size
         return 0.0
 
     async def _checkpoint(self):
         await self._context.checkpoint(self._checkpoint_with_user_message)
 
-    async def run(self, user_input: str):
-        if self._agent_globals.llm is None:
+    async def run(self, user_input: str | list[ContentPart]):
+        if self._runtime.llm is None:
             raise LLMNotSet()
 
+        if (
+            isinstance(user_input, list)
+            and any(isinstance(part, ImageURLPart) for part in user_input)
+            and not self._runtime.llm.supports_image_in
+        ):
+            raise LLMNotSupported(self._runtime.llm, ["image_in"])
+
         await self._checkpoint()  # this creates the checkpoint 0 on first run
         await self._context.append_message(Message(role="user", content=user_input))
         logger.debug("Appended user message to context")
@@ -107,7 +122,7 @@ class KimiSoul(Soul):
 
     async def _agent_loop(self):
         """The main agent loop for one run."""
-        assert self._agent_globals.llm is not None
+        assert self._runtime.llm is not None
 
         async def _pipe_approval_to_wire():
             while True:
@@ -126,7 +141,7 @@ class KimiSoul(Soul):
                 # compact the context if needed
                 if (
                     self._context.token_count + self._reserved_tokens
-                    >= self._agent_globals.llm.max_context_size
+                    >= self._runtime.llm.max_context_size
                 ):
                     logger.info("Context too long, compacting...")
                     wire_send(CompactionBegin())
@@ -159,8 +174,8 @@ class KimiSoul(Soul):
     async def _step(self) -> bool:
         """Run an single step and return whether the run should be stopped."""
         # already checked in `run`
-        assert self._agent_globals.llm is not None
-        chat_provider = self._agent_globals.llm.chat_provider
+        assert self._runtime.llm is not None
+        chat_provider = self._runtime.llm.chat_provider
 
         @tenacity.retry(
             retry=retry_if_exception(self._is_retryable_error),
@@ -255,9 +270,9 @@ class KimiSoul(Soul):
             reraise=True,
         )
         async def _compact_with_retry() -> Sequence[Message]:
-            if self._agent_globals.llm is None:
+            if self._runtime.llm is None:
                 raise LLMNotSet()
-            return await self._compaction.compact(self._context.history, self._agent_globals.llm)
+            return await self._compaction.compact(self._context.history, self._runtime.llm)
 
         compacted_messages = await _compact_with_retry()
         await self._context.revert_to(0)

kimi_cli/soul/message.py

@@ -14,7 +14,7 @@ def tool_result_to_messages(tool_result: ToolResult) -> list[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: list[ContentPart] = [system(message)]
+        content: list[ContentPart] = [system(f"ERROR: {message}")]
         if tool_result.result.output:
             content.append(TextPart(text=tool_result.result.output))
         return [

kimi_cli/soul/{globals.py → runtime.py}

@@ -6,7 +6,7 @@ from typing import NamedTuple
 
 from kimi_cli.config import Config
 from kimi_cli.llm import LLM
-from kimi_cli.metadata import Session
+from kimi_cli.session import Session
 from kimi_cli.soul.approval import Approval
 from kimi_cli.soul.denwarenji import DenwaRenji
 from kimi_cli.utils.logging import logger
@@ -58,8 +58,8 @@ def _list_work_dir(work_dir: Path) -> str:
     return ls.stdout.strip()
 
 
-class AgentGlobals(NamedTuple):
-    """Agent globals."""
+class Runtime(NamedTuple):
+    """Agent runtime."""
 
     config: Config
     llm: LLM | None
@@ -68,22 +68,24 @@ class AgentGlobals(NamedTuple):
     denwa_renji: DenwaRenji
     approval: Approval
 
-    @classmethod
+    @staticmethod
     async def create(
-        cls, config: Config, llm: LLM | None, session: Session, yolo: bool
-    ) -> "AgentGlobals":
-        work_dir = Path(session.work_dir.path)
+        config: Config,
+        llm: LLM | None,
+        session: Session,
+        yolo: bool,
+    ) -> "Runtime":
         # FIXME: do these asynchronously
-        ls_output = _list_work_dir(work_dir)
-        agents_md = load_agents_md(work_dir) or ""
+        ls_output = _list_work_dir(session.work_dir)
+        agents_md = load_agents_md(session.work_dir) or ""
 
-        return cls(
+        return Runtime(
             config=config,
             llm=llm,
             session=session,
             builtin_args=BuiltinSystemPromptArgs(
                 KIMI_NOW=datetime.now().astimezone().isoformat(),
-                KIMI_WORK_DIR=work_dir,
+                KIMI_WORK_DIR=session.work_dir,
                 KIMI_WORK_DIR_LS=ls_output,
                 KIMI_AGENTS_MD=agents_md,
             ),

kimi_cli/tools/file/glob.py

@@ -8,7 +8,7 @@ import aiofiles.os
 from kosong.tooling import CallableTool2, ToolError, ToolOk, ToolReturnType
 from pydantic import BaseModel, Field
 
-from kimi_cli.soul.globals import BuiltinSystemPromptArgs
+from kimi_cli.soul.runtime import BuiltinSystemPromptArgs
 from kimi_cli.tools.utils import load_desc
 
 MAX_MATCHES = 1000

kimi_cli/tools/file/patch.py

@@ -7,7 +7,7 @@ from kosong.tooling import CallableTool2, ToolError, ToolOk, ToolReturnType
 from pydantic import BaseModel, Field
 
 from kimi_cli.soul.approval import Approval
-from kimi_cli.soul.globals import BuiltinSystemPromptArgs
+from kimi_cli.soul.runtime import BuiltinSystemPromptArgs
 from kimi_cli.tools.file import FileActions
 from kimi_cli.tools.utils import ToolRejectedError
 

kimi_cli/tools/file/read.py

@@ -5,7 +5,7 @@ import aiofiles
 from kosong.tooling import CallableTool2, ToolError, ToolOk, ToolReturnType
 from pydantic import BaseModel, Field
 
-from kimi_cli.soul.globals import BuiltinSystemPromptArgs
+from kimi_cli.soul.runtime import BuiltinSystemPromptArgs
 from kimi_cli.tools.utils import load_desc, truncate_line
 
 MAX_LINES = 1000

kimi_cli/tools/file/replace.py

@@ -6,7 +6,7 @@ from kosong.tooling import CallableTool2, ToolError, ToolOk, ToolReturnType
 from pydantic import BaseModel, Field
 
 from kimi_cli.soul.approval import Approval
-from kimi_cli.soul.globals import BuiltinSystemPromptArgs
+from kimi_cli.soul.runtime import BuiltinSystemPromptArgs
 from kimi_cli.tools.file import FileActions
 from kimi_cli.tools.utils import ToolRejectedError
 

kimi_cli/tools/file/write.py

@@ -6,7 +6,7 @@ from kosong.tooling import CallableTool2, ToolError, ToolOk, ToolReturnType
 from pydantic import BaseModel, Field
 
 from kimi_cli.soul.approval import Approval
-from kimi_cli.soul.globals import BuiltinSystemPromptArgs
+from kimi_cli.soul.runtime import BuiltinSystemPromptArgs
 from kimi_cli.tools.file import FileActions
 from kimi_cli.tools.utils import ToolRejectedError
 

kimi_cli/tools/task/__init__.py

@@ -5,12 +5,12 @@ from typing import override
 from kosong.tooling import CallableTool2, ToolError, ToolOk, ToolReturnType
 from pydantic import BaseModel, Field
 
-from kimi_cli.agentspec import ResolvedAgentSpec
+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.globals import AgentGlobals
 from kimi_cli.soul.kimisoul import KimiSoul
+from kimi_cli.soul.runtime import Runtime
 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
@@ -49,28 +49,37 @@ class Task(CallableTool2[Params]):
     name: str = "Task"
     params: type[Params] = Params
 
-    def __init__(self, agent_spec: ResolvedAgentSpec, agent_globals: AgentGlobals, **kwargs):
-        subagents: dict[str, Agent] = {}
-        descs = []
-
-        # load all 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):
         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."""
@@ -85,6 +94,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}",
@@ -117,12 +130,7 @@ class Task(CallableTool2[Params]):
 
         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,
-        )
+        soul = KimiSoul(agent, runtime=self._runtime, context=context)
 
         try:
             await run_soul(soul, prompt, _ui_loop_fn, asyncio.Event())

kimi_cli/tools/web/search.py

@@ -44,9 +44,11 @@ class SearchWeb(CallableTool2[Params]):
         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()
+            self._custom_headers = config.services.moonshot_search.custom_headers or {}
         else:
             self._base_url = ""
             self._api_key = ""
+            self._custom_headers = {}
 
     @override
     async def __call__(self, params: Params) -> ToolReturnType:
@@ -69,6 +71,7 @@ class SearchWeb(CallableTool2[Params]):
                     "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/acp/__init__.py

@@ -172,33 +172,29 @@ class ACPAgent:
             self.run_state.cancel_event.set()
 
     async def _stream_events(self, wire: WireUISide):
-        try:
-            # expect a StepBegin
-            assert isinstance(await wire.receive(), StepBegin)
-
-            while True:
-                msg = await wire.receive()
-
-                if isinstance(msg, TextPart):
-                    await self._send_text(msg.text)
-                elif isinstance(msg, ContentPart):
-                    logger.warning("Unsupported content part: {part}", part=msg)
-                    await self._send_text(f"[{msg.__class__.__name__}]")
-                elif isinstance(msg, ToolCall):
-                    await self._send_tool_call(msg)
-                elif isinstance(msg, ToolCallPart):
-                    await self._send_tool_call_part(msg)
-                elif isinstance(msg, ToolResult):
-                    await self._send_tool_result(msg)
-                elif isinstance(msg, ApprovalRequest):
-                    await self._handle_approval_request(msg)
-                elif isinstance(msg, StatusUpdate):
-                    # TODO: stream status if needed
-                    pass
-                elif isinstance(msg, StepInterrupted):
-                    break
-        except asyncio.QueueShutDown:
-            logger.debug("Event stream loop shutting down")
+        assert isinstance(await wire.receive(), StepBegin)
+
+        while True:
+            msg = await wire.receive()
+
+            if isinstance(msg, TextPart):
+                await self._send_text(msg.text)
+            elif isinstance(msg, ContentPart):
+                logger.warning("Unsupported content part: {part}", part=msg)
+                await self._send_text(f"[{msg.__class__.__name__}]")
+            elif isinstance(msg, ToolCall):
+                await self._send_tool_call(msg)
+            elif isinstance(msg, ToolCallPart):
+                await self._send_tool_call_part(msg)
+            elif isinstance(msg, ToolResult):
+                await self._send_tool_result(msg)
+            elif isinstance(msg, ApprovalRequest):
+                await self._handle_approval_request(msg)
+            elif isinstance(msg, StatusUpdate):
+                # TODO: stream status if needed
+                pass
+            elif isinstance(msg, StepInterrupted):
+                break
 
     async def _send_text(self, text: str):
         """Send text chunk to client."""
@@ -321,7 +317,7 @@ class ACPAgent:
         # Create permission request with options
         permission_request = acp.RequestPermissionRequest(
             sessionId=self.session_id,
-            toolCall=acp.schema.ToolCallUpdate(
+            toolCall=acp.schema.ToolCall(
                 toolCallId=state.acp_tool_call_id,
                 content=[
                     acp.schema.ContentToolCallContent(

kimi_cli/ui/print/__init__.py

@@ -1,6 +1,5 @@
 import asyncio
 import json
-import signal
 import sys
 from functools import partial
 from pathlib import Path
@@ -9,10 +8,12 @@ from typing import Literal
 import aiofiles
 from kosong.base.message import Message
 from kosong.chat_provider import ChatProviderError
+from rich import print
 
 from kimi_cli.soul import LLMNotSet, MaxStepsReached, RunCancelled, Soul, run_soul
 from kimi_cli.utils.logging import logger
 from kimi_cli.utils.message import message_extract_text
+from kimi_cli.utils.signals import install_sigint_handler
 from kimi_cli.wire import WireUISide
 from kimi_cli.wire.message import StepInterrupted
 
@@ -51,7 +52,7 @@ class PrintApp:
             cancel_event.set()
 
         loop = asyncio.get_running_loop()
-        loop.add_signal_handler(signal.SIGINT, _handler)
+        remove_sigint = install_sigint_handler(loop, _handler)
 
         if command is None and not sys.stdin.isatty() and self.input_format == "text":
             command = sys.stdin.read().strip()
@@ -98,7 +99,7 @@ class PrintApp:
             print(f"Unknown error: {e}")
             raise
         finally:
-            loop.remove_signal_handler(signal.SIGINT)
+            remove_sigint()
         return False
 
     def _read_next_command(self) -> str | None:
@@ -127,33 +128,29 @@ class PrintApp:
                 logger.warning("Ignoring invalid user message: {json_line}", json_line=json_line)
 
     async def _visualize_text(self, wire: WireUISide):
-        try:
-            while True:
-                msg = await wire.receive()
-                print(msg)
-                if isinstance(msg, StepInterrupted):
-                    break
-        except asyncio.QueueShutDown:
-            logger.debug("Visualization loop shutting down")
+        while True:
+            msg = await wire.receive()
+            print(msg)
+            if isinstance(msg, StepInterrupted):
+                break
 
     async def _visualize_stream_json(self, wire: WireUISide, start_position: int):
         # TODO: be aware of context compaction
-        try:
-            async with aiofiles.open(self.context_file, encoding="utf-8") as f:
-                await f.seek(start_position)
-                while True:
-                    should_end = False
-                    while wire._queue.qsize() > 0:
-                        msg = wire._queue.get_nowait()
-                        if isinstance(msg, StepInterrupted):
-                            should_end = True
-
-                    line = await f.readline()
-                    if not line:
-                        if should_end:
-                            break
-                        await asyncio.sleep(0.1)
-                        continue
-                    print(line, end="")
-        except asyncio.QueueShutDown:
-            logger.debug("Visualization loop shutting down")
+        # FIXME: this is only a temporary impl, may miss the last lines of the context file
+        if not self.context_file.exists():
+            self.context_file.touch()
+        async with aiofiles.open(self.context_file, encoding="utf-8") as f:
+            await f.seek(start_position)
+            while True:
+                should_end = False
+                while (msg := wire.receive_nowait()) is not None:
+                    if isinstance(msg, StepInterrupted):
+                        should_end = True
+
+                line = await f.readline()
+                if not line:
+                    if should_end:
+                        break
+                    await asyncio.sleep(0.1)
+                    continue
+                print(line, end="")