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

kimi_cli/ui/acp/__init__.py

@@ -1,29 +1,33 @@
 import asyncio
 import uuid
+from typing import Any
 
-import acp
-import streamingjson
-from kosong.base.message import (
+import acp  # pyright: ignore[reportMissingTypeStubs]
+import streamingjson  # pyright: ignore[reportMissingTypeStubs]
+from kosong.chat_provider import ChatProviderError
+from kosong.message import (
     ContentPart,
     TextPart,
+    ThinkPart,
     ToolCall,
     ToolCallPart,
 )
-from kosong.chat_provider import ChatProviderError
 from kosong.tooling import ToolError, ToolOk, ToolResult
 
-from kimi_cli.soul import LLMNotSet, MaxStepsReached, Soul
-from kimi_cli.soul.wire import (
+from kimi_cli.soul import LLMNotSet, MaxStepsReached, RunCancelled, Soul, run_soul
+from kimi_cli.tools import extract_key_argument
+from kimi_cli.utils.logging import logger
+from kimi_cli.wire import WireUISide
+from kimi_cli.wire.message import (
     ApprovalRequest,
     ApprovalResponse,
+    CompactionBegin,
+    CompactionEnd,
     StatusUpdate,
     StepBegin,
     StepInterrupted,
-    Wire,
+    SubagentEvent,
 )
-from kimi_cli.tools import extract_subtitle
-from kimi_cli.ui import RunCancelled, run_soul
-from kimi_cli.utils.logging import logger
 
 
 class _ToolCallState:
@@ -50,7 +54,7 @@ class _ToolCallState:
     def get_title(self) -> str:
         """Get the current title with subtitle if available."""
         tool_name = self.tool_call.function.name
-        subtitle = extract_subtitle(self.lexer, tool_name)
+        subtitle = extract_key_argument(self.lexer, tool_name)
         if subtitle:
             return f"{tool_name}: {subtitle}"
         return tool_name
@@ -62,9 +66,10 @@ class _RunState:
         """Map of tool call ID (LLM-side ID) to tool call state."""
         self.last_tool_call: _ToolCallState | None = None
         self.cancel_event = asyncio.Event()
+        self.in_thinking = False
 
 
-class ACPAgentImpl:
+class ACPAgent:
     """Implementation of the ACP Agent protocol."""
 
     def __init__(self, soul: Soul, connection: acp.AgentSideConnection):
@@ -117,12 +122,12 @@ class ACPAgentImpl:
         logger.warning("Set session mode: {mode}", mode=params.modeId)
         return None
 
-    async def extMethod(self, method: str, params: dict) -> dict:
+    async def extMethod(self, method: str, params: dict[str, Any]) -> dict[str, Any]:
         """Handle extension method."""
         logger.warning("Unsupported extension method: {method}", method=method)
         return {}
 
-    async def extNotification(self, method: str, params: dict) -> None:
+    async def extNotification(self, method: str, params: dict[str, Any]) -> None:
         """Handle extension notification."""
         logger.warning("Unsupported extension notification: {method}", method=method)
 
@@ -172,34 +177,46 @@ class ACPAgentImpl:
             logger.info("Cancelling running prompt")
             self.run_state.cancel_event.set()
 
-    async def _stream_events(self, wire: Wire):
-        try:
-            # expect a StepBegin
-            assert isinstance(await wire.receive(), StepBegin)
+    async def _stream_events(self, wire: WireUISide):
+        while True:
+            msg = await wire.receive()
 
-            while True:
-                msg = await wire.receive()
+            assert self.run_state is not None
+            if isinstance(msg, ThinkPart) and not self.run_state.in_thinking:
+                await self._send_text("<think>\n")
+                self.run_state.in_thinking = True
+            if not isinstance(msg, ThinkPart) and self.run_state.in_thinking:
+                await self._send_text("\n\n</think>\n\n")
+                self.run_state.in_thinking = False
 
-                if isinstance(msg, TextPart):
-                    await self._send_text(msg.text)
-                elif isinstance(msg, ContentPart):
+            match msg:
+                case StepBegin():
+                    pass
+                case StepInterrupted():
+                    break
+                case CompactionBegin():
+                    pass
+                case CompactionEnd():
+                    pass
+                case StatusUpdate():
+                    pass
+                case ThinkPart(think=think):
+                    await self._send_text(think)
+                case TextPart(text=text):
+                    await self._send_text(text)
+                case ContentPart():
                     logger.warning("Unsupported content part: {part}", part=msg)
                     await self._send_text(f"[{msg.__class__.__name__}]")
-                elif isinstance(msg, ToolCall):
+                case ToolCall():
                     await self._send_tool_call(msg)
-                elif isinstance(msg, ToolCallPart):
+                case ToolCallPart():
                     await self._send_tool_call_part(msg)
-                elif isinstance(msg, ToolResult):
+                case 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
+                case SubagentEvent():
                     pass
-                elif isinstance(msg, StepInterrupted):
-                    break
-        except asyncio.QueueShutDown:
-            logger.debug("Event stream loop shutting down")
+                case ApprovalRequest():
+                    await self._handle_approval_request(msg)
 
     async def _send_text(self, text: str):
         """Send text chunk to client."""
@@ -322,7 +339,7 @@ class ACPAgentImpl:
         # 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(
@@ -388,7 +405,13 @@ def _tool_result_to_acp_content(
     | acp.schema.FileEditToolCallContent
     | acp.schema.TerminalToolCallContent
 ]:
-    def _to_acp_content(part: ContentPart) -> acp.schema.ContentToolCallContent:
+    def _to_acp_content(
+        part: ContentPart,
+    ) -> (
+        acp.schema.ContentToolCallContent
+        | acp.schema.FileEditToolCallContent
+        | acp.schema.TerminalToolCallContent
+    ):
         if isinstance(part, TextPart):
             return acp.schema.ContentToolCallContent(
                 type="content", content=acp.schema.TextContentBlock(type="text", text=part.text)
@@ -402,7 +425,13 @@ def _tool_result_to_acp_content(
                 ),
             )
 
-    content = []
+    content: list[
+        (
+            acp.schema.ContentToolCallContent
+            | acp.schema.FileEditToolCallContent
+            | acp.schema.TerminalToolCallContent
+        )
+    ] = []
     if isinstance(tool_result.output, str):
         content.append(_to_acp_content(TextPart(text=tool_result.output)))
     elif isinstance(tool_result.output, ContentPart):
@@ -428,7 +457,7 @@ class ACPServer:
 
         # Create connection - the library handles all JSON-RPC details!
         _ = acp.AgentSideConnection(
-            lambda conn: ACPAgentImpl(self.soul, conn),
+            lambda conn: ACPAgent(self.soul, conn),
             writer,
             reader,
         )

kimi_cli/ui/print/__init__.py

@@ -1,23 +1,21 @@
 import asyncio
 import json
-import signal
 import sys
 from functools import partial
-from typing import Literal
+from pathlib import Path
 
 import aiofiles
-from kosong.base.message import Message
 from kosong.chat_provider import ChatProviderError
+from kosong.message import Message
+from rich import print
 
-from kimi_cli.soul import LLMNotSet, MaxStepsReached
-from kimi_cli.soul.kimisoul import KimiSoul
-from kimi_cli.soul.wire import StepInterrupted, Wire
-from kimi_cli.ui import RunCancelled, run_soul
+from kimi_cli.cli import InputFormat, OutputFormat
+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
-
-InputFormat = Literal["text", "stream-json"]
-OutputFormat = Literal["text", "stream-json"]
+from kimi_cli.utils.signals import install_sigint_handler
+from kimi_cli.wire import WireUISide
+from kimi_cli.wire.message import StepInterrupted
 
 
 class PrintApp:
@@ -25,17 +23,23 @@ class PrintApp:
     An app implementation that prints the agent behavior to the console.
 
     Args:
-        soul (KimiSoul): The soul to run. Only `KimiSoul` is supported.
+        soul (Soul): The soul to run.
         input_format (InputFormat): The input format to use.
         output_format (OutputFormat): The output format to use.
+        context_file (Path): The file to store the context.
     """
 
-    def __init__(self, soul: KimiSoul, input_format: InputFormat, output_format: OutputFormat):
+    def __init__(
+        self,
+        soul: Soul,
+        input_format: InputFormat,
+        output_format: OutputFormat,
+        context_file: Path,
+    ):
         self.soul = soul
         self.input_format = input_format
         self.output_format = output_format
-        self.soul._approval.set_yolo(True)
-        # TODO(approval): proper approval request handling
+        self.context_file = context_file
 
     async def run(self, command: str | None = None) -> bool:
         cancel_event = asyncio.Event()
@@ -45,7 +49,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()
@@ -92,33 +96,9 @@ class PrintApp:
             print(f"Unknown error: {e}")
             raise
         finally:
-            loop.remove_signal_handler(signal.SIGINT)
+            remove_sigint()
         return False
 
-    # TODO: unify with `_soul_run` in `ShellApp` and `ACPAgentImpl`
-    async def _soul_run(self, user_input: str):
-        wire = Wire()
-        logger.debug("Starting visualization loop")
-
-        if self.output_format == "text":
-            vis_task = asyncio.create_task(self._visualize_text(wire))
-        else:
-            assert self.output_format == "stream-json"
-            if not self.soul._context._file_backend.exists():
-                self.soul._context._file_backend.touch()
-            start_position = self.soul._context._file_backend.stat().st_size
-            vis_task = asyncio.create_task(self._visualize_stream_json(wire, start_position))
-
-        try:
-            await self.soul.run(user_input, wire)
-        finally:
-            wire.shutdown()
-            # shutting down the event queue should break the visualization loop
-            try:
-                await asyncio.wait_for(vis_task, timeout=0.5)
-            except TimeoutError:
-                logger.warning("Visualization loop timed out")
-
     def _read_next_command(self) -> str | None:
         while True:
             json_line = sys.stdin.readline()
@@ -144,33 +124,30 @@ class PrintApp:
             except Exception:
                 logger.warning("Ignoring invalid user message: {json_line}", json_line=json_line)
 
-    async def _visualize_text(self, wire: Wire):
-        try:
+    async def _visualize_text(self, wire: WireUISide):
+        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
+        # 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:
-                msg = await wire.receive()
-                print(msg)
-                if isinstance(msg, StepInterrupted):
-                    break
-        except asyncio.QueueShutDown:
-            logger.debug("Visualization loop shutting down")
-
-    async def _visualize_stream_json(self, wire: Wire, start_position: int):
-        try:
-            async with aiofiles.open(self.soul._context._file_backend, 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")
+                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="")