klaude-code 1.9.0__py3-none-any.whl → 2.0.0__py3-none-any.whl

klaude_code/ui/modes/repl/event_handler.py

@@ -1,16 +1,22 @@
 from __future__ import annotations
 
+from collections.abc import Callable
 from dataclasses import dataclass
 
 from rich.rule import Rule
 from rich.text import Text
 
-from klaude_code import const
+from klaude_code.const import MARKDOWN_LEFT_MARGIN, MARKDOWN_STREAM_LIVE_REPAINT_ENABLED, STATUS_DEFAULT_TEXT
 from klaude_code.protocol import events
 from klaude_code.ui.core.stage_manager import Stage, StageManager
 from klaude_code.ui.modes.repl.renderer import REPLRenderer
 from klaude_code.ui.renderers.assistant import ASSISTANT_MESSAGE_MARK
-from klaude_code.ui.renderers.thinking import THINKING_MESSAGE_MARK, normalize_thinking_content
+from klaude_code.ui.renderers.thinking import (
+    THINKING_MESSAGE_MARK,
+    extract_last_bold_header,
+    normalize_thinking_content,
+)
+from klaude_code.ui.renderers.tools import get_tool_active_form
 from klaude_code.ui.rich import status as r_status
 from klaude_code.ui.rich.markdown import MarkdownStream, ThinkingMarkdown
 from klaude_code.ui.rich.theme import ThemeKey
@@ -18,39 +24,24 @@ from klaude_code.ui.terminal.notifier import Notification, NotificationType, Ter
 from klaude_code.ui.terminal.progress_bar import OSC94States, emit_osc94
 
 
-def extract_last_bold_header(text: str) -> str | None:
-    """Extract the latest complete bold header ("**...**") from text.
-
-    We treat a bold segment as a "header" only if it appears at the beginning
-    of a line (ignoring leading whitespace). This avoids picking up incidental
-    emphasis inside paragraphs.
-
-    Returns None if no complete bold segment is available yet.
-    """
-
-    last: str | None = None
-    i = 0
-    while True:
-        start = text.find("**", i)
-        if start < 0:
-            break
-
-        line_start = text.rfind("\n", 0, start) + 1
-        if text[line_start:start].strip():
-            i = start + 2
-            continue
-
-        end = text.find("**", start + 2)
-        if end < 0:
-            break
+@dataclass
+class SubAgentThinkingHeaderState:
+    buffer: str = ""
+    last_header: str | None = None
 
-        inner = " ".join(text[start + 2 : end].split())
-        if inner and "\n" not in inner:
-            last = inner
+    def append_and_extract_new_header(self, content: str) -> str | None:
+        self.buffer += content
 
-        i = end + 2
+        # Sub-agent thinking does not need full streaming; keep a bounded tail.
+        max_chars = 8192
+        if len(self.buffer) > max_chars:
+            self.buffer = self.buffer[-max_chars:]
 
-    return last
+        header = extract_last_bold_header(normalize_thinking_content(self.buffer))
+        if header and header != self.last_header:
+            self.last_header = header
+            return header
+        return None
 
 
 @dataclass
@@ -107,6 +98,31 @@ class StreamState:
         """End the current streaming session."""
         self._active = None
 
+    def render(self, *, transform: Callable[[str], str] | None = None, final: bool = False) -> bool:
+        """Render the current buffer to the markdown stream.
+
+        Returns:
+            bool: True if an active stream was rendered.
+        """
+
+        if self._active is None:
+            return False
+
+        text = self._active.buffer
+        if transform is not None:
+            text = transform(text)
+        self._active.mdstream.update(text, final=final)
+
+        if final:
+            self.finish()
+
+        return True
+
+    def finalize(self, *, transform: Callable[[str], str] | None = None) -> bool:
+        """Finalize rendering and end the current streaming session."""
+
+        return self.render(transform=transform, final=True)
+
 
 class ActivityState:
     """Represents the current activity state for spinner display.
@@ -185,7 +201,7 @@ class SpinnerStatusState:
     - context_percent: Context usage percentage, updated during task execution
 
     Display logic:
-    - If activity: show base + activity (if base exists) or activity + "..."
+    - If activity: show base + activity (if base exists) or activity + "…"
     - Elif base_status: show base_status
     - Else: show "Thinking …"
     - Context percent is appended at the end if available
@@ -260,7 +276,7 @@ class SpinnerStatusState:
             activity_text.append(" …")
             result = activity_text
         else:
-            result = Text(const.STATUS_DEFAULT_TEXT, style=ThemeKey.STATUS_TEXT)
+            result = Text(STATUS_DEFAULT_TEXT, style=ThemeKey.STATUS_TEXT)
 
         return result
 
@@ -295,6 +311,7 @@ class DisplayEventHandler:
         self.notifier = notifier
         self.assistant_stream = StreamState()
         self.thinking_stream = StreamState()
+        self._sub_agent_thinking_headers: dict[str, SubAgentThinkingHeaderState] = {}
         self.spinner_status = SpinnerStatusState()
 
         self.stage_manager = StageManager(
@@ -302,6 +319,31 @@ class DisplayEventHandler:
             finish_thinking=self._finish_thinking_stream,
         )
 
+    def _new_thinking_mdstream(self) -> MarkdownStream:
+        return MarkdownStream(
+            mdargs={
+                "code_theme": self.renderer.themes.code_theme,
+                "style": ThemeKey.THINKING,
+            },
+            theme=self.renderer.themes.thinking_markdown_theme,
+            console=self.renderer.console,
+            live_sink=self.renderer.set_stream_renderable if MARKDOWN_STREAM_LIVE_REPAINT_ENABLED else None,
+            mark=THINKING_MESSAGE_MARK,
+            mark_style=ThemeKey.THINKING,
+            left_margin=MARKDOWN_LEFT_MARGIN,
+            markdown_class=ThinkingMarkdown,
+        )
+
+    def _new_assistant_mdstream(self) -> MarkdownStream:
+        return MarkdownStream(
+            mdargs={"code_theme": self.renderer.themes.code_theme},
+            theme=self.renderer.themes.markdown_theme,
+            console=self.renderer.console,
+            live_sink=self.renderer.set_stream_renderable if MARKDOWN_STREAM_LIVE_REPAINT_ENABLED else None,
+            mark=ASSISTANT_MESSAGE_MARK,
+            left_margin=MARKDOWN_LEFT_MARGIN,
+        )
+
     async def consume_event(self, event: events.Event) -> None:
         match event:
             case events.ReplayHistoryEvent() as e:
@@ -316,16 +358,16 @@ class DisplayEventHandler:
                 self._on_developer_message(e)
             case events.TurnStartEvent() as e:
                 self._on_turn_start(e)
-            case events.ThinkingEvent() as e:
-                await self._on_thinking(e)
             case events.ThinkingDeltaEvent() as e:
                 await self._on_thinking_delta(e)
-            case events.AssistantMessageDeltaEvent() as e:
+            case events.AssistantTextDeltaEvent() as e:
                 await self._on_assistant_delta(e)
+            case events.AssistantImageDeltaEvent() as e:
+                await self._on_assistant_image_delta(e)
             case events.AssistantMessageEvent() as e:
                 await self._on_assistant_message(e)
             case events.TurnToolCallStartEvent() as e:
-                self._on_tool_call_start(e)
+                await self._on_tool_call_start(e)
             case events.ToolCallEvent() as e:
                 await self._on_tool_call(e)
             case events.ToolResultEvent() as e:
@@ -350,8 +392,8 @@ class DisplayEventHandler:
                 await self._on_end(e)
 
     async def stop(self) -> None:
-        await self._flush_assistant_buffer(self.assistant_stream)
-        await self._flush_thinking_buffer(self.thinking_stream)
+        self._flush_assistant_buffer()
+        self._flush_thinking_buffer()
 
     # ─────────────────────────────────────────────────────────────────────────────
     # Private event handlers
@@ -370,6 +412,8 @@ class DisplayEventHandler:
     def _on_task_start(self, event: events.TaskStartEvent) -> None:
         if event.sub_agent_state is None:
             r_status.set_task_start()
+        else:
+            self._sub_agent_thinking_headers[event.session_id] = SubAgentThinkingHeaderState()
         self.renderer.spinner_start()
         self.renderer.display_task_start(event)
         emit_osc94(OSC94States.INDETERMINATE)
@@ -385,41 +429,20 @@ class DisplayEventHandler:
         self.spinner_status.set_reasoning_status(None)
         self._update_spinner()
 
-    async def _on_thinking(self, event: events.ThinkingEvent) -> None:
-        if self.renderer.is_sub_agent_session(event.session_id):
-            return
-        # If streaming was active, finalize it
-        if self.thinking_stream.is_active:
-            await self._finish_thinking_stream()
-        else:
-            # Non-streaming path (history replay or models without delta support)
-            reasoning_status = extract_last_bold_header(normalize_thinking_content(event.content))
-            if reasoning_status:
-                self.spinner_status.set_reasoning_status(reasoning_status)
-                self._update_spinner()
-            await self.stage_manager.enter_thinking_stage()
-            self.renderer.display_thinking(event.content)
-
     async def _on_thinking_delta(self, event: events.ThinkingDeltaEvent) -> None:
         if self.renderer.is_sub_agent_session(event.session_id):
+            if not self.renderer.should_display_sub_agent_thinking_header(event.session_id):
+                return
+            state = self._sub_agent_thinking_headers.setdefault(event.session_id, SubAgentThinkingHeaderState())
+            header = state.append_and_extract_new_header(event.content)
+            if header:
+                with self.renderer.session_print_context(event.session_id):
+                    self.renderer.display_thinking_header(header)
             return
 
         first_delta = not self.thinking_stream.is_active
         if first_delta:
-            mdstream = MarkdownStream(
-                mdargs={
-                    "code_theme": self.renderer.themes.code_theme,
-                    "style": ThemeKey.THINKING,
-                },
-                theme=self.renderer.themes.thinking_markdown_theme,
-                console=self.renderer.console,
-                live_sink=self.renderer.set_stream_renderable if const.MARKDOWN_STREAM_LIVE_REPAINT_ENABLED else None,
-                mark=THINKING_MESSAGE_MARK,
-                mark_style=ThemeKey.THINKING,
-                left_margin=const.MARKDOWN_LEFT_MARGIN,
-                markdown_class=ThinkingMarkdown,
-            )
-            self.thinking_stream.start(mdstream)
+            self.thinking_stream.start(self._new_thinking_mdstream())
 
         self.thinking_stream.append(event.content)
 
@@ -428,62 +451,52 @@ class DisplayEventHandler:
             self.spinner_status.set_reasoning_status(reasoning_status)
             self._update_spinner()
 
-        if first_delta and self.thinking_stream.mdstream is not None:
-            self.thinking_stream.mdstream.update(normalize_thinking_content(self.thinking_stream.buffer))
+        if first_delta:
+            self.thinking_stream.render(transform=normalize_thinking_content)
 
         await self.stage_manager.enter_thinking_stage()
-        await self._flush_thinking_buffer(self.thinking_stream)
+        self._flush_thinking_buffer()
 
-    async def _on_assistant_delta(self, event: events.AssistantMessageDeltaEvent) -> None:
+    async def _on_assistant_delta(self, event: events.AssistantTextDeltaEvent) -> None:
         if self.renderer.is_sub_agent_session(event.session_id):
             self.spinner_status.set_composing(True)
             self._update_spinner()
             return
+
         if len(event.content.strip()) == 0 and self.stage_manager.current_stage != Stage.ASSISTANT:
+            await self.stage_manager.transition_to(Stage.WAITING)
             return
+
+        await self.stage_manager.transition_to(Stage.ASSISTANT)
         first_delta = not self.assistant_stream.is_active
         if first_delta:
             self.spinner_status.set_composing(True)
             self.spinner_status.clear_tool_calls()
             self._update_spinner()
-            mdstream = MarkdownStream(
-                mdargs={"code_theme": self.renderer.themes.code_theme},
-                theme=self.renderer.themes.markdown_theme,
-                console=self.renderer.console,
-                live_sink=self.renderer.set_stream_renderable if const.MARKDOWN_STREAM_LIVE_REPAINT_ENABLED else None,
-                mark=ASSISTANT_MESSAGE_MARK,
-                left_margin=const.MARKDOWN_LEFT_MARGIN,
-            )
-            self.assistant_stream.start(mdstream)
+            self.assistant_stream.start(self._new_assistant_mdstream())
         self.assistant_stream.append(event.content)
         self.spinner_status.set_buffer_length(len(self.assistant_stream.buffer))
         if not first_delta:
             self._update_spinner()
-        if first_delta and self.assistant_stream.mdstream is not None:
-            self.assistant_stream.mdstream.update(self.assistant_stream.buffer)
-        await self.stage_manager.transition_to(Stage.ASSISTANT)
-        await self._flush_assistant_buffer(self.assistant_stream)
+        if first_delta:
+            self.assistant_stream.render()
+        self._flush_assistant_buffer()
 
     async def _on_assistant_message(self, event: events.AssistantMessageEvent) -> None:
         if self.renderer.is_sub_agent_session(event.session_id):
             return
-        await self.stage_manager.transition_to(Stage.ASSISTANT)
-        if self.assistant_stream.is_active:
-            mdstream = self.assistant_stream.mdstream
-            assert mdstream is not None
-            mdstream.update(event.content.strip(), final=True)
-        else:
-            self.renderer.display_assistant_message(event.content)
-        self.assistant_stream.finish()
+
+        await self.stage_manager.transition_to(Stage.WAITING)
         self.spinner_status.set_composing(False)
         self._update_spinner()
-        await self.stage_manager.transition_to(Stage.WAITING)
-        self.renderer.print()
         self.renderer.spinner_start()
 
-    def _on_tool_call_start(self, event: events.TurnToolCallStartEvent) -> None:
-        from klaude_code.ui.renderers.tools import get_tool_active_form
+    async def _on_assistant_image_delta(self, event: events.AssistantImageDeltaEvent) -> None:
+        await self.stage_manager.transition_to(Stage.ASSISTANT)
+        self.renderer.display_image(event.file_path)
 
+    async def _on_tool_call_start(self, event: events.TurnToolCallStartEvent) -> None:
+        self._flush_assistant_buffer()
         self.spinner_status.set_composing(False)
         self.spinner_status.add_tool_call(get_tool_active_form(event.tool_name))
         self._update_spinner()
@@ -494,18 +507,18 @@ class DisplayEventHandler:
             self.renderer.display_tool_call(event)
 
     async def _on_tool_result(self, event: events.ToolResultEvent) -> None:
-        if self.renderer.is_sub_agent_session(event.session_id) and event.status == "success":
+        is_sub_agent = self.renderer.is_sub_agent_session(event.session_id)
+        if is_sub_agent and event.status == "success":
             return
         await self.stage_manager.transition_to(Stage.TOOL_RESULT)
         with self.renderer.session_print_context(event.session_id):
-            self.renderer.display_tool_call_result(event)
+            self.renderer.display_tool_call_result(event, is_sub_agent=is_sub_agent)
 
     def _on_task_metadata(self, event: events.TaskMetadataEvent) -> None:
         self.renderer.display_task_metadata(event)
 
     def _on_todo_change(self, event: events.TodoChangeEvent) -> None:
-        active_form_status_text = self._extract_active_form_text(event)
-        self.spinner_status.set_todo_status(active_form_status_text if active_form_status_text else None)
+        self.spinner_status.set_todo_status(self._extract_active_form_text(event))
         # Clear tool calls when todo changes, as the tool execution has advanced
         self.spinner_status.clear_for_new_turn()
         self._update_spinner()
@@ -525,6 +538,8 @@ class DisplayEventHandler:
             self.renderer.spinner_stop()
             self.renderer.console.print(Rule(characters="─", style=ThemeKey.LINES))
             emit_tmux_signal()  # Signal test harness if KLAUDE_TEST_SIGNAL is set
+        else:
+            self._sub_agent_thinking_headers.pop(event.session_id, None)
         await self.stage_manager.transition_to(Stage.WAITING)
         self._maybe_notify_task_finish(event)
 
@@ -555,13 +570,6 @@ class DisplayEventHandler:
     # Private helper methods
     # ─────────────────────────────────────────────────────────────────────────────
 
-    async def _finish_assistant_stream(self) -> None:
-        if self.assistant_stream.is_active:
-            mdstream = self.assistant_stream.mdstream
-            assert mdstream is not None
-            mdstream.update(self.assistant_stream.buffer, final=True)
-            self.assistant_stream.finish()
-
     def _update_spinner(self) -> None:
         """Update spinner text from current status state."""
         status_text = self.spinner_status.get_status()
@@ -571,27 +579,23 @@ class DisplayEventHandler:
             right_text,
         )
 
-    async def _flush_assistant_buffer(self, state: StreamState) -> None:
-        if state.is_active:
-            mdstream = state.mdstream
-            assert mdstream is not None
-            mdstream.update(state.buffer)
+    def _flush_thinking_buffer(self) -> None:
+        self.thinking_stream.render(transform=normalize_thinking_content)
 
-    async def _flush_thinking_buffer(self, state: StreamState) -> None:
-        if state.is_active:
-            mdstream = state.mdstream
-            assert mdstream is not None
-            mdstream.update(normalize_thinking_content(state.buffer))
+    def _flush_assistant_buffer(self) -> None:
+        self.assistant_stream.render()
 
     async def _finish_thinking_stream(self) -> None:
-        if self.thinking_stream.is_active:
-            mdstream = self.thinking_stream.mdstream
-            assert mdstream is not None
-            mdstream.update(normalize_thinking_content(self.thinking_stream.buffer), final=True)
-            self.thinking_stream.finish()
+        finalized = self.thinking_stream.finalize(transform=normalize_thinking_content)
+        if finalized:
             self.renderer.print()
             self.renderer.spinner_start()
 
+    async def _finish_assistant_stream(self) -> None:
+        finalized = self.assistant_stream.finalize()
+        if finalized:
+            self.renderer.print()
+
     def _maybe_notify_task_finish(self, event: events.TaskFinishEvent) -> None:
         if self.notifier is None:
             return
@@ -617,12 +621,17 @@ class DisplayEventHandler:
             return squashed[:197] + "…"
         return squashed
 
-    def _extract_active_form_text(self, todo_event: events.TodoChangeEvent) -> str:
-        status_text = ""
+    def _extract_active_form_text(self, todo_event: events.TodoChangeEvent) -> str | None:
+        status_text: str | None = None
         for todo in todo_event.todos:
             if todo.status == "in_progress":
                 if len(todo.active_form) > 0:
                     status_text = todo.active_form
                 if len(todo.content) > 0:
                     status_text = todo.content
-        return status_text.replace("\n", " ").strip()
+
+        if status_text is None:
+            return None
+
+        normalized = status_text.replace("\n", " ").strip()
+        return normalized if normalized else None

klaude_code/ui/modes/repl/input_prompt_toolkit.py

@@ -34,7 +34,7 @@ from klaude_code.config.thinking import (
 )
 from klaude_code.protocol import llm_param
 from klaude_code.protocol.commands import CommandInfo
-from klaude_code.protocol.model import UserInputPayload
+from klaude_code.protocol.message import UserInputPayload
 from klaude_code.ui.core.input import InputProviderABC
 from klaude_code.ui.modes.repl.clipboard import capture_clipboard_tag, copy_to_clipboard, extract_images_from_text
 from klaude_code.ui.modes.repl.completers import AT_TOKEN_PATTERN, create_repl_completer

klaude_code/ui/modes/repl/key_bindings.py

@@ -198,7 +198,7 @@ def create_key_bindings(
         """Ensure completions refresh on backspace when editing an @token.
 
         We delete the character before cursor (default behavior), then explicitly
-        trigger completion refresh if the caret is still within an @... token.
+        trigger completion refresh if the caret is still within an @… token.
         """
         buf = event.current_buffer  # type: ignore
         # Handle selection: cut selection if present, otherwise delete one character