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

klaude_code/ui/modes/repl/renderer.py

@@ -12,17 +12,22 @@ from rich.spinner import Spinner
 from rich.style import Style, StyleType
 from rich.text import Text
 
-from klaude_code import const
-from klaude_code.protocol import events, model
+from klaude_code.const import (
+    MARKDOWN_STREAM_LIVE_REPAINT_ENABLED,
+    STATUS_DEFAULT_TEXT,
+    STREAM_MAX_HEIGHT_SHRINK_RESET_LINES,
+)
+from klaude_code.protocol import events, model, tools
 from klaude_code.ui.renderers import assistant as r_assistant
 from klaude_code.ui.renderers import developer as r_developer
 from klaude_code.ui.renderers import errors as r_errors
+from klaude_code.ui.renderers import mermaid_viewer as r_mermaid_viewer
 from klaude_code.ui.renderers import metadata as r_metadata
 from klaude_code.ui.renderers import sub_agent as r_sub_agent
 from klaude_code.ui.renderers import thinking as r_thinking
 from klaude_code.ui.renderers import tools as r_tools
 from klaude_code.ui.renderers import user_input as r_user_input
-from klaude_code.ui.renderers.common import truncate_display
+from klaude_code.ui.renderers.common import truncate_head, truncate_middle
 from klaude_code.ui.rich import status as r_status
 from klaude_code.ui.rich.live import CropAboveLive, SingleLine
 from klaude_code.ui.rich.quote import Quote
@@ -51,7 +56,7 @@ class REPLRenderer:
         self._stream_last_width: int = 0
         self._spinner_visible: bool = False
 
-        self._status_text: ShimmerStatusText = ShimmerStatusText(const.STATUS_DEFAULT_TEXT)
+        self._status_text: ShimmerStatusText = ShimmerStatusText(STATUS_DEFAULT_TEXT)
         self._status_spinner: Spinner = BreathingSpinner(
             r_status.spinner_name(),
             text=SingleLine(self._status_text),
@@ -75,6 +80,13 @@ class REPLRenderer:
     def is_sub_agent_session(self, session_id: str) -> bool:
         return session_id in self.session_map and self.session_map[session_id].sub_agent_state is not None
 
+    def should_display_sub_agent_thinking_header(self, session_id: str) -> bool:
+        # Hardcoded: only show sub-agent thinking headers for ImageGen.
+        status = self.session_map.get(session_id)
+        if status is None or status.sub_agent_state is None:
+            return False
+        return status.sub_agent_state.sub_agent_type == "ImageGen"
+
     def _advance_sub_agent_color_index(self) -> None:
         palette_size = len(self.themes.sub_agent_colors)
         if palette_size == 0:
@@ -127,10 +139,26 @@ class REPLRenderer:
         if renderable is not None:
             self.print(renderable)
 
-    def display_tool_call_result(self, e: events.ToolResultEvent) -> None:
+    def display_tool_call_result(self, e: events.ToolResultEvent, *, is_sub_agent: bool = False) -> None:
         if r_tools.is_sub_agent_tool(e.tool_name):
             return
-        renderable = r_tools.render_tool_result(e, code_theme=self.themes.code_theme)
+        # Sub-agent errors: show only first 2 lines
+        if is_sub_agent and e.status == "error":
+            error_msg = truncate_head(e.result)
+            self.print(r_errors.render_tool_error(error_msg))
+            return
+        if not is_sub_agent and e.tool_name == tools.MERMAID and isinstance(e.ui_extra, model.MermaidLinkUIExtra):
+            image_path = r_mermaid_viewer.download_mermaid_png(
+                link=e.ui_extra.link,
+                tool_call_id=e.tool_call_id,
+                session_id=e.session_id,
+            )
+            if image_path is not None:
+                self.display_image(str(image_path), height=None)
+
+            renderable = r_tools.render_tool_result(e, code_theme=self.themes.code_theme, session_id=e.session_id)
+        else:
+            renderable = r_tools.render_tool_result(e, code_theme=self.themes.code_theme, session_id=e.session_id)
         if renderable is not None:
             self.print(renderable)
 
@@ -146,8 +174,26 @@ class REPLRenderer:
             self.console.pop_theme()
             self.print()
 
+    def display_thinking_header(self, header: str) -> None:
+        """Display a single thinking header line.
+
+        Used by sub-agent sessions to avoid verbose thinking streaming.
+        """
+
+        stripped = header.strip()
+        if not stripped:
+            return
+        self.print(
+            Text.assemble(
+                (r_thinking.THINKING_MESSAGE_MARK, ThemeKey.THINKING),
+                " ",
+                (stripped, ThemeKey.THINKING_BOLD),
+            )
+        )
+
     async def replay_history(self, history_events: events.ReplayHistoryEvent) -> None:
         tool_call_dict: dict[str, events.ToolCallEvent] = {}
+        self.print()
         for event in history_events.events:
             event_session_id = getattr(event, "session_id", history_events.session_id)
             is_sub_agent = self.is_sub_agent_session(event_session_id)
@@ -158,17 +204,23 @@ class REPLRenderer:
                         self.display_task_start(e)
                     case events.TurnStartEvent():
                         self.print()
+                    case events.AssistantImageDeltaEvent() as e:
+                        self.display_image(e.file_path)
                     case events.AssistantMessageEvent() as e:
                         if is_sub_agent:
+                            if self.should_display_sub_agent_thinking_header(event_session_id) and e.thinking_text:
+                                header = r_thinking.extract_last_bold_header(
+                                    r_thinking.normalize_thinking_content(e.thinking_text)
+                                )
+                                if header:
+                                    self.display_thinking_header(header)
                             continue
+                        if e.thinking_text:
+                            self.display_thinking(e.thinking_text)
                         renderable = r_assistant.render_assistant_message(e.content, code_theme=self.themes.code_theme)
                         if renderable is not None:
                             self.print(renderable)
                             self.print()
-                    case events.ThinkingEvent() as e:
-                        if is_sub_agent:
-                            continue
-                        self.display_thinking(e.content)
                     case events.DeveloperMessageEvent() as e:
                         self.display_developer_message(e)
                         self.display_command_output(e)
@@ -187,6 +239,7 @@ class REPLRenderer:
                             continue
                         self.display_tool_call_result(e)
                     case events.TaskMetadataEvent() as e:
+                        self.print()
                         self.print(r_metadata.render_task_metadata(e))
                         self.print()
                     case events.InterruptEvent():
@@ -237,6 +290,37 @@ class REPLRenderer:
             self.print(renderable)
             self.print()
 
+    def display_image(self, file_path: str, height: int | None = 40) -> None:
+        """Display an image in the terminal.
+
+        Args:
+            file_path: Path to the image file.
+            height: Height in terminal lines for displaying the image.
+        """
+        from klaude_code.ui.terminal.image import print_kitty_image
+
+        # Suspend the Live status bar while emitting raw terminal output to avoid
+        # interleaving refreshes with Kitty graphics escape sequences.
+        had_live = self._bottom_live is not None
+        was_spinner_visible = self._spinner_visible
+        has_stream = MARKDOWN_STREAM_LIVE_REPAINT_ENABLED and self._stream_renderable is not None
+        resume_live = had_live and (was_spinner_visible or has_stream)
+
+        if self._bottom_live is not None:
+            with contextlib.suppress(Exception):
+                self._bottom_live.stop()
+            self._bottom_live = None
+
+        try:
+            print_kitty_image(file_path, height=height, file=self.console.file)
+        finally:
+            if resume_live:
+                if was_spinner_visible:
+                    self.spinner_start()
+                else:
+                    self._ensure_bottom_live_started()
+                    self._refresh_bottom_live()
+
     def display_task_metadata(self, event: events.TaskMetadataEvent) -> None:
         with self.session_print_context(event.session_id):
             self.print(r_metadata.render_task_metadata(event))
@@ -268,9 +352,9 @@ class REPLRenderer:
     def display_error(self, event: events.ErrorEvent) -> None:
         if event.session_id:
             with self.session_print_context(event.session_id):
-                self.print(r_errors.render_error(truncate_display(event.error_message)))
+                self.print(r_errors.render_error(truncate_middle(event.error_message)))
         else:
-            self.print(r_errors.render_error(truncate_display(event.error_message)))
+            self.print(r_errors.render_error(truncate_middle(event.error_message)))
 
     # -------------------------------------------------------------------------
     # Spinner control methods
@@ -314,7 +398,11 @@ class REPLRenderer:
         height = len(self.console.render_lines(renderable, self.console.options, pad=False))
         self._stream_last_height = height
         self._stream_last_width = self.console.size.width
-        self._stream_max_height = max(self._stream_max_height, height)
+
+        if self._stream_max_height - height > STREAM_MAX_HEIGHT_SHRINK_RESET_LINES:
+            self._stream_max_height = height
+        else:
+            self._stream_max_height = max(self._stream_max_height, height)
         self._refresh_bottom_live()
 
     def _ensure_bottom_live_started(self) -> None:
@@ -334,7 +422,7 @@ class REPLRenderer:
         stream_part: RenderableType = Group()
         gap_part: RenderableType = Group()
 
-        if const.MARKDOWN_STREAM_LIVE_REPAINT_ENABLED:
+        if MARKDOWN_STREAM_LIVE_REPAINT_ENABLED:
             stream = self._stream_renderable
             if stream is not None:
                 current_width = self.console.size.width
@@ -342,7 +430,11 @@ class REPLRenderer:
                     height = len(self.console.render_lines(stream, self.console.options, pad=False))
                     self._stream_last_height = height
                     self._stream_last_width = current_width
-                    self._stream_max_height = max(self._stream_max_height, height)
+
+                    if self._stream_max_height - height > STREAM_MAX_HEIGHT_SHRINK_RESET_LINES:
+                        self._stream_max_height = height
+                    else:
+                        self._stream_max_height = max(self._stream_max_height, height)
                 else:
                     height = self._stream_last_height
 

klaude_code/ui/renderers/assistant.py

@@ -2,7 +2,7 @@ from rich.console import RenderableType
 from rich.padding import Padding
 from rich.text import Text
 
-from klaude_code import const
+from klaude_code.const import MARKDOWN_RIGHT_MARGIN
 from klaude_code.ui.renderers.common import create_grid
 from klaude_code.ui.rich.markdown import NoInsetMarkdown
 from klaude_code.ui.rich.theme import ThemeKey
@@ -23,6 +23,6 @@ def render_assistant_message(content: str, *, code_theme: str) -> RenderableType
     grid = create_grid()
     grid.add_row(
         Text(ASSISTANT_MESSAGE_MARK, style=ThemeKey.ASSISTANT_MESSAGE_MARK),
-        Padding(NoInsetMarkdown(stripped, code_theme=code_theme), (0, const.MARKDOWN_RIGHT_MARGIN, 0, 0)),
+        Padding(NoInsetMarkdown(stripped, code_theme=code_theme), (0, MARKDOWN_RIGHT_MARGIN, 0, 0)),
     )
     return grid

klaude_code/ui/renderers/bash_syntax.py

@@ -7,6 +7,8 @@ from pygments.lexers import BashLexer  # pyright: ignore[reportUnknownVariableTy
 from pygments.token import Token
 from rich.text import Text
 
+from klaude_code.const import BASH_MULTILINE_STRING_TRUNCATE_MAX_LINES
+from klaude_code.ui.renderers.common import truncate_head
 from klaude_code.ui.rich.theme import ThemeKey
 
 # Token types for bash syntax highlighting
@@ -110,13 +112,34 @@ def _append_heredoc(result: Text, token_value: str) -> None:
         # Extra content on first line (e.g., "> file.py")
         if extra:
             result.append(extra, style=ThemeKey.BASH_ARGUMENT)
-        # Body content
-        result.append(body, style=ThemeKey.BASH_STRING)
+
+        # Body content (truncate to keep tool call rendering compact)
+        body_inner = body.strip("\n")
+        result.append("\n")
+        if body_inner:
+            body_text = truncate_head(
+                body_inner,
+                max_lines=BASH_MULTILINE_STRING_TRUNCATE_MAX_LINES,
+                base_style=ThemeKey.BASH_STRING,
+                truncated_style=ThemeKey.TOOL_RESULT_TRUNCATED,
+            )
+            result.append_text(body_text)
+            result.append("\n")
+
         # End delimiter
         result.append(end_delimiter, style=ThemeKey.BASH_HEREDOC_DELIMITER)
     else:
         # Fallback: couldn't parse heredoc structure
-        result.append(token_value, style=ThemeKey.BASH_STRING)
+        if "\n" in token_value and len(token_value.splitlines()) > BASH_MULTILINE_STRING_TRUNCATE_MAX_LINES:
+            truncated = truncate_head(
+                token_value,
+                max_lines=BASH_MULTILINE_STRING_TRUNCATE_MAX_LINES,
+                base_style=ThemeKey.BASH_STRING,
+                truncated_style=ThemeKey.TOOL_RESULT_TRUNCATED,
+            )
+            result.append_text(truncated)
+        else:
+            result.append(token_value, style=ThemeKey.BASH_STRING)
 
 
 def highlight_bash_command(command: str) -> Text:
@@ -145,7 +168,16 @@ def highlight_bash_command(command: str) -> Text:
             if token_value.startswith("<<"):
                 _append_heredoc(result, token_value)
             else:
-                result.append(token_value, style=ThemeKey.BASH_STRING)
+                if "\n" in token_value and len(token_value.splitlines()) > BASH_MULTILINE_STRING_TRUNCATE_MAX_LINES:
+                    truncated = truncate_head(
+                        token_value,
+                        max_lines=BASH_MULTILINE_STRING_TRUNCATE_MAX_LINES,
+                        base_style=ThemeKey.BASH_STRING,
+                        truncated_style=ThemeKey.TOOL_RESULT_TRUNCATED,
+                    )
+                    result.append_text(truncated)
+                else:
+                    result.append(token_value, style=ThemeKey.BASH_STRING)
             expect_subcommand = False
         elif token_type in _OPERATOR_TOKENS:
             result.append(token_value, style=ThemeKey.BASH_OPERATOR)

klaude_code/ui/renderers/common.py

@@ -1,22 +1,30 @@
+from typing import Literal
+
 from rich.style import Style
 from rich.table import Table
 from rich.text import Text
 
-from klaude_code import const
+from klaude_code.const import (
+    MIN_HIDDEN_LINES_FOR_INDICATOR,
+    TAB_EXPAND_WIDTH,
+    TRUNCATE_DISPLAY_MAX_LINE_LENGTH,
+    TRUNCATE_DISPLAY_MAX_LINES,
+    TRUNCATE_HEAD_MAX_LINES,
+)
 from klaude_code.ui.rich.theme import ThemeKey
 
 
-def create_grid() -> Table:
+def create_grid(*, overflow: Literal["fold", "crop", "ellipsis", "ignore"] = "fold") -> Table:
     grid = Table.grid(padding=(0, 1))
     grid.add_column(no_wrap=True)
-    grid.add_column(overflow="fold")
+    grid.add_column(overflow=overflow)
     return grid
 
 
-def truncate_display(
+def truncate_middle(
     text: str,
-    max_lines: int = const.TRUNCATE_DISPLAY_MAX_LINES,
-    max_line_length: int = const.TRUNCATE_DISPLAY_MAX_LINE_LENGTH,
+    max_lines: int = TRUNCATE_DISPLAY_MAX_LINES,
+    max_line_length: int = TRUNCATE_DISPLAY_MAX_LINE_LENGTH,
     *,
     base_style: str | Style | None = None,
 ) -> Text:
@@ -25,12 +33,12 @@ def truncate_display(
     Applies `ThemeKey.TOOL_RESULT_TRUNCATED` style to truncation indicators.
     """
     # Expand tabs to spaces to ensure correct alignment when Rich applies padding.
-    text = text.expandtabs(8)
+    text = text.expandtabs(TAB_EXPAND_WIDTH)
 
     if max_lines <= 0:
         truncated_lines = text.split("\n")
         remaining = max(0, len(truncated_lines))
-        return Text(f"… (more {remaining} lines)", style=ThemeKey.TOOL_RESULT_TRUNCATED)
+        return Text(f"  … (more {remaining} lines)", style=ThemeKey.TOOL_RESULT_TRUNCATED)
 
     lines = text.split("\n")
     truncated_lines = 0
@@ -42,7 +50,7 @@ def truncate_display(
 
         # If the hidden section is too small, show everything instead of inserting
         # the "(more N lines)" indicator.
-        if truncated_lines < 5:
+        if truncated_lines < MIN_HIDDEN_LINES_FOR_INDICATOR:
             truncated_lines = 0
             head_lines = lines
         else:
@@ -76,7 +84,7 @@ def truncate_display(
             out.append("\n")
 
     if truncated_lines > 0:
-        out.append_text(Text(f"⋮ (more {truncated_lines} lines)\n", style=ThemeKey.TOOL_RESULT_TRUNCATED))
+        out.append_text(Text(f"  … (more {truncated_lines} lines)\n", style=ThemeKey.TOOL_RESULT_TRUNCATED))
 
     for idx, line in enumerate(tail_lines):
         append_line(out, line)
@@ -84,3 +92,55 @@ def truncate_display(
             out.append("\n")
 
     return out
+
+
+def truncate_head(
+    text: str,
+    max_lines: int = TRUNCATE_HEAD_MAX_LINES,
+    max_line_length: int = TRUNCATE_DISPLAY_MAX_LINE_LENGTH,
+    *,
+    base_style: str | Style | None = None,
+    truncated_style: str | Style | None = None,
+) -> Text:
+    """Truncate text to show only the first N lines."""
+    text = text.expandtabs(TAB_EXPAND_WIDTH)
+    lines = [line for line in text.split("\n") if line.strip()]
+
+    out = Text()
+    if base_style is not None:
+        out.style = base_style
+
+    if len(lines) <= max_lines:
+        for idx, line in enumerate(lines):
+            if len(line) > max_line_length:
+                out.append(line[:max_line_length])
+                out.append_text(
+                    Text(
+                        f" … (more {len(line) - max_line_length} characters)",
+                        style=truncated_style or ThemeKey.TOOL_RESULT_TRUNCATED,
+                    )
+                )
+            else:
+                out.append(line)
+            if idx < len(lines) - 1:
+                out.append("\n")
+        return out
+
+    for idx in range(max_lines):
+        line = lines[idx]
+        if len(line) > max_line_length:
+            out.append(line[:max_line_length])
+            out.append_text(
+                Text(
+                    f" … (more {len(line) - max_line_length} characters)",
+                    style=truncated_style or ThemeKey.TOOL_RESULT_TRUNCATED,
+                )
+            )
+        else:
+            out.append(line)
+        out.append("\n")
+
+    remaining = len(lines) - max_lines
+    out.append_text(Text(f"  … (more {remaining} lines)", style=truncated_style or ThemeKey.TOOL_RESULT_TRUNCATED))
+
+    return out

klaude_code/ui/renderers/developer.py

@@ -3,8 +3,8 @@ from rich.padding import Padding
 from rich.table import Table
 from rich.text import Text
 
-from klaude_code.protocol import commands, events, model
-from klaude_code.ui.renderers.common import create_grid, truncate_display
+from klaude_code.protocol import commands, events, message, model
+from klaude_code.ui.renderers.common import create_grid, truncate_middle
 from klaude_code.ui.renderers.tools import render_path
 from klaude_code.ui.rich.markdown import NoInsetMarkdown
 from klaude_code.ui.rich.theme import ThemeKey
@@ -124,19 +124,20 @@ def render_command_output(e: events.DeveloperMessageEvent) -> RenderableType:
     if not e.item.command_output:
         return Text("")
 
+    content = message.join_text_parts(e.item.parts)
     match e.item.command_output.command_name:
         case commands.CommandName.HELP:
-            return Padding.indent(Text.from_markup(e.item.content or ""), level=2)
+            return Padding.indent(Text.from_markup(content or ""), level=2)
         case commands.CommandName.STATUS:
             return _render_status_output(e.item.command_output)
         case commands.CommandName.RELEASE_NOTES:
-            return Padding.indent(NoInsetMarkdown(e.item.content or ""), level=2)
+            return Padding.indent(NoInsetMarkdown(content or ""), level=2)
         case commands.CommandName.FORK_SESSION:
             return _render_fork_session_output(e.item.command_output)
         case _:
-            content = e.item.content or "(no content)"
+            content = content or "(no content)"
             style = ThemeKey.TOOL_RESULT if not e.item.command_output.is_error else ThemeKey.ERROR
-            return Padding.indent(truncate_display(content, base_style=style), level=2)
+            return Padding.indent(truncate_middle(content, base_style=style), level=2)
 
 
 def _format_tokens(tokens: int) -> str:

klaude_code/ui/renderers/diffs.py

@@ -4,7 +4,7 @@ from rich.padding import Padding
 from rich.panel import Panel
 from rich.text import Text
 
-from klaude_code import const
+from klaude_code.const import DIFF_PREFIX_WIDTH, MAX_DIFF_LINES
 from klaude_code.protocol import model
 from klaude_code.ui.renderers.common import create_grid
 from klaude_code.ui.rich.theme import ThemeKey
@@ -65,7 +65,7 @@ def render_diff(diff_text: str, show_file_name: bool = False) -> RenderableType:
             elif line.strip():  # Non-empty line in untracked section
                 file_text = Text(line.strip(), style=ThemeKey.TOOL_PARAM_BOLD)
                 grid.add_row(
-                    Text(f"{'+':>{const.DIFF_PREFIX_WIDTH}}", style=ThemeKey.TOOL_PARAM_BOLD),
+                    Text(f"{'+':>{DIFF_PREFIX_WIDTH}}", style=ThemeKey.TOOL_PARAM_BOLD),
                     file_text,
                 )
                 continue
@@ -130,7 +130,7 @@ def render_diff(diff_text: str, show_file_name: bool = False) -> RenderableType:
                 file_mark = "±"
 
             grid.add_row(
-                Text(f"{file_mark:>{const.DIFF_PREFIX_WIDTH}}  ", style=ThemeKey.DIFF_FILE_NAME),
+                Text(f"{file_mark:>{DIFF_PREFIX_WIDTH}}  ", style=ThemeKey.DIFF_FILE_NAME),
                 file_line,
             )
             has_rendered_file_header = True
@@ -151,7 +151,7 @@ def render_diff(diff_text: str, show_file_name: bool = False) -> RenderableType:
             except (IndexError, ValueError):
                 new_ln = None
             if has_rendered_diff_content:
-                grid.add_row(Text(f"{'⋮':>{const.DIFF_PREFIX_WIDTH}}", style=ThemeKey.TOOL_RESULT), "")
+                grid.add_row(Text(f"{'⋮':>{DIFF_PREFIX_WIDTH}}", style=ThemeKey.TOOL_RESULT), "")
             continue
 
         # Skip +++ lines (already handled above)
@@ -159,12 +159,12 @@ def render_diff(diff_text: str, show_file_name: bool = False) -> RenderableType:
             continue
 
         # Only handle unified diff hunk lines; ignore other metadata like
-        # "diff --git" or "index ..." which would otherwise skew counters.
+        # "diff --git" or "index …" which would otherwise skew counters.
         if not line or line[:1] not in {" ", "+", "-"}:
             continue
 
         # Compute line number prefix and style diff content
-        prefix, new_ln = _make_diff_prefix(line, new_ln, const.DIFF_PREFIX_WIDTH)
+        prefix, new_ln = _make_diff_prefix(line, new_ln, DIFF_PREFIX_WIDTH)
 
         if line.startswith("-"):
             text = Text(line[1:])
@@ -197,7 +197,7 @@ def render_structured_diff(ui_extra: model.DiffUIExtra, show_file_name: bool = F
             grid.add_row(*_render_file_header(file_diff))
 
         for line in file_diff.lines:
-            prefix = _make_structured_prefix(line, const.DIFF_PREFIX_WIDTH)
+            prefix = _make_structured_prefix(line, DIFF_PREFIX_WIDTH)
             text = _render_structured_line(line)
             grid.add_row(Text(prefix, ThemeKey.TOOL_RESULT), text)
 
@@ -213,9 +213,9 @@ def render_diff_panel(
 ) -> RenderableType:
     lines = diff_text.splitlines()
     truncated_notice: Text | None = None
-    if len(lines) > const.MAX_DIFF_LINES:
-        truncated_lines = len(lines) - const.MAX_DIFF_LINES
-        diff_text = "\n".join(lines[: const.MAX_DIFF_LINES])
+    if len(lines) > MAX_DIFF_LINES:
+        truncated_lines = len(lines) - MAX_DIFF_LINES
+        diff_text = "\n".join(lines[:MAX_DIFF_LINES])
         truncated_notice = Text(f"… truncated {truncated_lines} lines", style=ThemeKey.TOOL_MARK)
 
     diff_body = render_diff(diff_text, show_file_name=show_file_name)
@@ -261,7 +261,7 @@ def _render_file_header(file_diff: model.DiffFileDiff) -> tuple[Text, Text]:
     else:
         file_mark = "±"
 
-    prefix = Text(f"{file_mark:>{const.DIFF_PREFIX_WIDTH}}  ", style=ThemeKey.DIFF_FILE_NAME)
+    prefix = Text(f"{file_mark:>{DIFF_PREFIX_WIDTH}}  ", style=ThemeKey.DIFF_FILE_NAME)
     return prefix, file_line
 
 

klaude_code/ui/renderers/mermaid_viewer.py

@@ -5,11 +5,58 @@ import importlib.resources
 from functools import lru_cache
 from pathlib import Path
 
-from klaude_code import const
+import httpx
+
+from klaude_code.const import TOOL_OUTPUT_TRUNCATION_DIR
+from klaude_code.llm.image import get_assistant_image_output_dir
+
+_MERMAID_INK_PREFIX = "https://mermaid.ink/img/pako:"
+_MERMAID_DEFAULT_PNG_WIDTH = 1600
+_MERMAID_DEFAULT_PNG_SCALE = 2
 
 
 def artifacts_dir() -> Path:
-    return Path(const.TOOL_OUTPUT_TRUNCATION_DIR) / "mermaid"
+    return Path(TOOL_OUTPUT_TRUNCATION_DIR) / "mermaid"
+
+
+def _extract_pako_from_link(link: str) -> str | None:
+    """Extract pako encoded string from mermaid.live link."""
+    # link format: https://mermaid.live/view#pako:xxxx
+    if "#pako:" not in link:
+        return None
+    return link.split("#pako:", 1)[1]
+
+
+def download_mermaid_png(
+    *,
+    link: str,
+    tool_call_id: str,
+    session_id: str | None = None,
+) -> Path | None:
+    """Download PNG image from mermaid.ink and save locally."""
+    pako = _extract_pako_from_link(link)
+    if not pako:
+        return None
+
+    safe_id = tool_call_id.replace("/", "_")
+    output_dir = get_assistant_image_output_dir(session_id)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    image_path = output_dir / f"mermaid-{safe_id}.png"
+
+    if image_path.exists():
+        return image_path
+
+    png_url = (
+        f"{_MERMAID_INK_PREFIX}{pako}?type=png&width={_MERMAID_DEFAULT_PNG_WIDTH}&scale={_MERMAID_DEFAULT_PNG_SCALE}"
+    )
+    try:
+        with httpx.Client(timeout=10.0) as client:
+            resp = client.get(png_url)
+            resp.raise_for_status()
+            image_path.write_bytes(resp.content)
+        return image_path
+    except Exception:
+        return None
 
 
 @lru_cache(maxsize=1)

klaude_code/ui/renderers/metadata.py

@@ -6,7 +6,7 @@ from rich.padding import Padding
 from rich.panel import Panel
 from rich.text import Text
 
-from klaude_code import const
+from klaude_code.const import DEFAULT_MAX_TOKENS
 from klaude_code.protocol import events, model
 from klaude_code.trace import is_debug_enabled
 from klaude_code.ui.renderers.common import create_grid
@@ -47,7 +47,7 @@ def _render_task_metadata_block(
     # First column: mark only
     mark = Text("└", style=ThemeKey.METADATA_DIM) if is_sub_agent else Text("⇅", style=ThemeKey.METADATA)
 
-    # Second column: model@provider / tokens / cost / ...
+    # Second column: model@provider / tokens / cost / …
     content = Text()
     content.append_text(Text(metadata.model_name, style=ThemeKey.METADATA_BOLD))
     if metadata.provider is not None:
@@ -82,6 +82,13 @@ def _render_task_metadata_block(
                     (format_number(metadata.usage.reasoning_tokens), ThemeKey.METADATA),
                 )
             )
+        if metadata.usage.image_tokens > 0:
+            token_parts.append(
+                Text.assemble(
+                    ("image ", ThemeKey.METADATA_DIM),
+                    (format_number(metadata.usage.image_tokens), ThemeKey.METADATA),
+                )
+            )
         parts.append(Text(" · ").join(token_parts))
 
     # Cost
@@ -97,9 +104,7 @@ def _render_task_metadata_block(
         if show_context_and_time and metadata.usage.context_usage_percent is not None:
             context_size = format_number(metadata.usage.context_size or 0)
             # Calculate effective limit (same as Usage.context_usage_percent)
-            effective_limit = (metadata.usage.context_limit or 0) - (
-                metadata.usage.max_tokens or const.DEFAULT_MAX_TOKENS
-            )
+            effective_limit = (metadata.usage.context_limit or 0) - (metadata.usage.max_tokens or DEFAULT_MAX_TOKENS)
             effective_limit_str = format_number(effective_limit) if effective_limit > 0 else "?"
             parts.append(
                 Text.assemble(
@@ -167,6 +172,29 @@ def render_task_metadata(e: events.TaskMetadataEvent) -> RenderableType:
     for meta in e.metadata.sub_agent_task_metadata:
         renderables.append(_render_task_metadata_block(meta, is_sub_agent=True, show_context_and_time=False))
 
+    # Add total cost line when there are sub-agents
+    if e.metadata.sub_agent_task_metadata:
+        total_cost = 0.0
+        currency = "USD"
+        # Sum up costs from main agent and all sub-agents
+        if e.metadata.main_agent.usage and e.metadata.main_agent.usage.total_cost:
+            total_cost += e.metadata.main_agent.usage.total_cost
+            currency = e.metadata.main_agent.usage.currency
+        for meta in e.metadata.sub_agent_task_metadata:
+            if meta.usage and meta.usage.total_cost:
+                total_cost += meta.usage.total_cost
+
+        currency_symbol = "¥" if currency == "CNY" else "$"
+        total_line = Text.assemble(
+            ("Σ ", ThemeKey.METADATA_DIM),
+            ("total ", ThemeKey.METADATA_DIM),
+            (currency_symbol, ThemeKey.METADATA_DIM),
+            (f"{total_cost:.4f}", ThemeKey.METADATA_BOLD),
+        )
+        grid = create_grid()
+        grid.add_row(Text(" ", style=ThemeKey.METADATA_DIM), total_line)
+        renderables.append(Padding(grid, (0, 0, 0, 2)))
+
     return Group(*renderables)