bhu-cli 1.46.0__py3-none-any.whl

kimi_cli/acp/session.py

@@ -0,0 +1,499 @@
+from __future__ import annotations
+
+import asyncio
+import uuid
+from contextvars import ContextVar
+
+import acp
+import streamingjson  # type: ignore[reportMissingTypeStubs]
+from kaos import Kaos, reset_current_kaos, set_current_kaos
+from kosong.chat_provider import APIStatusError, ChatProviderError
+
+from kimi_cli.acp.convert import (
+    acp_blocks_to_content_parts,
+    display_block_to_acp_content,
+    tool_result_to_acp_content,
+)
+from kimi_cli.acp.types import ACPContentBlock
+from kimi_cli.app import KimiCLI
+from kimi_cli.soul import LLMNotSet, LLMNotSupported, MaxStepsReached, RunCancelled
+from kimi_cli.tools import extract_key_argument
+from kimi_cli.utils.logging import logger
+from kimi_cli.wire.types import (
+    ApprovalRequest,
+    ApprovalResponse,
+    CompactionBegin,
+    CompactionEnd,
+    ContentPart,
+    MCPLoadingBegin,
+    MCPLoadingEnd,
+    Notification,
+    PlanDisplay,
+    QuestionRequest,
+    StatusUpdate,
+    SteerInput,
+    StepBegin,
+    StepInterrupted,
+    StepRetry,
+    SubagentEvent,
+    TextPart,
+    ThinkPart,
+    TodoDisplayBlock,
+    ToolCall,
+    ToolCallPart,
+    ToolCallRequest,
+    ToolResult,
+    TurnBegin,
+    TurnEnd,
+)
+
+_current_turn_id = ContextVar[str | None]("current_turn_id", default=None)
+_terminal_tool_call_ids = ContextVar[set[str] | None]("terminal_tool_call_ids", default=None)
+
+
+def get_current_acp_tool_call_id_or_none() -> str | None:
+    """See `_ToolCallState.acp_tool_call_id`."""
+    from kimi_cli.soul.toolset import get_current_tool_call_or_none
+
+    turn_id = _current_turn_id.get()
+    if turn_id is None:
+        return None
+    tool_call = get_current_tool_call_or_none()
+    if tool_call is None:
+        return None
+    return f"{turn_id}/{tool_call.id}"
+
+
+def register_terminal_tool_call_id(tool_call_id: str) -> None:
+    calls = _terminal_tool_call_ids.get()
+    if calls is not None:
+        calls.add(tool_call_id)
+
+
+def should_hide_terminal_output(tool_call_id: str) -> bool:
+    calls = _terminal_tool_call_ids.get()
+    return calls is not None and tool_call_id in calls
+
+
+class _ToolCallState:
+    """Manages the state of a single tool call for streaming updates."""
+
+    def __init__(self, tool_call: ToolCall):
+        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)
+
+    @property
+    def acp_tool_call_id(self) -> str:
+        # 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 ensure the uniqueness by prefixing with the turn ID.
+        turn_id = _current_turn_id.get()
+        assert turn_id is not None
+        return f"{turn_id}/{self.tool_call.id}"
+
+    def append_args_part(self, args_part: str) -> None:
+        """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_key_argument(self.lexer, tool_name)
+        if subtitle:
+            return f"{tool_name}: {subtitle}"
+        return tool_name
+
+
+class _TurnState:
+    def __init__(self):
+        self.id = str(uuid.uuid4())
+        """Unique ID for the turn."""
+        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()
+
+
+class ACPSession:
+    def __init__(
+        self,
+        id: str,
+        cli: KimiCLI,
+        acp_conn: acp.Client,
+        kaos: Kaos | None = None,
+    ) -> None:
+        self._id = id
+        self._cli = cli
+        self._conn = acp_conn
+        self._kaos = kaos
+        self._turn_state: _TurnState | None = None
+
+    @property
+    def id(self) -> str:
+        """The ID of the ACP session."""
+        return self._id
+
+    @property
+    def cli(self) -> KimiCLI:
+        """The Kimi Code CLI instance bound to this ACP session."""
+        return self._cli
+
+    def _is_oauth_session(self) -> bool:
+        """Return True if the current session uses OAuth-based authentication."""
+        try:
+            llm = self._cli.soul.runtime.llm
+            return llm is not None and getattr(llm.provider_config, "oauth", None) is not None
+        except AttributeError:
+            return False
+
+    async def prompt(self, prompt: list[ACPContentBlock]) -> acp.PromptResponse:
+        user_input = acp_blocks_to_content_parts(prompt)
+        self._turn_state = _TurnState()
+        token = _current_turn_id.set(self._turn_state.id)
+        kaos_token = set_current_kaos(self._kaos) if self._kaos is not None else None
+        terminal_tool_calls_token = _terminal_tool_call_ids.set(set())
+        try:
+            async for msg in self._cli.run(user_input, self._turn_state.cancel_event):
+                match msg:
+                    case TurnBegin():
+                        pass
+                    case SteerInput():
+                        pass
+                    case TurnEnd():
+                        pass
+                    case StepBegin():
+                        pass
+                    case StepInterrupted():
+                        break
+                    case StepRetry():
+                        pass
+                    case CompactionBegin():
+                        pass
+                    case CompactionEnd():
+                        pass
+                    case MCPLoadingBegin():
+                        pass
+                    case MCPLoadingEnd():
+                        pass
+                    case StatusUpdate():
+                        pass
+                    case Notification():
+                        await self._send_notification(msg)
+                    case ThinkPart(think=think):
+                        await self._send_thinking(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__}]")
+                    case ToolCall():
+                        await self._send_tool_call(msg)
+                    case ToolCallPart():
+                        await self._send_tool_call_part(msg)
+                    case ToolResult():
+                        await self._send_tool_result(msg)
+                    case ApprovalResponse():
+                        pass
+                    case SubagentEvent():
+                        pass
+                    case PlanDisplay():
+                        pass
+                    case ApprovalRequest():
+                        await self._handle_approval_request(msg)
+                    case ToolCallRequest():
+                        logger.warning("Unexpected ToolCallRequest in ACP session: {msg}", msg=msg)
+                    case QuestionRequest():
+                        logger.warning(
+                            "QuestionRequest is unsupported in ACP session; resolving empty answer."
+                        )
+                        msg.resolve({})
+                    case _:
+                        pass
+        except LLMNotSet as e:
+            logger.exception("LLM not set:")
+            raise acp.RequestError.auth_required() from e
+        except LLMNotSupported as e:
+            logger.exception("LLM not supported:")
+            raise acp.RequestError.internal_error({"error": str(e)}) from e
+        except APIStatusError as e:
+            if e.status_code == 401 and self._is_oauth_session():
+                logger.warning("Authentication failed (401), prompting re-login")
+                raise acp.RequestError.auth_required() from e
+            logger.exception("LLM API status error:")
+            raise acp.RequestError.internal_error({"error": str(e)}) from e
+        except ChatProviderError as e:
+            logger.exception("LLM provider error:")
+            raise acp.RequestError.internal_error({"error": str(e)}) from e
+        except MaxStepsReached as e:
+            logger.warning("Max steps reached: {n_steps}", n_steps=e.n_steps)
+            return acp.PromptResponse(stop_reason="max_turn_requests")
+        except RunCancelled:
+            logger.info("Prompt cancelled by user")
+            return acp.PromptResponse(stop_reason="cancelled")
+        except Exception as e:
+            logger.exception("Unexpected error during prompt:")
+            raise acp.RequestError.internal_error({"error": str(e)}) from e
+        finally:
+            self._turn_state = None
+            if kaos_token is not None:
+                reset_current_kaos(kaos_token)
+            _terminal_tool_call_ids.reset(terminal_tool_calls_token)
+            _current_turn_id.reset(token)
+        return acp.PromptResponse(stop_reason="end_turn")
+
+    async def cancel(self) -> None:
+        if self._turn_state is None:
+            logger.warning("Cancel requested but no prompt is running")
+            return
+
+        self._turn_state.cancel_event.set()
+
+    async def _send_thinking(self, think: str):
+        """Send thinking content to client."""
+        if not self._id or not self._conn:
+            return
+
+        await self._conn.session_update(
+            self._id,
+            acp.schema.AgentThoughtChunk(
+                content=acp.schema.TextContentBlock(type="text", text=think),
+                session_update="agent_thought_chunk",
+            ),
+        )
+
+    async def _send_text(self, text: str):
+        """Send text chunk to client."""
+        if not self._id or not self._conn:
+            return
+
+        await self._conn.session_update(
+            session_id=self._id,
+            update=acp.schema.AgentMessageChunk(
+                content=acp.schema.TextContentBlock(type="text", text=text),
+                session_update="agent_message_chunk",
+            ),
+        )
+
+    async def _send_notification(self, notification: Notification):
+        """Send a system notification to the client as a text chunk."""
+        body = notification.body.strip()
+        text = f"[Notification] {notification.title}"
+        if body:
+            text = f"{text}\n{body}"
+        await self._send_text(text)
+
+    async def _send_tool_call(self, tool_call: ToolCall):
+        """Send tool call to client."""
+        assert self._turn_state is not None
+        if not self._id or not self._conn:
+            return
+
+        # Create and store tool call state
+        state = _ToolCallState(tool_call)
+        self._turn_state.tool_calls[tool_call.id] = state
+        self._turn_state.last_tool_call = state
+
+        await self._conn.session_update(
+            session_id=self._id,
+            update=acp.schema.ToolCallStart(
+                session_update="tool_call",
+                tool_call_id=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._turn_state is not None
+        if (
+            not self._id
+            or not self._conn
+            or not part.arguments_part
+            or self._turn_state.last_tool_call is None
+        ):
+            return
+
+        # Append new arguments part to the last tool call
+        self._turn_state.last_tool_call.append_args_part(part.arguments_part)
+
+        # Update the tool call with new content and title
+        update = acp.schema.ToolCallProgress(
+            session_update="tool_call_update",
+            tool_call_id=self._turn_state.last_tool_call.acp_tool_call_id,
+            title=self._turn_state.last_tool_call.get_title(),
+            status="in_progress",
+            content=[
+                acp.schema.ContentToolCallContent(
+                    type="content",
+                    content=acp.schema.TextContentBlock(
+                        type="text", text=self._turn_state.last_tool_call.args
+                    ),
+                )
+            ],
+        )
+
+        await self._conn.session_update(session_id=self._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._turn_state is not None
+        if not self._id or not self._conn:
+            return
+
+        tool_ret = result.return_value
+
+        state = self._turn_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(
+            session_update="tool_call_update",
+            tool_call_id=state.acp_tool_call_id,
+            status="failed" if tool_ret.is_error else "completed",
+        )
+
+        contents = (
+            []
+            if should_hide_terminal_output(state.acp_tool_call_id)
+            else tool_result_to_acp_content(tool_ret)
+        )
+        if contents:
+            update.content = contents
+
+        await self._conn.session_update(session_id=self._id, update=update)
+        logger.debug("Sent tool result: {id}", id=result.tool_call_id)
+
+        for block in tool_ret.display:
+            if isinstance(block, TodoDisplayBlock):
+                await self._send_plan_update(block)
+
+    async def _handle_approval_request(self, request: ApprovalRequest):
+        """Handle approval request by sending permission request to client."""
+        assert self._turn_state is not None
+        if not self._id or not self._conn:
+            logger.warning("No session ID, auto-rejecting approval request")
+            request.resolve("reject")
+            return
+
+        state = self._turn_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("reject")
+            return
+
+        try:
+            content: list[
+                acp.schema.ContentToolCallContent
+                | acp.schema.FileEditToolCallContent
+                | acp.schema.TerminalToolCallContent
+            ] = []
+            if request.display:
+                for block in request.display:
+                    diff_content = display_block_to_acp_content(block)
+                    if diff_content is not None:
+                        content.append(diff_content)
+            if not content:
+                content.append(
+                    acp.schema.ContentToolCallContent(
+                        type="content",
+                        content=acp.schema.TextContentBlock(
+                            type="text",
+                            text=f"Requesting approval to perform: {request.description}",
+                        ),
+                    )
+                )
+
+            # Send permission request and wait for response
+            logger.debug("Requesting permission for action: {action}", action=request.action)
+            response = await self._conn.request_permission(
+                [
+                    acp.schema.PermissionOption(
+                        option_id="approve",
+                        name="Approve once",
+                        kind="allow_once",
+                    ),
+                    acp.schema.PermissionOption(
+                        option_id="approve_for_session",
+                        name="Approve for this session",
+                        kind="allow_always",
+                    ),
+                    acp.schema.PermissionOption(
+                        option_id="reject",
+                        name="Reject",
+                        kind="reject_once",
+                    ),
+                ],
+                self._id,
+                acp.schema.ToolCallUpdate(
+                    tool_call_id=state.acp_tool_call_id,
+                    title=state.get_title(),
+                    content=content,
+                ),
+            )
+            logger.debug("Received permission response: {response}", response=response)
+
+            # Process the outcome
+            if isinstance(response.outcome, acp.schema.AllowedOutcome):
+                # selected
+                option_id = response.outcome.option_id
+                if option_id == "approve":
+                    logger.debug("Permission granted for: {action}", action=request.action)
+                    request.resolve("approve")
+                elif option_id == "approve_for_session":
+                    logger.debug("Permission granted for session: {action}", action=request.action)
+                    request.resolve("approve_for_session")
+                else:
+                    logger.debug("Permission denied for: {action}", action=request.action)
+                    request.resolve("reject")
+            else:
+                # cancelled
+                logger.debug("Permission request cancelled for: {action}", action=request.action)
+                request.resolve("reject")
+        except Exception:
+            logger.exception("Error handling approval request:")
+            # On error, reject the request
+            request.resolve("reject")
+
+    async def _send_plan_update(self, block: TodoDisplayBlock) -> None:
+        """Send todo list updates as ACP agent plan updates."""
+
+        status_map: dict[str, acp.schema.PlanEntryStatus] = {
+            "pending": "pending",
+            "in progress": "in_progress",
+            "in_progress": "in_progress",
+            "done": "completed",
+            "completed": "completed",
+        }
+        entries: list[acp.schema.PlanEntry] = [
+            acp.schema.PlanEntry(
+                content=todo.title,
+                priority="medium",
+                status=status_map.get(todo.status.lower(), "pending"),
+            )
+            for todo in block.items
+            if todo.title
+        ]
+
+        if not entries:
+            logger.warning("No valid todo items to send in plan update: {todos}", todos=block.items)
+            return
+
+        await self._conn.session_update(
+            session_id=self._id,
+            update=acp.schema.AgentPlanUpdate(session_update="plan", entries=entries),
+        )

kimi_cli/acp/tools.py

@@ -0,0 +1,167 @@
+import asyncio
+from contextlib import suppress
+
+import acp
+from kaos import get_current_kaos
+from kaos.local import local_kaos
+from kosong.tooling import CallableTool2, ToolReturnValue
+
+from kimi_cli.soul.agent import Runtime
+from kimi_cli.soul.approval import Approval
+from kimi_cli.soul.toolset import KimiToolset
+from kimi_cli.tools.shell import Params as ShellParams
+from kimi_cli.tools.shell import Shell
+from kimi_cli.tools.utils import ToolResultBuilder
+from kimi_cli.wire.types import DisplayBlock
+
+
+def replace_tools(
+    client_capabilities: acp.schema.ClientCapabilities,
+    acp_conn: acp.Client,
+    acp_session_id: str,
+    toolset: KimiToolset,
+    runtime: Runtime,
+) -> None:
+    current_kaos = get_current_kaos().name
+    if current_kaos not in (local_kaos.name, "acp"):
+        # Only replace tools when running locally or under ACPKaos.
+        return
+
+    if client_capabilities.terminal and (shell_tool := toolset.find(Shell)):
+        # Replace the Shell tool with the ACP Terminal tool if supported.
+        toolset.add(
+            Terminal(
+                shell_tool,
+                acp_conn,
+                acp_session_id,
+                runtime.approval,
+            )
+        )
+
+
+class HideOutputDisplayBlock(DisplayBlock):
+    """A special DisplayBlock that indicates output should be hidden in ACP clients."""
+
+    type: str = "acp/hide_output"
+
+
+class Terminal(CallableTool2[ShellParams]):
+    def __init__(
+        self,
+        shell_tool: Shell,
+        acp_conn: acp.Client,
+        acp_session_id: str,
+        approval: Approval,
+    ) -> None:
+        # Use the `name`, `description`, and `params` from the existing Shell tool,
+        # so that when this is added to the toolset, it replaces the original Shell tool.
+        super().__init__(shell_tool.name, shell_tool.description, shell_tool.params)
+        self._acp_conn = acp_conn
+        self._acp_session_id = acp_session_id
+        self._approval = approval
+
+    async def __call__(self, params: ShellParams) -> ToolReturnValue:
+        from kimi_cli.acp.session import get_current_acp_tool_call_id_or_none
+
+        builder = ToolResultBuilder()
+        # Hide tool output because we use `TerminalToolCallContent` which already streams output
+        # directly to the user.
+        builder.display(HideOutputDisplayBlock())
+
+        if not params.command:
+            return builder.error("Command cannot be empty.", brief="Empty command")
+
+        approval_result = await self._approval.request(
+            self.name,
+            "run shell command",
+            f"Run command `{params.command}`",
+        )
+        if not approval_result:
+            return approval_result.rejection_error()
+
+        timeout_seconds = float(params.timeout)
+        timeout_label = f"{timeout_seconds:g}s"
+        terminal_id: str | None = None
+        exit_status: (
+            acp.schema.WaitForTerminalExitResponse | acp.schema.TerminalExitStatus | None
+        ) = None
+        timed_out = False
+
+        try:
+            resp = await self._acp_conn.create_terminal(
+                command=params.command,
+                session_id=self._acp_session_id,
+                output_byte_limit=builder.max_chars,
+            )
+            terminal_id = resp.terminal_id
+
+            acp_tool_call_id = get_current_acp_tool_call_id_or_none()
+            assert acp_tool_call_id, "Expected to have an ACP tool call ID in context"
+            await self._acp_conn.session_update(
+                session_id=self._acp_session_id,
+                update=acp.schema.ToolCallProgress(
+                    session_update="tool_call_update",
+                    tool_call_id=acp_tool_call_id,
+                    status="in_progress",
+                    content=[
+                        acp.schema.TerminalToolCallContent(
+                            type="terminal",
+                            terminal_id=terminal_id,
+                        )
+                    ],
+                ),
+            )
+
+            try:
+                async with asyncio.timeout(timeout_seconds):
+                    exit_status = await self._acp_conn.wait_for_terminal_exit(
+                        session_id=self._acp_session_id,
+                        terminal_id=terminal_id,
+                    )
+            except TimeoutError:
+                timed_out = True
+                await self._acp_conn.kill_terminal(
+                    session_id=self._acp_session_id,
+                    terminal_id=terminal_id,
+                )
+
+            output_response = await self._acp_conn.terminal_output(
+                session_id=self._acp_session_id,
+                terminal_id=terminal_id,
+            )
+            builder.write(output_response.output)
+            if output_response.exit_status:
+                exit_status = output_response.exit_status
+
+            exit_code = exit_status.exit_code if exit_status else None
+            exit_signal = exit_status.signal if exit_status else None
+
+            truncated_note = (
+                " Output was truncated by the client output limit."
+                if output_response.truncated
+                else ""
+            )
+
+            if timed_out:
+                return builder.error(
+                    f"Command killed by timeout ({timeout_label}){truncated_note}",
+                    brief=f"Killed by timeout ({timeout_label})",
+                )
+            if exit_signal:
+                return builder.error(
+                    f"Command terminated by signal: {exit_signal}.{truncated_note}",
+                    brief=f"Signal: {exit_signal}",
+                )
+            if exit_code not in (None, 0):
+                return builder.error(
+                    f"Command failed with exit code: {exit_code}.{truncated_note}",
+                    brief=f"Failed with exit code: {exit_code}",
+                )
+            return builder.ok(f"Command executed successfully.{truncated_note}")
+        finally:
+            if terminal_id is not None:
+                with suppress(Exception):
+                    await self._acp_conn.release_terminal(
+                        session_id=self._acp_session_id,
+                        terminal_id=terminal_id,
+                    )

kimi_cli/acp/types.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+import acp
+
+MCPServer = acp.schema.HttpMcpServer | acp.schema.SseMcpServer | acp.schema.McpServerStdio
+
+ACPContentBlock = (
+    acp.schema.TextContentBlock
+    | acp.schema.ImageContentBlock
+    | acp.schema.AudioContentBlock
+    | acp.schema.ResourceContentBlock
+    | acp.schema.EmbeddedResourceContentBlock
+)