klaude-code 1.2.6__py3-none-any.whl → 1.8.0__py3-none-any.whl

klaude_code/ui/renderers/mermaid_viewer.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+import html
+import importlib.resources
+from functools import lru_cache
+from pathlib import Path
+
+from klaude_code import const
+
+
+def artifacts_dir() -> Path:
+    return Path(const.TOOL_OUTPUT_TRUNCATION_DIR) / "mermaid"
+
+
+@lru_cache(maxsize=1)
+def load_template() -> str:
+    template_file = importlib.resources.files("klaude_code.session.templates").joinpath("mermaid_viewer.html")
+    return template_file.read_text(encoding="utf-8")
+
+
+def ensure_viewer_file(*, code: str, link: str, tool_call_id: str) -> Path | None:
+    """Create a local HTML viewer with large preview + editor."""
+
+    if not tool_call_id:
+        return None
+
+    safe_id = tool_call_id.replace("/", "_")
+    path = artifacts_dir() / f"mermaid-viewer-{safe_id}.html"
+    if path.exists():
+        return path
+
+    try:
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+        escaped_code = html.escape(code)
+        escaped_view_link = html.escape(link, quote=True)
+        escaped_edit_link = html.escape(link.replace("/view#pako:", "/edit#pako:"), quote=True)
+
+        template = load_template()
+        content = (
+            template.replace("__KLAUDE_VIEW_LINK__", escaped_view_link)
+            .replace("__KLAUDE_EDIT_LINK__", escaped_edit_link)
+            .replace("__KLAUDE_CODE__", escaped_code)
+        )
+        path.write_text(content, encoding="utf-8")
+    except OSError:
+        return None
+
+    return path
+
+
+def build_viewer(*, code: str, link: str, tool_call_id: str) -> Path | None:
+    """Create a local Mermaid viewer HTML file."""
+
+    if not code:
+        return None
+    return ensure_viewer_file(code=code, link=link, tool_call_id=tool_call_id)

klaude_code/ui/renderers/metadata.py

@@ -1,14 +1,15 @@
-from importlib.metadata import version
+from importlib.metadata import PackageNotFoundError, version
 
 from rich import box
-from rich.box import Box
 from rich.console import Group, RenderableType
 from rich.padding import Padding
 from rich.panel import Panel
 from rich.text import Text
 
-from klaude_code.protocol import events
+from klaude_code import const
+from klaude_code.protocol import events, model
 from klaude_code.trace import is_debug_enabled
+from klaude_code.ui.renderers.common import create_grid
 from klaude_code.ui.rich.theme import ThemeKey
 from klaude_code.ui.utils.common import format_number
 
@@ -17,80 +18,96 @@ def _get_version() -> str:
     """Get the current version of klaude-code."""
     try:
         return version("klaude-code")
-    except Exception:
+    except PackageNotFoundError:
         return "unknown"
 
 
-def render_response_metadata(e: events.ResponseMetadataEvent) -> RenderableType:
-    metadata = e.metadata
+def _render_task_metadata_block(
+    metadata: model.TaskMetadata,
+    *,
+    is_sub_agent: bool = False,
+    show_context_and_time: bool = True,
+) -> RenderableType:
+    """Render a single TaskMetadata block.
 
-    # Line 1: Model and Provider
-    model_text = Text()
-    model_text.append_text(Text("- ", style=ThemeKey.METADATA_BOLD)).append_text(
-        Text(metadata.model_name, style=ThemeKey.METADATA_BOLD)
-    )
+    Args:
+        metadata: The TaskMetadata to render.
+        is_sub_agent: Whether this is a sub-agent block.
+        show_context_and_time: Whether to show context usage percent and time.
+
+    Returns:
+        A renderable for this metadata block.
+    """
+    grid = create_grid()
+
+    # Get currency symbol
+    currency = metadata.usage.currency if metadata.usage else "USD"
+    currency_symbol = "¥" if currency == "CNY" else "$"
+
+    # 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 / ...
+    content = Text()
+    content.append_text(Text(metadata.model_name, style=ThemeKey.METADATA_BOLD))
     if metadata.provider is not None:
-        model_text.append_text(Text("@", style=ThemeKey.METADATA)).append_text(
+        content.append_text(Text("@", style=ThemeKey.METADATA)).append_text(
             Text(metadata.provider.lower().replace(" ", "-"), style=ThemeKey.METADATA)
         )
 
-    renderables: list[RenderableType] = [model_text]
-
-    # Line 2: Token consumption, Context, TPS, Cost
+    # All info parts (tokens, cost, context, etc.)
     parts: list[Text] = []
 
     if metadata.usage is not None:
-        # Input
-        input_parts: list[tuple[str, str]] = [
-            ("input:", ThemeKey.METADATA_DIM),
-            (format_number(metadata.usage.input_tokens), ThemeKey.METADATA_DIM),
+        # Tokens: ↑ 37k cache 5k ↓ 907 think 45k
+        token_parts: list[Text] = [
+            Text.assemble(("↑", ThemeKey.METADATA_DIM), (format_number(metadata.usage.input_tokens), ThemeKey.METADATA))
         ]
-        if metadata.usage.input_cost is not None:
-            input_parts.append((f"(${metadata.usage.input_cost:.4f})", ThemeKey.METADATA_DIM))
-        parts.append(Text.assemble(*input_parts))
-
-        # Cached
         if metadata.usage.cached_tokens > 0:
-            cached_parts: list[tuple[str, str]] = [
-                ("cached:", ThemeKey.METADATA_DIM),
-                (format_number(metadata.usage.cached_tokens), ThemeKey.METADATA_DIM),
-            ]
-            if metadata.usage.cache_read_cost is not None:
-                cached_parts.append((f"(${metadata.usage.cache_read_cost:.4f})", ThemeKey.METADATA_DIM))
-            parts.append(Text.assemble(*cached_parts))
-
-        # Output
-        output_parts: list[tuple[str, str]] = [
-            ("output:", ThemeKey.METADATA_DIM),
-            (format_number(metadata.usage.output_tokens), ThemeKey.METADATA_DIM),
-        ]
-        if metadata.usage.output_cost is not None:
-            output_parts.append((f"(${metadata.usage.output_cost:.4f})", ThemeKey.METADATA_DIM))
-        parts.append(Text.assemble(*output_parts))
-
-        # Reasoning
+            token_parts.append(
+                Text.assemble(
+                    Text("cache ", style=ThemeKey.METADATA_DIM),
+                    Text(format_number(metadata.usage.cached_tokens), style=ThemeKey.METADATA),
+                )
+            )
+        token_parts.append(
+            Text.assemble(
+                ("↓", ThemeKey.METADATA_DIM), (format_number(metadata.usage.output_tokens), ThemeKey.METADATA)
+            )
+        )
         if metadata.usage.reasoning_tokens > 0:
-            parts.append(
+            token_parts.append(
                 Text.assemble(
-                    ("thinking", ThemeKey.METADATA_DIM),
-                    (":", ThemeKey.METADATA_DIM),
-                    (
-                        format_number(metadata.usage.reasoning_tokens),
-                        ThemeKey.METADATA_DIM,
-                    ),
+                    ("think ", ThemeKey.METADATA_DIM),
+                    (format_number(metadata.usage.reasoning_tokens), ThemeKey.METADATA),
                 )
             )
+        parts.append(Text(" · ").join(token_parts))
 
-        # Context
-        if metadata.usage.context_usage_percent is not None:
+    # Cost
+    if metadata.usage is not None and metadata.usage.total_cost is not None:
+        parts.append(
+            Text.assemble(
+                (currency_symbol, ThemeKey.METADATA_DIM),
+                (f"{metadata.usage.total_cost:.4f}", ThemeKey.METADATA),
+            )
+        )
+    if metadata.usage is not None:
+        # Context (only for main agent)
+        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_str = format_number(effective_limit) if effective_limit > 0 else "?"
             parts.append(
                 Text.assemble(
-                    ("context", ThemeKey.METADATA_DIM),
-                    (":", ThemeKey.METADATA_DIM),
-                    (
-                        f"{metadata.usage.context_usage_percent:.1f}%",
-                        ThemeKey.METADATA_DIM,
-                    ),
+                    ("context ", ThemeKey.METADATA_DIM),
+                    (context_size, ThemeKey.METADATA),
+                    ("/", ThemeKey.METADATA_DIM),
+                    (effective_limit_str, ThemeKey.METADATA),
+                    (f" ({metadata.usage.context_usage_percent:.1f}%)", ThemeKey.METADATA_DIM),
                 )
             )
 
@@ -98,44 +115,63 @@ def render_response_metadata(e: events.ResponseMetadataEvent) -> RenderableType:
         if metadata.usage.throughput_tps is not None:
             parts.append(
                 Text.assemble(
-                    ("tps", ThemeKey.METADATA_DIM),
-                    (":", ThemeKey.METADATA_DIM),
-                    (f"{metadata.usage.throughput_tps:.1f}", ThemeKey.METADATA_DIM),
+                    (f"{metadata.usage.throughput_tps:.1f} ", ThemeKey.METADATA),
+                    ("avg-tps", ThemeKey.METADATA_DIM),
+                )
+            )
+
+        # First token latency
+        if metadata.usage.first_token_latency_ms is not None:
+            parts.append(
+                Text.assemble(
+                    (f"{metadata.usage.first_token_latency_ms:.0f}", ThemeKey.METADATA),
+                    ("ms avg-ftl", ThemeKey.METADATA_DIM),
                 )
             )
 
     # Duration
-    if metadata.task_duration_s is not None:
+    if show_context_and_time and metadata.task_duration_s is not None:
         parts.append(
             Text.assemble(
-                ("time", ThemeKey.METADATA_DIM),
-                (":", ThemeKey.METADATA_DIM),
-                (f"{metadata.task_duration_s:.1f}s", ThemeKey.METADATA_DIM),
+                (f"{metadata.task_duration_s:.1f}", ThemeKey.METADATA),
+                ("s", ThemeKey.METADATA_DIM),
             )
         )
 
-    # Cost (USD)
-    if metadata.usage is not None and metadata.usage.total_cost is not None:
+    # Turn count
+    if show_context_and_time and metadata.turn_count > 0:
         parts.append(
             Text.assemble(
-                ("cost", ThemeKey.METADATA_DIM),
-                (":", ThemeKey.METADATA_DIM),
-                (f"${metadata.usage.total_cost:.4f}", ThemeKey.METADATA_DIM),
+                (str(metadata.turn_count), ThemeKey.METADATA),
+                (" turns", ThemeKey.METADATA_DIM),
             )
         )
 
     if parts:
-        line2 = Text("/", style=ThemeKey.METADATA_DIM).join(parts)
-        renderables.append(Padding(line2, (0, 0, 0, 2)))
+        content.append_text(Text(" · ", style=ThemeKey.METADATA_DIM))
+        content.append_text(Text(" · ", style=ThemeKey.METADATA_DIM).join(parts))
+
+    grid.add_row(mark, content)
+    return grid if not is_sub_agent else Padding(grid, (0, 0, 0, 2))
+
+
+def render_task_metadata(e: events.TaskMetadataEvent) -> RenderableType:
+    """Render task metadata including main agent and sub-agents."""
+    renderables: list[RenderableType] = []
+
+    renderables.append(
+        _render_task_metadata_block(e.metadata.main_agent, is_sub_agent=False, show_context_and_time=True)
+    )
+
+    # Render each sub-agent metadata block
+    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))
 
     return Group(*renderables)
 
 
-def render_welcome(e: events.WelcomeEvent, *, box_style: Box | None = None) -> RenderableType:
+def render_welcome(e: events.WelcomeEvent) -> RenderableType:
     """Render the welcome panel with model info and settings."""
-    if box_style is None:
-        box_style = box.ROUNDED
-
     debug_mode = is_debug_enabled()
 
     # First line: Klaude Code version
@@ -185,6 +221,6 @@ def render_welcome(e: events.WelcomeEvent, *, box_style: Box | None = None) -> R
 
     border_style = ThemeKey.WELCOME_DEBUG_BORDER if debug_mode else ThemeKey.LINES
     return Group(
-        Panel.fit(panel_content, border_style=border_style, box=box_style),
+        Panel.fit(panel_content, border_style=border_style, box=box.ROUNDED),
         "",  # empty line
     )

klaude_code/ui/renderers/sub_agent.py

@@ -1,6 +1,9 @@
 import json
+from typing import Any, cast
 
+from rich import box
 from rich.console import Group, RenderableType
+from rich.json import JSON
 from rich.panel import Panel
 from rich.style import Style
 from rich.text import Text
@@ -12,37 +15,112 @@ from klaude_code.ui.rich.markdown import NoInsetMarkdown
 from klaude_code.ui.rich.theme import ThemeKey
 
 
+def _compact_schema_value(value: dict[str, Any]) -> str | list[Any] | dict[str, Any]:
+    """Convert a JSON Schema value to compact representation."""
+    value_type = value.get("type", "any").lower()
+    desc = value.get("description", "")
+
+    if value_type == "object":
+        props = value.get("properties", {})
+        return {k: _compact_schema_value(v) for k, v in props.items()}
+    elif value_type == "array":
+        items = value.get("items", {})
+        # If items have no description, use the array's description
+        if desc and not items.get("description"):
+            items = {**items, "description": desc}
+        return [_compact_schema_value(items)]
+    else:
+        if desc:
+            return f"{value_type} // {desc}"
+        return value_type
+
+
+def _compact_schema(schema: dict[str, Any]) -> dict[str, Any] | list[Any] | str:
+    """Convert JSON Schema to compact representation for display."""
+    return _compact_schema_value(schema)
+
+
 def render_sub_agent_call(e: model.SubAgentState, style: Style | None = None) -> RenderableType:
     """Render sub-agent tool call header and prompt body."""
     desc = Text(
         f" {e.sub_agent_desc} ",
         style=Style(color=style.color if style else None, bold=True, reverse=True),
     )
-    return Group(
+    elements: list[RenderableType] = [
         Text.assemble((e.sub_agent_type, ThemeKey.TOOL_NAME), " ", desc),
         Text(e.sub_agent_prompt, style=style or ""),
-    )
+    ]
+    if e.output_schema:
+        elements.append(Text("\nExpected Output Format JSON:", style=style or ""))
+        compact = _compact_schema(e.output_schema)
+        schema_text = json.dumps(compact, ensure_ascii=False, indent=2)
+        elements.append(JSON(schema_text))
+    return Group(*elements)
 
 
-def render_sub_agent_result(result: str, *, code_theme: str, style: Style | None = None) -> RenderableType:
+def render_sub_agent_result(
+    result: str,
+    *,
+    code_theme: str,
+    style: Style | None = None,
+    has_structured_output: bool = False,
+    description: str | None = None,
+    panel_style: Style | None = None,
+) -> RenderableType:
     stripped_result = result.strip()
+    result_panel_style = panel_style or ThemeKey.SUB_AGENT_RESULT_PANEL
+
+    # Use rich JSON for structured output
+    if has_structured_output:
+        try:
+            group_elements: list[RenderableType] = [
+                Text(
+                    "use /export to view full output",
+                    style=ThemeKey.TOOL_RESULT,
+                ),
+                JSON(stripped_result),
+            ]
+            if description:
+                group_elements.insert(0, NoInsetMarkdown(f"# {description}", code_theme=code_theme, style=style or ""))
+            return Panel.fit(
+                Group(*group_elements),
+                box=box.SIMPLE,
+                border_style=ThemeKey.LINES,
+                style=result_panel_style,
+            )
+        except json.JSONDecodeError:
+            # Fall back to markdown if not valid JSON
+            pass
+
+    # Add markdown heading if description is provided for non-structured output
+    if description:
+        stripped_result = f"# {description}\n\n{stripped_result}"
+
     lines = stripped_result.splitlines()
     if len(lines) > const.SUB_AGENT_RESULT_MAX_LINES:
         hidden_count = len(lines) - const.SUB_AGENT_RESULT_MAX_LINES
-        truncated_text = "\n".join(lines[-const.SUB_AGENT_RESULT_MAX_LINES :])
+        head_count = const.SUB_AGENT_RESULT_MAX_LINES // 2
+        tail_count = const.SUB_AGENT_RESULT_MAX_LINES - head_count
+        head_text = "\n".join(lines[:head_count])
+        tail_text = "\n".join(lines[-tail_count:])
         return Panel.fit(
             Group(
+                NoInsetMarkdown(head_text, code_theme=code_theme, style=style or ""),
                 Text(
-                    f"… more {hidden_count} lines — use /export to view full output",
-                    style=ThemeKey.TOOL_RESULT,
+                    f"\n( … more {hidden_count} lines — use /export to view full output )\n",
+                    style=ThemeKey.TOOL_RESULT_TRUNCATED,
                 ),
-                NoInsetMarkdown(truncated_text, code_theme=code_theme, style=style or ""),
+                NoInsetMarkdown(tail_text, code_theme=code_theme, style=style or ""),
             ),
+            box=box.SIMPLE,
             border_style=ThemeKey.LINES,
+            style=result_panel_style,
         )
     return Panel.fit(
         NoInsetMarkdown(stripped_result, code_theme=code_theme),
+        box=box.SIMPLE,
         border_style=ThemeKey.LINES,
+        style=result_panel_style,
     )
 
 
@@ -53,6 +131,7 @@ def build_sub_agent_state_from_tool_call(e: events.ToolCallEvent) -> model.SubAg
         return None
     description = profile.name
     prompt = ""
+    output_schema: dict[str, Any] | None = None
     if e.arguments:
         try:
             payload: dict[str, object] = json.loads(e.arguments)
@@ -64,8 +143,14 @@ def build_sub_agent_state_from_tool_call(e: events.ToolCallEvent) -> model.SubAg
         prompt_value = payload.get("prompt") or payload.get("task")
         if isinstance(prompt_value, str):
             prompt = prompt_value.strip()
+        # Extract output_schema if profile supports it
+        if profile.output_schema_arg:
+            schema_value = payload.get(profile.output_schema_arg)
+            if isinstance(schema_value, dict):
+                output_schema = cast(dict[str, Any], schema_value)
     return model.SubAgentState(
         sub_agent_type=profile.name,
         sub_agent_desc=description,
         sub_agent_prompt=prompt,
+        output_schema=output_schema,
     )

klaude_code/ui/renderers/thinking.py

@@ -1,27 +1,44 @@
+import re
+
 from rich.console import RenderableType
 from rich.padding import Padding
 from rich.text import Text
 
-from klaude_code.ui.rich.markdown import NoInsetMarkdown
+from klaude_code import const
+from klaude_code.ui.renderers.common import create_grid
+from klaude_code.ui.rich.markdown import ThinkingMarkdown
 from klaude_code.ui.rich.theme import ThemeKey
 
-
-def thinking_prefix() -> Text:
-    return Text.from_markup("[not italic]⸫[/not italic] Thinking …", style=ThemeKey.THINKING)
+# UI markers
+THINKING_MESSAGE_MARK = "∴"
 
 
-def _normalize_thinking_content(content: str) -> str:
+def normalize_thinking_content(content: str) -> str:
     """Normalize thinking content for display."""
-    return (
-        content.rstrip()
-        .replace("**\n\n", "**  \n")
-        .replace("\\n\\n\n\n", "")  # Weird case of Gemini 3
-        .replace("****", "**\n\n**")  # remove extra newlines after bold titles
-    )
+    text = content.rstrip()
+
+    # Weird case of Gemini 3
+    text = text.replace("\\n\\n\n\n", "")
+
+    # Fix OpenRouter OpenAI reasoning formatting where segments like
+    # "text**Title**\n\n" lose the blank line between segments.
+    # We want: "text\n**Title**\n" so that each bold title starts on
+    # its own line and uses a single trailing newline.
+    text = re.sub(r"([^\n])(\*\*[^*]+?\*\*)\n\n", r"\1  \n\n\2  \n", text)
+
+    # Remove extra newlines between back-to-back bold titles, eg
+    # "**Title1****Title2**" -> "**Title1**\n\n**Title2**".
+    text = text.replace("****", "**\n\n**")
+
+    # Compact double-newline after bold so the body text follows
+    # directly after the title line, using a markdown line break.
+    text = text.replace("**\n\n", "**  \n")
+
+    return text
 
 
 def render_thinking(content: str, *, code_theme: str, style: str) -> RenderableType | None:
-    """Render thinking content as indented markdown.
+    """Render thinking content as markdown with left mark.
 
     Returns None if content is empty.
     Note: Caller should push thinking_markdown_theme before printing.
@@ -29,11 +46,16 @@ def render_thinking(content: str, *, code_theme: str, style: str) -> RenderableT
     if len(content.strip()) == 0:
         return None
 
-    return Padding.indent(
-        NoInsetMarkdown(
-            _normalize_thinking_content(content),
-            code_theme=code_theme,
-            style=style,
+    grid = create_grid()
+    grid.add_row(
+        Text(THINKING_MESSAGE_MARK, style=ThemeKey.THINKING),
+        Padding(
+            ThinkingMarkdown(
+                normalize_thinking_content(content),
+                code_theme=code_theme,
+                style=style,
+            ),
+            (0, const.MARKDOWN_RIGHT_MARGIN, 0, 0),
         ),
-        level=2,
     )
+    return grid