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

kimi_cli/ui/acp/__init__.py

@@ -1,436 +1,89 @@
-import asyncio
-import uuid
+from __future__ import annotations
+
+from typing import Any, NoReturn
 
 import acp
-import streamingjson
-from kosong.base.message import (
-    ContentPart,
-    TextPart,
-    ToolCall,
-    ToolCallPart,
-)
-from kosong.chat_provider import ChatProviderError
-from kosong.tooling import ToolError, ToolOk, ToolResult
 
-from kimi_cli.soul import LLMNotSet, MaxStepsReached, RunCancelled, Soul, run_soul
-from kimi_cli.tools import extract_subtitle
+from kimi_cli.acp.types import ACPContentBlock, MCPServer
+from kimi_cli.soul import Soul
 from kimi_cli.utils.logging import logger
-from kimi_cli.wire import WireUISide
-from kimi_cli.wire.message import (
-    ApprovalRequest,
-    ApprovalResponse,
-    StatusUpdate,
-    StepBegin,
-    StepInterrupted,
-)
-
-
-class _ToolCallState:
-    """Manages the state of a single tool call for streaming updates."""
-
-    def __init__(self, tool_call: ToolCall):
-        # When the user rejected or cancelled a tool call, the step result may not
-        # be appended to the context. In this case, future step may emit tool call
-        # with the same tool call ID (on the LLM side). To avoid confusion of the
-        # ACP client, we need to ensure the uniqueness in the ACP connection.
-        self.acp_tool_call_id = str(uuid.uuid4())
-
-        self.tool_call = tool_call
-        self.args = tool_call.function.arguments or ""
-        self.lexer = streamingjson.Lexer()
-        if tool_call.function.arguments is not None:
-            self.lexer.append_string(tool_call.function.arguments)
-
-    def append_args_part(self, args_part: str):
-        """Append a new arguments part to the accumulated args and lexer."""
-        self.args += args_part
-        self.lexer.append_string(args_part)
-
-    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)
-        if subtitle:
-            return f"{tool_name}: {subtitle}"
-        return tool_name
-
-
-class _RunState:
-    def __init__(self):
-        self.tool_calls: dict[str, _ToolCallState] = {}
-        """Map of tool call ID (LLM-side ID) to tool call state."""
-        self.last_tool_call: _ToolCallState | None = None
-        self.cancel_event = asyncio.Event()
 
+_DEPRECATED_MESSAGE = (
+    "`kimi --acp` is deprecated. "
+    "Update your ACP client settings to use `kimi acp` without any flags or options."
+)
 
-class ACPAgent:
-    """Implementation of the ACP Agent protocol."""
 
-    def __init__(self, soul: Soul, connection: acp.AgentSideConnection):
+class ACPServerSingleSession:
+    def __init__(self, soul: Soul):
         self.soul = soul
-        self.connection = connection
-        self.session_id: str | None = None
-        self.run_state: _RunState | None = None
-
-    async def initialize(self, params: acp.InitializeRequest) -> acp.InitializeResponse:
-        """Handle initialize request."""
-        logger.info(
-            "ACP server initialized with protocol version: {version}",
-            version=params.protocolVersion,
-        )
 
-        return acp.InitializeResponse(
-            protocolVersion=params.protocolVersion,
-            agentCapabilities=acp.schema.AgentCapabilities(
-                loadSession=False,
-                promptCapabilities=acp.schema.PromptCapabilities(
-                    embeddedContext=False, image=False, audio=False
-                ),
-            ),
-            authMethods=[],
-        )
-
-    async def authenticate(self, params: acp.AuthenticateRequest) -> None:
-        """Handle authenticate request."""
-        logger.info("Authenticate with method: {method}", method=params.methodId)
-
-    async def newSession(self, params: acp.NewSessionRequest) -> acp.NewSessionResponse:
-        """Handle new session request."""
-        self.session_id = f"sess_{uuid.uuid4().hex[:16]}"
-        logger.info("Created session {id} with cwd: {cwd}", id=self.session_id, cwd=params.cwd)
-        return acp.NewSessionResponse(sessionId=self.session_id)
-
-    async def loadSession(self, params: acp.LoadSessionRequest) -> None:
-        """Handle load session request."""
-        self.session_id = params.sessionId
-        logger.info("Loaded session: {id}", id=self.session_id)
-
-    async def setSessionModel(self, params: acp.SetSessionModelRequest) -> None:
-        """Handle set session model request."""
-        logger.warning("Set session model: {model}", model=params.modelId)
-
-    async def setSessionMode(
-        self, params: acp.SetSessionModeRequest
+    def on_connect(self, conn: acp.Client) -> None:
+        logger.info("ACP client connected")
+
+    def _raise(self) -> NoReturn:
+        logger.error(_DEPRECATED_MESSAGE)
+        raise acp.RequestError.invalid_params({"error": _DEPRECATED_MESSAGE})
+
+    async def initialize(
+        self,
+        protocol_version: int,
+        client_capabilities: acp.schema.ClientCapabilities | None = None,
+        client_info: acp.schema.Implementation | None = None,
+        **kwargs: Any,
+    ) -> acp.InitializeResponse:
+        self._raise()
+
+    async def new_session(
+        self, cwd: str, mcp_servers: list[MCPServer], **kwargs: Any
+    ) -> acp.NewSessionResponse:
+        self._raise()
+
+    async def load_session(
+        self, cwd: str, mcp_servers: list[MCPServer], session_id: str, **kwargs: Any
+    ) -> None:
+        self._raise()
+
+    async def list_sessions(
+        self, cursor: str | None = None, cwd: str | None = None, **kwargs: Any
+    ) -> acp.schema.ListSessionsResponse:
+        self._raise()
+
+    async def set_session_mode(
+        self, mode_id: str, session_id: str, **kwargs: Any
     ) -> acp.SetSessionModeResponse | None:
-        """Handle set session mode request."""
-        logger.warning("Set session mode: {mode}", mode=params.modeId)
-        return None
-
-    async def extMethod(self, method: str, params: dict) -> dict:
-        """Handle extension method."""
-        logger.warning("Unsupported extension method: {method}", method=method)
-        return {}
-
-    async def extNotification(self, method: str, params: dict) -> None:
-        """Handle extension notification."""
-        logger.warning("Unsupported extension notification: {method}", method=method)
-
-    async def prompt(self, params: acp.PromptRequest) -> acp.PromptResponse:
-        """Handle prompt request with streaming support."""
-        # Extract text from prompt content blocks
-        prompt_text = "\n".join(
-            block.text for block in params.prompt if isinstance(block, acp.schema.TextContentBlock)
-        )
-
-        if not prompt_text:
-            raise acp.RequestError.invalid_params({"reason": "No text in prompt"})
-
-        logger.info("Processing prompt: {text}", text=prompt_text[:100])
-
-        self.run_state = _RunState()
-        try:
-            await run_soul(self.soul, prompt_text, self._stream_events, self.run_state.cancel_event)
-            return acp.PromptResponse(stopReason="end_turn")
-        except LLMNotSet:
-            logger.error("LLM not set")
-            raise acp.RequestError.internal_error({"error": "LLM not set"}) from None
-        except ChatProviderError as e:
-            logger.exception("LLM provider error:")
-            raise acp.RequestError.internal_error({"error": f"LLM provider error: {e}"}) from e
-        except MaxStepsReached as e:
-            logger.warning("Max steps reached: {n}", n=e.n_steps)
-            return acp.PromptResponse(stopReason="max_turn_requests")
-        except RunCancelled:
-            logger.info("Prompt cancelled by user")
-            return acp.PromptResponse(stopReason="cancelled")
-        except BaseException as e:
-            logger.exception("Unknown error:")
-            raise acp.RequestError.internal_error({"error": f"Unknown error: {e}"}) from e
-        finally:
-            self.run_state = None
-
-    async def cancel(self, params: acp.CancelNotification) -> None:
-        """Handle cancel notification."""
-        logger.info("Cancel for session: {id}", id=params.sessionId)
-
-        if self.run_state is None:
-            logger.warning("No running prompt to cancel")
-            return
-
-        if not self.run_state.cancel_event.is_set():
-            logger.info("Cancelling running prompt")
-            self.run_state.cancel_event.set()
+        self._raise()
 
-    async def _stream_events(self, wire: WireUISide):
-        assert isinstance(await wire.receive(), StepBegin)
+    async def set_session_model(
+        self, model_id: str, session_id: str, **kwargs: Any
+    ) -> acp.SetSessionModelResponse | None:
+        self._raise()
 
-        while True:
-            msg = await wire.receive()
+    async def authenticate(self, method_id: str, **kwargs: Any) -> acp.AuthenticateResponse | None:
+        self._raise()
 
-            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 prompt(
+        self, prompt: list[ACPContentBlock], session_id: str, **kwargs: Any
+    ) -> acp.PromptResponse:
+        self._raise()
 
-    async def _send_text(self, text: str):
-        """Send text chunk to client."""
-        if not self.session_id:
-            return
+    async def cancel(self, session_id: str, **kwargs: Any) -> None:
+        self._raise()
 
-        await self.connection.sessionUpdate(
-            acp.SessionNotification(
-                sessionId=self.session_id,
-                update=acp.schema.AgentMessageChunk(
-                    content=acp.schema.TextContentBlock(type="text", text=text),
-                    sessionUpdate="agent_message_chunk",
-                ),
-            )
-        )
+    async def ext_method(self, method: str, params: dict[str, Any]) -> dict[str, Any]:
+        self._raise()
 
-    async def _send_tool_call(self, tool_call: ToolCall):
-        """Send tool call to client."""
-        assert self.run_state is not None
-        if not self.session_id:
-            return
+    async def ext_notification(self, method: str, params: dict[str, Any]) -> None:
+        self._raise()
 
-        # Create and store tool call state
-        state = _ToolCallState(tool_call)
-        self.run_state.tool_calls[tool_call.id] = state
-        self.run_state.last_tool_call = state
 
-        await self.connection.sessionUpdate(
-            acp.SessionNotification(
-                sessionId=self.session_id,
-                update=acp.schema.ToolCallStart(
-                    sessionUpdate="tool_call",
-                    toolCallId=state.acp_tool_call_id,
-                    title=state.get_title(),
-                    status="in_progress",
-                    content=[
-                        acp.schema.ContentToolCallContent(
-                            type="content",
-                            content=acp.schema.TextContentBlock(type="text", text=state.args),
-                        )
-                    ],
-                ),
-            )
-        )
-        logger.debug("Sent tool call: {name}", name=tool_call.function.name)
-
-    async def _send_tool_call_part(self, part: ToolCallPart):
-        """Send tool call part (streaming arguments)."""
-        assert self.run_state is not None
-        if not self.session_id or not part.arguments_part or self.run_state.last_tool_call is None:
-            return
-
-        # Append new arguments part to the last tool call
-        self.run_state.last_tool_call.append_args_part(part.arguments_part)
-
-        # Update the tool call with new content and title
-        update = acp.schema.ToolCallProgress(
-            sessionUpdate="tool_call_update",
-            toolCallId=self.run_state.last_tool_call.acp_tool_call_id,
-            title=self.run_state.last_tool_call.get_title(),
-            status="in_progress",
-            content=[
-                acp.schema.ContentToolCallContent(
-                    type="content",
-                    content=acp.schema.TextContentBlock(
-                        type="text", text=self.run_state.last_tool_call.args
-                    ),
-                )
-            ],
-        )
-
-        await self.connection.sessionUpdate(
-            acp.SessionNotification(sessionId=self.session_id, update=update)
-        )
-        logger.debug("Sent tool call update: {delta}", delta=part.arguments_part[:50])
-
-    async def _send_tool_result(self, result: ToolResult):
-        """Send tool result to client."""
-        assert self.run_state is not None
-        if not self.session_id:
-            return
-
-        tool_result = result.result
-        is_error = isinstance(tool_result, ToolError)
-
-        state = self.run_state.tool_calls.pop(result.tool_call_id, None)
-        if state is None:
-            logger.warning("Tool call not found: {id}", id=result.tool_call_id)
-            return
-
-        update = acp.schema.ToolCallProgress(
-            sessionUpdate="tool_call_update",
-            toolCallId=state.acp_tool_call_id,
-            status="failed" if is_error else "completed",
-        )
-
-        if state.tool_call.function.name == "SetTodoList" and not is_error:
-            update.content = _tool_result_to_acp_content(tool_result)
-
-        await self.connection.sessionUpdate(
-            acp.SessionNotification(sessionId=self.session_id, update=update)
-        )
-
-        logger.debug("Sent tool result: {id}", id=result.tool_call_id)
-
-    async def _handle_approval_request(self, request: ApprovalRequest):
-        """Handle approval request by sending permission request to client."""
-        assert self.run_state is not None
-        if not self.session_id:
-            logger.warning("No session ID, auto-rejecting approval request")
-            request.resolve(ApprovalResponse.REJECT)
-            return
-
-        state = self.run_state.tool_calls.get(request.tool_call_id, None)
-        if state is None:
-            logger.warning("Tool call not found: {id}", id=request.tool_call_id)
-            request.resolve(ApprovalResponse.REJECT)
-            return
-
-        # Create permission request with options
-        permission_request = acp.RequestPermissionRequest(
-            sessionId=self.session_id,
-            toolCall=acp.schema.ToolCall(
-                toolCallId=state.acp_tool_call_id,
-                content=[
-                    acp.schema.ContentToolCallContent(
-                        type="content",
-                        content=acp.schema.TextContentBlock(
-                            type="text",
-                            text=f"Requesting approval to perform: {request.description}",
-                        ),
-                    ),
-                ],
-            ),
-            options=[
-                acp.schema.PermissionOption(
-                    optionId="approve",
-                    name="Approve",
-                    kind="allow_once",
-                ),
-                acp.schema.PermissionOption(
-                    optionId="approve_for_session",
-                    name="Approve for this session",
-                    kind="allow_always",
-                ),
-                acp.schema.PermissionOption(
-                    optionId="reject",
-                    name="Reject",
-                    kind="reject_once",
-                ),
-            ],
-        )
-
-        try:
-            # Send permission request and wait for response
-            logger.debug("Requesting permission for action: {action}", action=request.action)
-            response = await self.connection.requestPermission(permission_request)
-            logger.debug("Received permission response: {response}", response=response)
-
-            # Process the outcome
-            if isinstance(response.outcome, acp.schema.AllowedOutcome):
-                # selected
-                if response.outcome.optionId == "approve":
-                    logger.debug("Permission granted for: {action}", action=request.action)
-                    request.resolve(ApprovalResponse.APPROVE)
-                elif response.outcome.optionId == "approve_for_session":
-                    logger.debug("Permission granted for session: {action}", action=request.action)
-                    request.resolve(ApprovalResponse.APPROVE_FOR_SESSION)
-                else:
-                    logger.debug("Permission denied for: {action}", action=request.action)
-                    request.resolve(ApprovalResponse.REJECT)
-            else:
-                # cancelled
-                logger.debug("Permission request cancelled for: {action}", action=request.action)
-                request.resolve(ApprovalResponse.REJECT)
-        except Exception:
-            logger.exception("Error handling approval request:")
-            # On error, reject the request
-            request.resolve(ApprovalResponse.REJECT)
-
-
-def _tool_result_to_acp_content(
-    tool_result: ToolOk | ToolError,
-) -> list[
-    acp.schema.ContentToolCallContent
-    | acp.schema.FileEditToolCallContent
-    | acp.schema.TerminalToolCallContent
-]:
-    def _to_acp_content(part: ContentPart) -> acp.schema.ContentToolCallContent:
-        if isinstance(part, TextPart):
-            return acp.schema.ContentToolCallContent(
-                type="content", content=acp.schema.TextContentBlock(type="text", text=part.text)
-            )
-        else:
-            logger.warning("Unsupported content part in tool result: {part}", part=part)
-            return acp.schema.ContentToolCallContent(
-                type="content",
-                content=acp.schema.TextContentBlock(
-                    type="text", text=f"[{part.__class__.__name__}]"
-                ),
-            )
-
-    content = []
-    if isinstance(tool_result.output, str):
-        content.append(_to_acp_content(TextPart(text=tool_result.output)))
-    elif isinstance(tool_result.output, ContentPart):
-        content.append(_to_acp_content(tool_result.output))
-    elif isinstance(tool_result.output, list):
-        content.extend(_to_acp_content(part) for part in tool_result.output)
-
-    return content
-
-
-class ACPServer:
+class ACP:
     """ACP server using the official acp library."""
 
     def __init__(self, soul: Soul):
         self.soul = soul
 
-    async def run(self) -> bool:
+    async def run(self):
         """Run the ACP server."""
-        logger.info("Starting ACP server on stdio")
-
-        # Get stdio streams
-        reader, writer = await acp.stdio_streams()
-
-        # Create connection - the library handles all JSON-RPC details!
-        _ = acp.AgentSideConnection(
-            lambda conn: ACPAgent(self.soul, conn),
-            writer,
-            reader,
-        )
-
-        logger.info("ACP server ready")
-
-        # Keep running - connection handles everything
-        await asyncio.Event().wait()
-
-        return True
+        logger.info("Starting ACP server (single session) on stdio")
+        await acp.run_agent(ACPServerSingleSession(self.soul))

kimi_cli/ui/print/__init__.py

@@ -1,24 +1,31 @@
+from __future__ import annotations
+
 import asyncio
 import json
 import sys
 from functools import partial
 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.cli import InputFormat, OutputFormat
-from kimi_cli.soul import LLMNotSet, MaxStepsReached, RunCancelled, Soul, run_soul
+from kimi_cli.soul import (
+    LLMNotSet,
+    LLMNotSupported,
+    MaxStepsReached,
+    RunCancelled,
+    Soul,
+    run_soul,
+)
+from kimi_cli.soul.kimisoul import KimiSoul
+from kimi_cli.ui.print.visualize import visualize
 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
 
 
-class PrintApp:
+class Print:
     """
     An app implementation that prints the agent behavior to the console.
 
@@ -27,6 +34,7 @@ class PrintApp:
         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.
+        final_only (bool): Whether to only print the final assistant message.
     """
 
     def __init__(
@@ -35,11 +43,14 @@ class PrintApp:
         input_format: InputFormat,
         output_format: OutputFormat,
         context_file: Path,
+        *,
+        final_only: bool = False,
     ):
         self.soul = soul
-        self.input_format = input_format
-        self.output_format = output_format
+        self.input_format: InputFormat = input_format
+        self.output_format: OutputFormat = output_format
         self.context_file = context_file
+        self.final_only = final_only
 
     async def run(self, command: str | None = None) -> bool:
         cancel_event = asyncio.Event()
@@ -68,26 +79,31 @@ class PrintApp:
 
                 if command:
                     logger.info("Running agent with command: {command}", command=command)
-                    if self.output_format == "text":
-                        visualize_fn = self._visualize_text
+                    if self.output_format == "text" and not self.final_only:
                         print(command)
-                    else:
-                        assert self.output_format == "stream-json"
-                        visualize_fn = partial(self._visualize_stream_json, start_position=0)
-                    await run_soul(self.soul, command, visualize_fn, cancel_event)
+                    await run_soul(
+                        self.soul,
+                        command,
+                        partial(visualize, self.output_format, self.final_only),
+                        cancel_event,
+                        self.soul.wire_file if isinstance(self.soul, KimiSoul) else None,
+                    )
                 else:
                     logger.info("Empty command, skipping")
 
                 command = None
-        except LLMNotSet:
-            logger.error("LLM not set")
-            print("LLM not set")
+        except LLMNotSet as e:
+            logger.exception("LLM not set:")
+            print(str(e))
+        except LLMNotSupported as e:
+            logger.exception("LLM not supported:")
+            print(str(e))
         except ChatProviderError as e:
             logger.exception("LLM provider error:")
-            print(f"LLM provider error: {e}")
+            print(str(e))
         except MaxStepsReached as e:
             logger.warning("Max steps reached: {n_steps}", n_steps=e.n_steps)
-            print(f"Max steps reached: {e.n_steps}")
+            print(str(e))
         except RunCancelled:
             logger.error("Interrupted by user")
             print("Interrupted by user")
@@ -115,7 +131,7 @@ class PrintApp:
                 data = json.loads(json_line)
                 message = Message.model_validate(data)
                 if message.role == "user":
-                    return message_extract_text(message)
+                    return message.extract_text(sep="\n")
                 logger.warning(
                     "Ignoring message with role `{role}`: {json_line}",
                     role=message.role,
@@ -123,31 +139,3 @@ class PrintApp:
                 )
             except Exception:
                 logger.warning("Ignoring invalid user message: {json_line}", json_line=json_line)
-
-    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:
-                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="")