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

kimi_cli/ui/shell/visualize.py

@@ -1,46 +1,58 @@
+from __future__ import annotations
+
 import asyncio
+from collections import deque
+from collections.abc import Callable
 from contextlib import asynccontextmanager, suppress
+from typing import NamedTuple
 
-from kosong.base.message import ContentPart, TextPart, ToolCall, ToolCallPart
-from kosong.tooling import ToolResult
+import streamingjson  # type: ignore[reportMissingTypeStubs]
+from kosong.message import Message
+from kosong.tooling import ToolError, ToolOk
+from rich.console import Group, RenderableType
+from rich.live import Live
+from rich.markup import escape
+from rich.panel import Panel
+from rich.spinner import Spinner
+from rich.text import Text
 
-from kimi_cli.soul import StatusSnapshot
+from kimi_cli.tools import extract_key_argument
 from kimi_cli.ui.shell.console import console
-from kimi_cli.ui.shell.keyboard import listen_for_keyboard
-from kimi_cli.ui.shell.liveview import StepLiveView, StepLiveViewWithMarkdown
+from kimi_cli.ui.shell.keyboard import KeyEvent, listen_for_keyboard
+from kimi_cli.utils.aioqueue import QueueShutDown
+from kimi_cli.utils.message import message_stringify
+from kimi_cli.utils.rich.columns import BulletColumns
+from kimi_cli.utils.rich.markdown import Markdown
 from kimi_cli.wire import WireUISide
-from kimi_cli.wire.message import (
+from kimi_cli.wire.types import (
     ApprovalRequest,
+    ApprovalRequestResolved,
+    BriefDisplayBlock,
     CompactionBegin,
     CompactionEnd,
+    ContentPart,
     StatusUpdate,
     StepBegin,
     StepInterrupted,
+    SubagentEvent,
+    TextPart,
+    ThinkPart,
+    TodoDisplayBlock,
+    ToolCall,
+    ToolCallPart,
+    ToolResult,
+    ToolReturnValue,
+    TurnBegin,
+    WireMessage,
 )
 
-
-@asynccontextmanager
-async def _keyboard_listener(step: StepLiveView):
-    async def _keyboard():
-        try:
-            async for event in listen_for_keyboard():
-                step.handle_keyboard_event(event)
-        except asyncio.CancelledError:
-            return
-
-    task = asyncio.create_task(_keyboard())
-    try:
-        yield
-    finally:
-        task.cancel()
-        with suppress(asyncio.CancelledError):
-            await task
+MAX_SUBAGENT_TOOL_CALLS_TO_SHOW = 4
 
 
 async def visualize(
     wire: WireUISide,
     *,
-    initial_status: StatusSnapshot,
+    initial_status: StatusUpdate,
     cancel_event: asyncio.Event | None = None,
 ):
     """
@@ -51,61 +63,551 @@ async def visualize(
         initial_status: Initial status snapshot
         cancel_event: Event that can be set (e.g., by ESC key) to cancel the run
     """
+    view = _LiveView(initial_status, cancel_event)
+    await view.visualize_loop(wire)
+
 
-    latest_status = initial_status
+class _ContentBlock:
+    def __init__(self, is_think: bool):
+        self.is_think = is_think
+        self._spinner = Spinner("dots", "Thinking..." if is_think else "Composing...")
+        self.raw_text = ""
 
-    # expect a StepBegin
-    assert isinstance(await wire.receive(), StepBegin)
+    def compose(self) -> RenderableType:
+        return self._spinner
 
-    while True:
-        # TODO: Maybe we can always have a StepLiveView here.
-        #       No need to recreate for each step.
-        with StepLiveViewWithMarkdown(latest_status, cancel_event) as step:
-            async with _keyboard_listener(step):
-                # spin the moon at the beginning of each step
-                with console.status("", spinner="moon"):
-                    msg = await wire.receive()
+    def compose_final(self) -> RenderableType:
+        return BulletColumns(
+            Markdown(
+                self.raw_text,
+                style="grey50 italic" if self.is_think else "",
+            ),
+            bullet_style="grey50" if self.is_think else None,
+        )
 
-                if isinstance(msg, CompactionBegin):
-                    with console.status("[cyan]Compacting...[/cyan]"):
+    def append(self, content: str) -> None:
+        self.raw_text += content
+
+
+class _ToolCallBlock:
+    class FinishedSubCall(NamedTuple):
+        call: ToolCall
+        result: ToolReturnValue
+
+    def __init__(self, tool_call: ToolCall):
+        self._tool_name = tool_call.function.name
+        self._lexer = streamingjson.Lexer()
+        if tool_call.function.arguments is not None:
+            self._lexer.append_string(tool_call.function.arguments)
+
+        self._argument = extract_key_argument(self._lexer, self._tool_name)
+        self._result: ToolReturnValue | None = None
+
+        self._ongoing_subagent_tool_calls: dict[str, ToolCall] = {}
+        self._last_subagent_tool_call: ToolCall | None = None
+        self._n_finished_subagent_tool_calls = 0
+        self._finished_subagent_tool_calls = deque[_ToolCallBlock.FinishedSubCall](
+            maxlen=MAX_SUBAGENT_TOOL_CALLS_TO_SHOW
+        )
+
+        self._spinning_dots = Spinner("dots", text="")
+        self._renderable: RenderableType = self._compose()
+
+    def compose(self) -> RenderableType:
+        return self._renderable
+
+    @property
+    def finished(self) -> bool:
+        return self._result is not None
+
+    def append_args_part(self, args_part: str):
+        if self.finished:
+            return
+        self._lexer.append_string(args_part)
+        # TODO: maybe don't extract detail if it's already stable
+        argument = extract_key_argument(self._lexer, self._tool_name)
+        if argument and argument != self._argument:
+            self._argument = argument
+            self._renderable = BulletColumns(
+                Text.from_markup(self._get_headline_markup()),
+                bullet=self._spinning_dots,
+            )
+
+    def finish(self, result: ToolReturnValue):
+        self._result = result
+        self._renderable = self._compose()
+
+    def append_sub_tool_call(self, tool_call: ToolCall):
+        self._ongoing_subagent_tool_calls[tool_call.id] = tool_call
+        self._last_subagent_tool_call = tool_call
+
+    def append_sub_tool_call_part(self, tool_call_part: ToolCallPart):
+        if self._last_subagent_tool_call is None:
+            return
+        if not tool_call_part.arguments_part:
+            return
+        if self._last_subagent_tool_call.function.arguments is None:
+            self._last_subagent_tool_call.function.arguments = tool_call_part.arguments_part
+        else:
+            self._last_subagent_tool_call.function.arguments += tool_call_part.arguments_part
+
+    def finish_sub_tool_call(self, tool_result: ToolResult):
+        self._last_subagent_tool_call = None
+        sub_tool_call = self._ongoing_subagent_tool_calls.pop(tool_result.tool_call_id, None)
+        if sub_tool_call is None:
+            return
+
+        self._finished_subagent_tool_calls.append(
+            _ToolCallBlock.FinishedSubCall(
+                call=sub_tool_call,
+                result=tool_result.return_value,
+            )
+        )
+        self._n_finished_subagent_tool_calls += 1
+        self._renderable = self._compose()
+
+    def _compose(self) -> RenderableType:
+        lines: list[RenderableType] = [
+            Text.from_markup(self._get_headline_markup()),
+        ]
+
+        if self._n_finished_subagent_tool_calls > MAX_SUBAGENT_TOOL_CALLS_TO_SHOW:
+            n_hidden = self._n_finished_subagent_tool_calls - MAX_SUBAGENT_TOOL_CALLS_TO_SHOW
+            lines.append(
+                BulletColumns(
+                    Text(
+                        f"{n_hidden} more tool call{'s' if n_hidden > 1 else ''} ...",
+                        style="grey50 italic",
+                    ),
+                    bullet_style="grey50",
+                )
+            )
+        for sub_call, sub_result in self._finished_subagent_tool_calls:
+            argument = extract_key_argument(
+                sub_call.function.arguments or "", sub_call.function.name
+            )
+            lines.append(
+                BulletColumns(
+                    Text.from_markup(
+                        f"Used [blue]{sub_call.function.name}[/blue]"
+                        + (f" [grey50]({argument})[/grey50]" if argument else "")
+                    ),
+                    bullet_style="green" if not sub_result.is_error else "red",
+                )
+            )
+
+        if self._result is not None:
+            for block in self._result.display:
+                if isinstance(block, BriefDisplayBlock):
+                    style = "grey50" if not self._result.is_error else "red"
+                    if block.text:
+                        lines.append(Markdown(block.text, style=style))
+                elif isinstance(block, TodoDisplayBlock):
+                    markdown = self._render_todo_markdown(block)
+                    if markdown:
+                        lines.append(Markdown(markdown, style="grey50"))
+
+        if self.finished:
+            assert self._result is not None
+            return BulletColumns(
+                Group(*lines),
+                bullet_style="green" if not self._result.is_error else "red",
+            )
+        else:
+            return BulletColumns(
+                Group(*lines),
+                bullet=self._spinning_dots,
+            )
+
+    def _get_headline_markup(self) -> str:
+        return f"{'Used' if self.finished else 'Using'} [blue]{self._tool_name}[/blue]" + (
+            f" [grey50]({escape(self._argument)})[/grey50]" if self._argument else ""
+        )
+
+    def _render_todo_markdown(self, block: TodoDisplayBlock) -> str:
+        lines: list[str] = []
+        for todo in block.items:
+            normalized = todo.status.replace("_", " ").lower()
+            match normalized:
+                case "pending":
+                    lines.append(f"- {todo.title}")
+                case "in progress":
+                    lines.append(f"- {todo.title} ←")
+                case "done":
+                    lines.append(f"- ~~{todo.title}~~")
+                case _:
+                    lines.append(f"- {todo.title}")
+        return "\n".join(lines)
+
+
+class _ApprovalRequestPanel:
+    def __init__(self, request: ApprovalRequest):
+        self.request = request
+        self.options: list[tuple[str, ApprovalRequest.Response]] = [
+            ("Approve once", "approve"),
+            ("Approve for this session", "approve_for_session"),
+            ("Reject, tell Kimi CLI what to do instead", "reject"),
+        ]
+        self.selected_index = 0
+
+    def render(self) -> RenderableType:
+        """Render the approval menu as a panel."""
+        lines: list[RenderableType] = []
+
+        # Add request details
+        lines.append(
+            Text.assemble(
+                Text.from_markup(f"[blue]{self.request.sender}[/blue]"),
+                Text(f' is requesting approval to "{self.request.description}".'),
+            )
+        )
+
+        lines.append(Text(""))  # Empty line
+
+        # Add menu options
+        for i, (option_text, _) in enumerate(self.options):
+            if i == self.selected_index:
+                lines.append(Text(f"→ {option_text}", style="cyan"))
+            else:
+                lines.append(Text(f"  {option_text}", style="grey50"))
+
+        content = Group(*lines)
+        return Panel.fit(
+            content,
+            title="[yellow]⚠ Approval Requested[/yellow]",
+            border_style="yellow",
+            padding=(1, 2),
+        )
+
+    def move_up(self):
+        """Move selection up."""
+        self.selected_index = (self.selected_index - 1) % len(self.options)
+
+    def move_down(self):
+        """Move selection down."""
+        self.selected_index = (self.selected_index + 1) % len(self.options)
+
+    def get_selected_response(self) -> ApprovalRequest.Response:
+        """Get the approval response based on selected option."""
+        return self.options[self.selected_index][1]
+
+
+class _StatusBlock:
+    def __init__(self, initial: StatusUpdate) -> None:
+        self.text = Text("", justify="right")
+        self.update(initial)
+
+    def render(self) -> RenderableType:
+        return self.text
+
+    def update(self, status: StatusUpdate) -> None:
+        if status.context_usage is not None:
+            self.text.plain = f"context: {status.context_usage:.1%}"
+
+
+@asynccontextmanager
+async def _keyboard_listener(handler: Callable[[KeyEvent], None]):
+    async def _keyboard():
+        async for event in listen_for_keyboard():
+            handler(event)
+
+    task = asyncio.create_task(_keyboard())
+    try:
+        yield
+    finally:
+        task.cancel()
+        with suppress(asyncio.CancelledError):
+            await task
+
+
+class _LiveView:
+    def __init__(self, initial_status: StatusUpdate, cancel_event: asyncio.Event | None = None):
+        self._cancel_event = cancel_event
+
+        self._mooning_spinner: Spinner | None = None
+        self._compacting_spinner: Spinner | None = None
+
+        self._current_content_block: _ContentBlock | None = None
+        self._tool_call_blocks: dict[str, _ToolCallBlock] = {}
+        self._last_tool_call_block: _ToolCallBlock | None = None
+        self._approval_request_queue = deque[ApprovalRequest]()
+        """
+        It is possible that multiple subagents request approvals at the same time,
+        in which case we will have to queue them up and show them one by one.
+        """
+        self._current_approval_request_panel: _ApprovalRequestPanel | None = None
+        self._reject_all_following = False
+        self._status_block = _StatusBlock(initial_status)
+
+        self._need_recompose = False
+
+    async def visualize_loop(self, wire: WireUISide):
+        with Live(
+            self.compose(),
+            console=console,
+            refresh_per_second=10,
+            transient=True,
+            vertical_overflow="visible",
+        ) as live:
+
+            def keyboard_handler(event: KeyEvent) -> None:
+                self.dispatch_keyboard_event(event)
+                if self._need_recompose:
+                    live.update(self.compose())
+                    self._need_recompose = False
+
+            async with _keyboard_listener(keyboard_handler):
+                while True:
+                    try:
                         msg = await wire.receive()
+                    except QueueShutDown:
+                        self.cleanup(is_interrupt=False)
+                        live.update(self.compose())
+                        break
+
                     if isinstance(msg, StepInterrupted):
+                        self.cleanup(is_interrupt=True)
+                        live.update(self.compose())
                         break
-                    assert isinstance(msg, CompactionEnd)
-                    continue
 
-                # visualization loop for one step
-                while True:
-                    match msg:
-                        case TextPart(text=text):
-                            step.append_text(text)
-                        case ContentPart():
-                            # TODO: support more content parts
-                            step.append_text(f"[{msg.__class__.__name__}]")
-                        case ToolCall():
-                            step.append_tool_call(msg)
-                        case ToolCallPart():
-                            step.append_tool_call_part(msg)
-                        case ToolResult():
-                            step.append_tool_result(msg)
-                        case ApprovalRequest():
-                            step.request_approval(msg)
-                        case StatusUpdate(status=status):
-                            latest_status = status
-                            step.update_status(latest_status)
-                        case _:
-                            break  # break the step loop
-                    msg = await wire.receive()
-
-                # cleanup the step live view
-                if isinstance(msg, StepInterrupted):
-                    step.interrupt()
-                else:
-                    step.finish()
-
-        if isinstance(msg, StepInterrupted):
-            # for StepInterrupted, the visualization loop should end immediately
+                    self.dispatch_wire_message(msg)
+                    if self._need_recompose:
+                        live.update(self.compose())
+                        self._need_recompose = False
+
+    def refresh_soon(self) -> None:
+        self._need_recompose = True
+
+    def compose(self) -> RenderableType:
+        """Compose the live view display content."""
+        blocks: list[RenderableType] = []
+        if self._mooning_spinner is not None:
+            blocks.append(self._mooning_spinner)
+        elif self._compacting_spinner is not None:
+            blocks.append(self._compacting_spinner)
+        else:
+            if self._current_content_block is not None:
+                blocks.append(self._current_content_block.compose())
+            for tool_call in self._tool_call_blocks.values():
+                blocks.append(tool_call.compose())
+        if self._current_approval_request_panel:
+            blocks.append(self._current_approval_request_panel.render())
+        blocks.append(self._status_block.render())
+        return Group(*blocks)
+
+    def dispatch_wire_message(self, msg: WireMessage) -> None:
+        """Dispatch the Wire message to UI components."""
+        assert not isinstance(msg, StepInterrupted)  # handled in visualize_loop
+
+        if isinstance(msg, StepBegin):
+            self.cleanup(is_interrupt=False)
+            self._mooning_spinner = Spinner("moon", "")
+            self.refresh_soon()
+            return
+
+        if self._mooning_spinner is not None:
+            # any message other than StepBegin should end the mooning state
+            self._mooning_spinner = None
+            self.refresh_soon()
+
+        match msg:
+            case TurnBegin():
+                self.flush_content()
+                console.print(
+                    Panel(
+                        Text(message_stringify(Message(role="user", content=msg.user_input))),
+                        padding=(0, 1),
+                    )
+                )
+            case CompactionBegin():
+                self._compacting_spinner = Spinner("balloon", "Compacting...")
+                self.refresh_soon()
+            case CompactionEnd():
+                self._compacting_spinner = None
+                self.refresh_soon()
+            case StatusUpdate():
+                self._status_block.update(msg)
+            case ContentPart():
+                self.append_content(msg)
+            case ToolCall():
+                self.append_tool_call(msg)
+            case ToolCallPart():
+                self.append_tool_call_part(msg)
+            case ToolResult():
+                self.append_tool_result(msg)
+            case SubagentEvent():
+                self.handle_subagent_event(msg)
+            case ApprovalRequestResolved():
+                # we don't need to handle this because the request is resolved on UI
+                pass
+            case ApprovalRequest():
+                self.request_approval(msg)
+
+    def dispatch_keyboard_event(self, event: KeyEvent) -> None:
+        # handle ESC key to cancel the run
+        if event == KeyEvent.ESCAPE and self._cancel_event is not None:
+            self._cancel_event.set()
+            return
+
+        if not self._current_approval_request_panel:
+            # just ignore any keyboard event when there's no approval request
+            return
+
+        match event:
+            case KeyEvent.UP:
+                self._current_approval_request_panel.move_up()
+                self.refresh_soon()
+            case KeyEvent.DOWN:
+                self._current_approval_request_panel.move_down()
+                self.refresh_soon()
+            case KeyEvent.ENTER:
+                resp = self._current_approval_request_panel.get_selected_response()
+                self._current_approval_request_panel.request.resolve(resp)
+                if resp == "approve_for_session":
+                    to_remove_from_queue: list[ApprovalRequest] = []
+                    for request in self._approval_request_queue:
+                        # approve all queued requests with the same action
+                        if request.action == self._current_approval_request_panel.request.action:
+                            request.resolve("approve_for_session")
+                            to_remove_from_queue.append(request)
+                    for request in to_remove_from_queue:
+                        self._approval_request_queue.remove(request)
+                elif resp == "reject":
+                    # one rejection should stop the step immediately
+                    while self._approval_request_queue:
+                        self._approval_request_queue.popleft().resolve("reject")
+                    self._reject_all_following = True
+                self.show_next_approval_request()
+            case _:
+                # just ignore any other keyboard event
+                return
+
+    def cleanup(self, is_interrupt: bool) -> None:
+        """Cleanup the live view on step end or interruption."""
+        self.flush_content()
+
+        for block in self._tool_call_blocks.values():
+            if not block.finished:
+                # this should not happen, but just in case
+                block.finish(
+                    ToolError(message="", brief="Interrupted")
+                    if is_interrupt
+                    else ToolOk(output="")
+                )
+        self._last_tool_call_block = None
+        self.flush_finished_tool_calls()
+
+        while self._approval_request_queue:
+            # should not happen, but just in case
+            self._approval_request_queue.popleft().resolve("reject")
+        self._current_approval_request_panel = None
+        self._reject_all_following = False
+
+    def flush_content(self) -> None:
+        """Flush the current content block."""
+        if self._current_content_block is not None:
+            console.print(self._current_content_block.compose_final())
+            self._current_content_block = None
+            self.refresh_soon()
+
+    def flush_finished_tool_calls(self) -> None:
+        """Flush all leading finished tool call blocks."""
+        tool_call_ids = list(self._tool_call_blocks.keys())
+        for tool_call_id in tool_call_ids:
+            block = self._tool_call_blocks[tool_call_id]
+            if not block.finished:
+                break
+
+            self._tool_call_blocks.pop(tool_call_id)
+            console.print(block.compose())
+            if self._last_tool_call_block == block:
+                self._last_tool_call_block = None
+            self.refresh_soon()
+
+    def append_content(self, part: ContentPart) -> None:
+        match part:
+            case ThinkPart(think=text) | TextPart(text=text):
+                if not text:
+                    return
+                is_think = isinstance(part, ThinkPart)
+                if self._current_content_block is None:
+                    self._current_content_block = _ContentBlock(is_think)
+                    self.refresh_soon()
+                elif self._current_content_block.is_think != is_think:
+                    self.flush_content()
+                    self._current_content_block = _ContentBlock(is_think)
+                    self.refresh_soon()
+                self._current_content_block.append(text)
+            case _:
+                # TODO: support more content part types
+                pass
+
+    def append_tool_call(self, tool_call: ToolCall) -> None:
+        self.flush_content()
+        self._tool_call_blocks[tool_call.id] = _ToolCallBlock(tool_call)
+        self._last_tool_call_block = self._tool_call_blocks[tool_call.id]
+        self.refresh_soon()
+
+    def append_tool_call_part(self, part: ToolCallPart) -> None:
+        if not part.arguments_part:
+            return
+        if self._last_tool_call_block is None:
+            return
+        self._last_tool_call_block.append_args_part(part.arguments_part)
+        self.refresh_soon()
+
+    def append_tool_result(self, result: ToolResult) -> None:
+        if block := self._tool_call_blocks.get(result.tool_call_id):
+            block.finish(result.return_value)
+            self.flush_finished_tool_calls()
+            self.refresh_soon()
+
+    def request_approval(self, request: ApprovalRequest) -> None:
+        # If we're rejecting all following requests, reject immediately
+        if self._reject_all_following:
+            request.resolve("reject")
+            return
+
+        self._approval_request_queue.append(request)
+
+        if self._current_approval_request_panel is None:
+            console.bell()
+            self.show_next_approval_request()
+
+    def show_next_approval_request(self) -> None:
+        """
+        Show the next approval request from the queue.
+        If there are no pending requests, clear the current approval panel.
+        """
+        if not self._approval_request_queue:
+            if self._current_approval_request_panel is not None:
+                self._current_approval_request_panel = None
+                self.refresh_soon()
+            return
+
+        while self._approval_request_queue:
+            request = self._approval_request_queue.popleft()
+            if request.resolved:
+                # skip resolved requests
+                continue
+            self._current_approval_request_panel = _ApprovalRequestPanel(request)
+            self.refresh_soon()
             break
 
-        assert isinstance(msg, StepBegin), "expect a StepBegin"
-        # start a new step
+    def handle_subagent_event(self, event: SubagentEvent) -> None:
+        block = self._tool_call_blocks.get(event.task_tool_call_id)
+        if block is None:
+            return
+
+        match event.event:
+            case ToolCall() as tool_call:
+                block.append_sub_tool_call(tool_call)
+            case ToolCallPart() as tool_call_part:
+                block.append_sub_tool_call_part(tool_call_part)
+            case ToolResult() as tool_result:
+                block.finish_sub_tool_call(tool_result)
+                self.refresh_soon()
+            case _:
+                # ignore other events for now
+                # TODO: may need to handle multi-level nested subagents
+                pass