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

klaude_code/session/export.py

@@ -63,7 +63,7 @@ def get_default_export_path(session: Session) -> Path:
     """Get default export path for a session."""
     from klaude_code.session.session import Session as SessionClass
 
-    exports_dir = SessionClass._exports_dir()  # pyright: ignore[reportPrivateUsage]
+    exports_dir = SessionClass.exports_dir()
     timestamp = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
     first_msg = get_first_user_message(session.conversation_history)
     sanitized_msg = _sanitize_filename(first_msg)
@@ -154,67 +154,94 @@ def _format_token_count(count: int) -> str:
     return f"{m}M" if rem == 0 else f"{m}M{rem}k"
 
 
-def _format_cost(cost: float) -> str:
-    return f"${cost:.4f}"
+def _format_cost(cost: float, currency: str = "USD") -> str:
+    symbol = "¥" if currency == "CNY" else "$"
+    return f"{symbol}{cost:.4f}"
 
 
-def _render_metadata_item(item: model.ResponseMetadataItem) -> str:
-    # Model Name [@ Provider]
+def _render_single_metadata(
+    metadata: model.TaskMetadata,
+    *,
+    indent: int = 0,
+    show_context: bool = True,
+) -> str:
+    """Render a single TaskMetadata block as HTML.
+
+    Args:
+        metadata: The TaskMetadata to render.
+        indent: Number of spaces to indent (0 for main, 2 for sub-agents).
+        show_context: Whether to show context usage percent.
+
+    Returns:
+        HTML string for this metadata block.
+    """
     parts: list[str] = []
 
-    model_parts = [f'<span class="metadata-model">{_escape_html(item.model_name)}</span>']
-    if item.provider:
-        provider = _escape_html(item.provider.lower().replace(" ", "-"))
+    # Model Name [@ Provider]
+    model_parts = [f'<span class="metadata-model">{_escape_html(metadata.model_name)}</span>']
+    if metadata.provider:
+        provider = _escape_html(metadata.provider.lower().replace(" ", "-"))
         model_parts.append(f'<span class="metadata-provider">@{provider}</span>')
 
     parts.append("".join(model_parts))
 
     # Stats
-    if item.usage:
-        u = item.usage
+    if metadata.usage:
+        u = metadata.usage
         # Input with cost
         input_stat = f"input: {_format_token_count(u.input_tokens)}"
         if u.input_cost is not None:
-            input_stat += f"({_format_cost(u.input_cost)})"
+            input_stat += f"({_format_cost(u.input_cost, u.currency)})"
         parts.append(f'<span class="metadata-stat">{input_stat}</span>')
 
         # Cached with cost
         if u.cached_tokens > 0:
             cached_stat = f"cached: {_format_token_count(u.cached_tokens)}"
             if u.cache_read_cost is not None:
-                cached_stat += f"({_format_cost(u.cache_read_cost)})"
+                cached_stat += f"({_format_cost(u.cache_read_cost, u.currency)})"
             parts.append(f'<span class="metadata-stat">{cached_stat}</span>')
 
         # Output with cost
         output_stat = f"output: {_format_token_count(u.output_tokens)}"
         if u.output_cost is not None:
-            output_stat += f"({_format_cost(u.output_cost)})"
+            output_stat += f"({_format_cost(u.output_cost, u.currency)})"
         parts.append(f'<span class="metadata-stat">{output_stat}</span>')
 
         if u.reasoning_tokens > 0:
-            parts.append(
-                f'<span class="metadata-stat">thinking: {_format_token_count(u.reasoning_tokens)}</span>'
-            )
-        if u.context_usage_percent is not None:
+            parts.append(f'<span class="metadata-stat">thinking: {_format_token_count(u.reasoning_tokens)}</span>')
+        if show_context and u.context_usage_percent is not None:
             parts.append(f'<span class="metadata-stat">context: {u.context_usage_percent:.1f}%</span>')
         if u.throughput_tps is not None:
             parts.append(f'<span class="metadata-stat">tps: {u.throughput_tps:.1f}</span>')
 
-    if item.task_duration_s is not None:
-        parts.append(f'<span class="metadata-stat">time: {item.task_duration_s:.1f}s</span>')
+    if metadata.task_duration_s is not None:
+        parts.append(f'<span class="metadata-stat">time: {metadata.task_duration_s:.1f}s</span>')
 
     # Total cost
-    if item.usage is not None and item.usage.total_cost is not None:
-        parts.append(f'<span class="metadata-stat">cost: {_format_cost(item.usage.total_cost)}</span>')
+    if metadata.usage is not None and metadata.usage.total_cost is not None:
+        parts.append(
+            f'<span class="metadata-stat">cost: {_format_cost(metadata.usage.total_cost, metadata.usage.currency)}</span>'
+        )
 
     divider = '<span class="metadata-divider">/</span>'
     joined_html = divider.join(parts)
 
-    return (
-        f'<div class="response-metadata">'
-        f'<div class="metadata-line">{joined_html}</div>'
-        f"</div>"
-    )
+    indent_style = f' style="padding-left: {indent}em;"' if indent > 0 else ""
+    return f'<div class="metadata-line"{indent_style}>{joined_html}</div>'
+
+
+def _render_metadata_item(item: model.TaskMetadataItem) -> str:
+    """Render TaskMetadataItem including main agent and sub-agents."""
+    lines: list[str] = []
+
+    # Main agent metadata
+    lines.append(_render_single_metadata(item.main_agent, indent=0, show_context=True))
+
+    # Sub-agent metadata with indent
+    for sub in item.sub_agent_task_metadata:
+        lines.append(_render_single_metadata(sub, indent=1, show_context=False))
+
+    return f'<div class="response-metadata">{"".join(lines)}</div>'
 
 
 def _render_assistant_message(index: int, content: str, timestamp: datetime) -> str:
@@ -240,20 +267,30 @@ def _render_assistant_message(index: int, content: str, timestamp: datetime) ->
     )
 
 
-def _try_render_todo_args(arguments: str) -> str | None:
+def _try_render_todo_args(arguments: str, tool_name: str) -> str | None:
     try:
         parsed = json.loads(arguments)
-        if not isinstance(parsed, dict) or "todos" not in parsed or not isinstance(parsed["todos"], list):
+        if not isinstance(parsed, dict):
             return None
 
-        todos = cast(list[dict[str, str]], parsed["todos"])
-        if not todos:
+        # Support both TodoWrite (todos/content) and update_plan (plan/step)
+        parsed_dict = cast(dict[str, Any], parsed)
+        if tool_name == "TodoWrite":
+            items = parsed_dict.get("todos")
+            content_key = "content"
+        elif tool_name == "update_plan":
+            items = parsed_dict.get("plan")
+            content_key = "step"
+        else:
+            return None
+
+        if not isinstance(items, list) or not items:
             return None
 
         items_html: list[str] = []
-        for todo in todos:
-            content = _escape_html(todo.get("content", ""))
-            status = todo.get("status", "pending")
+        for item in cast(list[dict[str, str]], items):
+            content = _escape_html(item.get(content_key, ""))
+            status = item.get("status", "pending")
             status_class = f"status-{status}"
 
             items_html.append(
@@ -267,23 +304,33 @@ def _try_render_todo_args(arguments: str) -> str | None:
             return None
 
         return f'<div class="todo-list">{"".join(items_html)}</div>'
-    except Exception:
+    except (json.JSONDecodeError, KeyError, TypeError):
         return None
 
 
-def _render_sub_agent_result(content: str) -> str:
-    encoded = _escape_html(content)
+def _render_sub_agent_result(content: str, description: str | None = None) -> str:
+    # Try to format as JSON for better readability
+    try:
+        parsed = json.loads(content)
+        formatted = "```json\n" + json.dumps(parsed, ensure_ascii=False, indent=2) + "\n```"
+    except (json.JSONDecodeError, TypeError):
+        formatted = content
+
+    if description:
+        formatted = f"# {description}\n\n{formatted}"
+
+    encoded = _escape_html(formatted)
     return (
-        f'<div class="subagent-result-container">'
-        f'<div class="subagent-toolbar">'
+        f'<div class="sub-agent-result-container">'
+        f'<div class="sub-agent-toolbar">'
         f'<button type="button" class="raw-toggle" aria-pressed="false" title="Toggle raw text view">Raw</button>'
         f'<button type="button" class="copy-raw-btn" title="Copy raw content">Copy</button>'
         f"</div>"
-        f'<div class="subagent-content">'
-        f'<div class="subagent-rendered markdown-content markdown-body" data-raw="{encoded}">'
+        f'<div class="sub-agent-content">'
+        f'<div class="sub-agent-rendered markdown-content markdown-body" data-raw="{encoded}">'
         f'<noscript><pre style="white-space: pre-wrap;">{encoded}</pre></noscript>'
         f"</div>"
-        f'<pre class="subagent-raw">{encoded}</pre>'
+        f'<pre class="sub-agent-raw">{encoded}</pre>'
         f"</div>"
         f"</div>"
     )
@@ -319,55 +366,141 @@ def _should_collapse(text: str) -> bool:
     return text.count("\n") + 1 > _COLLAPSIBLE_LINE_THRESHOLD or len(text) > _COLLAPSIBLE_CHAR_THRESHOLD
 
 
-def _render_diff_block(diff: str) -> str:
-    lines = diff.splitlines()
+def _render_diff_block(diff: model.DiffUIExtra) -> str:
     rendered: list[str] = []
-    for line in lines:
-        escaped = _escape_html(line)
-        if line.startswith("+"):
-            rendered.append(f'<span class="diff-line diff-plus">{escaped}</span>')
-        elif line.startswith("-"):
-            rendered.append(f'<span class="diff-line diff-minus">{escaped}</span>')
-        else:
-            rendered.append(f'<span class="diff-line diff-ctx">{escaped}</span>')
+    line_count = 0
+
+    for file_diff in diff.files:
+        header = _render_diff_file_header(file_diff)
+        if header:
+            rendered.append(header)
+        for line in file_diff.lines:
+            rendered.append(_render_diff_line(line))
+            line_count += 1
+
+    if line_count == 0:
+        rendered.append('<span class="diff-line diff-ctx">&nbsp;</span>')
+
     diff_content = f'<div class="diff-view">{"".join(rendered)}</div>'
-    open_attr = "" if _should_collapse(diff) else " open"
+    open_attr = "" if _should_collapse("\n" * max(1, line_count)) else " open"
     return (
         f'<details class="diff-collapsible"{open_attr}>'
-        f"<summary>Diff ({len(lines)} lines)</summary>"
+        f"<summary>Diff ({line_count} lines)</summary>"
         f"{diff_content}"
         "</details>"
     )
 
 
-def _get_diff_text(ui_extra: model.ToolResultUIExtra | None) -> str | None:
+def _render_diff_file_header(file_diff: model.DiffFileDiff) -> str:
+    stats_parts: list[str] = []
+    if file_diff.stats_add > 0:
+        stats_parts.append(f'<span class="diff-stats-add">+{file_diff.stats_add}</span>')
+    if file_diff.stats_remove > 0:
+        stats_parts.append(f'<span class="diff-stats-remove">-{file_diff.stats_remove}</span>')
+    stats_html = f' <span class="diff-stats">{" ".join(stats_parts)}</span>' if stats_parts else ""
+    file_name = _escape_html(file_diff.file_path)
+    return f'<div class="diff-file">{file_name}{stats_html}</div>'
+
+
+def _render_diff_line(line: model.DiffLine) -> str:
+    if line.kind == "gap":
+        line_class = "diff-ctx"
+        prefix = "⋮"
+    else:
+        line_class = "diff-plus" if line.kind == "add" else "diff-minus" if line.kind == "remove" else "diff-ctx"
+        prefix = "+" if line.kind == "add" else "-" if line.kind == "remove" else " "
+    spans = [_render_diff_span(span, line.kind) for span in line.spans]
+    content = "".join(spans)
+    if not content:
+        content = "&nbsp;"
+    return f'<span class="diff-line {line_class}">{prefix} {content}</span>'
+
+
+def _render_diff_span(span: model.DiffSpan, line_kind: str) -> str:
+    text = _escape_html(span.text)
+    if line_kind == "add" and span.op == "insert":
+        return f'<span class="diff-span diff-char-add">{text}</span>'
+    if line_kind == "remove" and span.op == "delete":
+        return f'<span class="diff-span diff-char-remove">{text}</span>'
+    return f'<span class="diff-span">{text}</span>'
+
+
+def _render_markdown_doc(doc: model.MarkdownDocUIExtra) -> str:
+    encoded = _escape_html(doc.content)
+    file_path = _escape_html(doc.file_path)
+    header = f'<div class="diff-file">{file_path} <span style="font-weight: normal; color: var(--text-dim); font-size: 12px; margin-left: 8px;">(markdown content)</span></div>'
+
+    # Using a container that mimics diff-view but for markdown
+    content = (
+        f'<div class="markdown-content markdown-body" data-raw="{encoded}" '
+        f'style="padding: 12px; border: 1px solid var(--border); border-radius: var(--radius-md); background: var(--bg-body); margin-top: 4px;">'
+        f'<noscript><pre style="white-space: pre-wrap;">{encoded}</pre></noscript>'
+        f"</div>"
+    )
+
+    line_count = doc.content.count("\n") + 1
+    open_attr = " open"
+
+    return (
+        f'<details class="diff-collapsible"{open_attr}>'
+        f"<summary>File Content ({line_count} lines)</summary>"
+        f'<div style="margin-top: 8px;">'
+        f"{header}"
+        f"{content}"
+        f"</div>"
+        f"</details>"
+    )
+
+
+def _collect_ui_extras(ui_extra: model.ToolResultUIExtra | None) -> list[model.ToolResultUIExtra]:
     if ui_extra is None:
-        return None
-    if ui_extra.type != model.ToolResultUIExtraType.DIFF_TEXT:
-        return None
-    return ui_extra.diff_text
+        return []
+    if isinstance(ui_extra, model.MultiUIExtra):
+        return list(ui_extra.items)
+    return [ui_extra]
+
+
+def _build_add_only_diff(text: str, file_path: str) -> model.DiffUIExtra:
+    lines: list[model.DiffLine] = []
+    new_line_no = 1
+    for line in text.splitlines():
+        lines.append(
+            model.DiffLine(
+                kind="add",
+                new_line_no=new_line_no,
+                spans=[model.DiffSpan(op="equal", text=line)],
+            )
+        )
+        new_line_no += 1
+    file_diff = model.DiffFileDiff(file_path=file_path, lines=lines, stats_add=len(lines), stats_remove=0)
+    return model.DiffUIExtra(files=[file_diff])
 
 
 def _get_mermaid_link_html(
     ui_extra: model.ToolResultUIExtra | None, tool_call: model.ToolCallItem | None = None
 ) -> str | None:
-    if tool_call and tool_call.name == "Mermaid":
+    code = ""
+    link: str | None = None
+    line_count = 0
+
+    if isinstance(ui_extra, model.MermaidLinkUIExtra):
+        code = ui_extra.code
+        link = ui_extra.link
+        line_count = ui_extra.line_count
+
+    if not code and tool_call and tool_call.name == "Mermaid":
         try:
             args = json.loads(tool_call.arguments)
             code = args.get("code", "")
-        except Exception:
+        except (json.JSONDecodeError, TypeError):
             code = ""
-    else:
-        code = ""
+        line_count = code.count("\n") + 1 if code else 0
 
-    if not code and (
-        ui_extra is None or ui_extra.type != model.ToolResultUIExtraType.MERMAID_LINK or not ui_extra.mermaid_link
-    ):
+    if not code and not link:
         return None
 
     # Prepare code for rendering and copy
     escaped_code = _escape_html(code) if code else ""
-    line_count = code.count("\n") + 1 if code else 0
 
     # Build Toolbar
     toolbar_items: list[str] = []
@@ -380,12 +513,9 @@ def _get_mermaid_link_html(
         buttons_html.append(
             f'<button type="button" class="copy-mermaid-btn" data-code="{escaped_code}" title="Copy Mermaid Code">Copy Code</button>'
         )
-
-    link = (
-        ui_extra.mermaid_link.link
-        if (ui_extra and ui_extra.type == model.ToolResultUIExtraType.MERMAID_LINK and ui_extra.mermaid_link)
-        else None
-    )
+        buttons_html.append(
+            '<button type="button" class="fullscreen-mermaid-btn" title="View Fullscreen">Fullscreen</button>'
+        )
 
     if link:
         link_url = _escape_html(link)
@@ -419,8 +549,8 @@ def _format_tool_call(tool_call: model.ToolCallItem, result: model.ToolResultIte
     is_todo_list = False
     ts_str = _format_msg_timestamp(tool_call.created_at)
 
-    if tool_call.name == "TodoWrite":
-        args_html = _try_render_todo_args(tool_call.arguments)
+    if tool_call.name in ("TodoWrite", "update_plan"):
+        args_html = _try_render_todo_args(tool_call.arguments, tool_call.name)
         if args_html:
             is_todo_list = True
 
@@ -428,7 +558,7 @@ def _format_tool_call(tool_call: model.ToolCallItem, result: model.ToolResultIte
         try:
             parsed = json.loads(tool_call.arguments)
             args_text = json.dumps(parsed, ensure_ascii=False, indent=2)
-        except Exception:
+        except (json.JSONDecodeError, TypeError):
             args_text = tool_call.arguments
 
         args_html = _escape_html(args_text or "")
@@ -450,7 +580,7 @@ def _format_tool_call(tool_call: model.ToolCallItem, result: model.ToolResultIte
                 parsed_args = json.loads(tool_call.arguments)
                 if parsed_args.get("command") in {"create", "str_replace", "insert"}:
                     force_collapse = True
-            except Exception:
+            except (json.JSONDecodeError, TypeError):
                 pass
 
         should_collapse = force_collapse or _should_collapse(args_html)
@@ -475,31 +605,50 @@ def _format_tool_call(tool_call: model.ToolCallItem, result: model.ToolResultIte
     ]
 
     if result:
-        diff_text = _get_diff_text(result.ui_extra)
-        mermaid_html = _get_mermaid_link_html(result.ui_extra, tool_call)
+        extras = _collect_ui_extras(result.ui_extra)
 
-        should_hide_text = tool_call.name == "TodoWrite" and result.status != "error"
+        mermaid_extra = next((x for x in extras if isinstance(x, model.MermaidLinkUIExtra)), None)
+        mermaid_source = mermaid_extra if mermaid_extra else result.ui_extra
+        mermaid_html = _get_mermaid_link_html(mermaid_source, tool_call)
 
-        if tool_call.name == "Edit" and not diff_text and result.status != "error":
+        should_hide_text = tool_call.name in ("TodoWrite", "update_plan") and result.status != "error"
+
+        if (
+            tool_call.name == "Edit"
+            and not any(isinstance(x, model.DiffUIExtra) for x in extras)
+            and result.status != "error"
+        ):
             try:
                 args_data = json.loads(tool_call.arguments)
+                file_path = args_data.get("file_path", "Unknown file")
                 old_string = args_data.get("old_string", "")
                 new_string = args_data.get("new_string", "")
                 if old_string == "" and new_string:
-                    diff_text = "\n".join(f"+{line}" for line in new_string.splitlines())
-            except Exception:
+                    extras.append(_build_add_only_diff(new_string, file_path))
+            except (json.JSONDecodeError, TypeError):
                 pass
 
         items_to_render: list[str] = []
 
         if result.output and not should_hide_text:
             if is_sub_agent_tool(tool_call.name):
-                items_to_render.append(_render_sub_agent_result(result.output))
+                description = None
+                try:
+                    args = json.loads(tool_call.arguments)
+                    if isinstance(args, dict):
+                        typed_args = cast(dict[str, Any], args)
+                        description = cast(str | None, typed_args.get("description"))
+                except (json.JSONDecodeError, TypeError):
+                    pass
+                items_to_render.append(_render_sub_agent_result(result.output, description))
             else:
                 items_to_render.append(_render_text_block(result.output))
 
-        if diff_text:
-            items_to_render.append(_render_diff_block(diff_text))
+        for extra in extras:
+            if isinstance(extra, model.DiffUIExtra):
+                items_to_render.append(_render_diff_block(extra))
+            elif isinstance(extra, model.MarkdownDocUIExtra):
+                items_to_render.append(_render_markdown_doc(extra))
 
         if mermaid_html:
             items_to_render.append(mermaid_html)
@@ -522,7 +671,13 @@ def _format_tool_call(tool_call: model.ToolCallItem, result: model.ToolResultIte
 def _build_messages_html(
     history: list[model.ConversationItem],
     tool_results: dict[str, model.ToolResultItem],
+    *,
+    seen_session_ids: set[str] | None = None,
+    nesting_level: int = 0,
 ) -> str:
+    if seen_session_ids is None:
+        seen_session_ids = set()
+
     blocks: list[str] = []
     assistant_counter = 0
 
@@ -549,7 +704,7 @@ def _build_messages_html(
         elif isinstance(item, model.AssistantMessageItem):
             assistant_counter += 1
             blocks.append(_render_assistant_message(assistant_counter, item.content or "", item.created_at))
-        elif isinstance(item, model.ResponseMetadataItem):
+        elif isinstance(item, model.TaskMetadataItem):
             blocks.append(_render_metadata_item(item))
         elif isinstance(item, model.DeveloperMessageItem):
             content = _escape_html(item.content or "")
@@ -574,9 +729,61 @@ def _build_messages_html(
             result = tool_results.get(item.call_id)
             blocks.append(_format_tool_call(item, result))
 
+            # Recursively render sub-agent session history
+            if result is not None:
+                sub_agent_html = _render_sub_agent_session(result, seen_session_ids, nesting_level)
+                if sub_agent_html:
+                    blocks.append(sub_agent_html)
+
     return "\n".join(blocks)
 
 
+def _render_sub_agent_session(
+    tool_result: model.ToolResultItem,
+    seen_session_ids: set[str],
+    nesting_level: int,
+) -> str | None:
+    """Render sub-agent session history when a tool result references it."""
+    from klaude_code.session.session import Session
+
+    ui_extra = tool_result.ui_extra
+    if not isinstance(ui_extra, model.SessionIdUIExtra):
+        return None
+
+    session_id = ui_extra.session_id
+    if not session_id or session_id in seen_session_ids:
+        return None
+
+    seen_session_ids.add(session_id)
+
+    try:
+        sub_session = Session.load(session_id)
+    except (OSError, json.JSONDecodeError, ValueError):
+        return None
+
+    sub_history = sub_session.conversation_history
+    sub_tool_results = {item.call_id: item for item in sub_history if isinstance(item, model.ToolResultItem)}
+
+    sub_html = _build_messages_html(
+        sub_history,
+        sub_tool_results,
+        seen_session_ids=seen_session_ids,
+        nesting_level=nesting_level + 1,
+    )
+
+    if not sub_html:
+        return None
+
+    # Wrap in a collapsible sub-agent container using same style as other collapsible sections
+    indent_style = f' style="margin-left: {nesting_level * 16}px;"' if nesting_level > 0 else ""
+    return (
+        f'<details class="sub-agent-session"{indent_style}>'
+        f"<summary>Sub-agent: {_escape_html(session_id)}</summary>"
+        f'<div class="sub-agent-content">{sub_html}</div>'
+        f"</details>"
+    )
+
+
 def build_export_html(
     session: Session,
     system_prompt: str,