klaude-code 1.2.1__py3-none-any.whl → 1.2.3__py3-none-any.whl

klaude_code/session/export.py

@@ -11,21 +11,8 @@ from pathlib import Path
 from string import Template
 from typing import TYPE_CHECKING, Any, Final, cast
 
-from klaude_code.core.sub_agent import is_sub_agent_tool
-from klaude_code.protocol.llm_parameter import ToolSchema
-from klaude_code.protocol.model import (
-    AssistantMessageItem,
-    ConversationItem,
-    DeveloperMessageItem,
-    ReasoningEncryptedItem,
-    ReasoningTextItem,
-    ResponseMetadataItem,
-    ToolCallItem,
-    ToolResultItem,
-    ToolResultUIExtra,
-    ToolResultUIExtraType,
-    UserMessageItem,
-)
+from klaude_code.protocol import llm_param, model
+from klaude_code.protocol.sub_agent import is_sub_agent_tool
 
 if TYPE_CHECKING:
     from klaude_code.session.session import Session
@@ -58,10 +45,14 @@ def _format_timestamp(value: float | None) -> str:
     return datetime.fromtimestamp(value).strftime("%Y-%m-%d %H:%M:%S")
 
 
-def get_first_user_message(history: list[ConversationItem]) -> str:
+def _format_msg_timestamp(dt: datetime) -> str:
+    return dt.strftime("%Y-%m-%d %H:%M:%S")
+
+
+def get_first_user_message(history: list[model.ConversationItem]) -> str:
     """Extract the first user message content from conversation history."""
     for item in history:
-        if isinstance(item, UserMessageItem) and item.content:
+        if isinstance(item, model.UserMessageItem) and item.content:
             content = item.content.strip()
             first_line = content.split("\n")[0]
             return first_line[:100] if len(first_line) > 100 else first_line
@@ -86,7 +77,7 @@ def _load_template() -> str:
     return template_file.read_text(encoding="utf-8")
 
 
-def _build_tools_html(tools: list[ToolSchema]) -> str:
+def _build_tools_html(tools: list[llm_param.ToolSchema]) -> str:
     if not tools:
         return '<div style="padding: 12px; font-style: italic;">No tools registered for this session.</div>'
     chunks: list[str] = []
@@ -163,7 +154,11 @@ def _format_token_count(count: int) -> str:
     return f"{m}M" if rem == 0 else f"{m}M{rem}k"
 
 
-def _render_metadata_item(item: ResponseMetadataItem) -> str:
+def _format_cost(cost: float) -> str:
+    return f"${cost:.4f}"
+
+
+def _render_metadata_item(item: model.ResponseMetadataItem) -> str:
     # Line 1: Model Name [@ Provider]
     model_parts = [f'<span class="metadata-model">{_escape_html(item.model_name)}</span>']
     if item.provider:
@@ -176,10 +171,25 @@ def _render_metadata_item(item: ResponseMetadataItem) -> str:
     stats_parts: list[str] = []
     if item.usage:
         u = item.usage
-        stats_parts.append(f'<span class="metadata-stat">input: {_format_token_count(u.input_tokens)}</span>')
+        # 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)})"
+        stats_parts.append(f'<span class="metadata-stat">{input_stat}</span>')
+
+        # Cached with cost
         if u.cached_tokens > 0:
-            stats_parts.append(f'<span class="metadata-stat">cached: {_format_token_count(u.cached_tokens)}</span>')
-        stats_parts.append(f'<span class="metadata-stat">output: {_format_token_count(u.output_tokens)}</span>')
+            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)})"
+            stats_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)})"
+        stats_parts.append(f'<span class="metadata-stat">{output_stat}</span>')
+
         if u.reasoning_tokens > 0:
             stats_parts.append(
                 f'<span class="metadata-stat">thinking: {_format_token_count(u.reasoning_tokens)}</span>'
@@ -190,7 +200,11 @@ def _render_metadata_item(item: ResponseMetadataItem) -> str:
             stats_parts.append(f'<span class="metadata-stat">tps: {u.throughput_tps:.1f}</span>')
 
     if item.task_duration_s is not None:
-        stats_parts.append(f'<span class="metadata-stat">cost: {item.task_duration_s:.1f}s</span>')
+        stats_parts.append(f'<span class="metadata-stat">time: {item.task_duration_s:.1f}s</span>')
+
+    # Total cost
+    if item.usage is not None and item.usage.total_cost is not None:
+        stats_parts.append(f'<span class="metadata-stat">cost: {_format_cost(item.usage.total_cost)}</span>')
 
     stats_html = ""
     if stats_parts:
@@ -205,13 +219,15 @@ def _render_metadata_item(item: ResponseMetadataItem) -> str:
     )
 
 
-def _render_assistant_message(index: int, content: str) -> str:
+def _render_assistant_message(index: int, content: str, timestamp: datetime) -> str:
     encoded = _escape_html(content)
+    ts_str = _format_msg_timestamp(timestamp)
     return (
         f'<div class="message-group assistant-message-group">'
         f'<div class="message-header">'
         f'<div class="role-label assistant">Assistant</div>'
         f'<div class="assistant-toolbar">'
+        f'<span class="timestamp">{ts_str}</span>'
         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>"
@@ -289,9 +305,9 @@ def _render_text_block(text: str) -> str:
     return (
         f'<div class="expandable-output expandable">'
         f'<div class="preview-text" style="white-space: pre-wrap; font-family: var(--font-mono);">{preview}</div>'
-        f'<div class="expand-hint expand-text">Click to expand full output ({len(lines)} lines)</div>'
+        f'<div class="expand-hint expand-text">click to expand full output ({len(lines)} lines)</div>'
         f'<div class="full-text" style="white-space: pre-wrap; font-family: var(--font-mono);">{full}</div>'
-        f'<div class="collapse-hint">Click to collapse</div>'
+        f'<div class="collapse-hint">click to collapse</div>'
         f"</div>"
     )
 
@@ -326,49 +342,85 @@ def _render_diff_block(diff: str) -> str:
     )
 
 
-def _get_diff_text(ui_extra: ToolResultUIExtra | None) -> str | None:
+def _get_diff_text(ui_extra: model.ToolResultUIExtra | None) -> str | None:
     if ui_extra is None:
         return None
-    if ui_extra.type != ToolResultUIExtraType.DIFF_TEXT:
+    if ui_extra.type != model.ToolResultUIExtraType.DIFF_TEXT:
         return None
     return ui_extra.diff_text
 
 
-def _get_mermaid_link_html(ui_extra: ToolResultUIExtra | None, tool_call: ToolCallItem | None = None) -> str | None:
-    if ui_extra is None:
-        return None
-    if ui_extra.type != ToolResultUIExtraType.MERMAID_LINK:
-        return None
-    if ui_extra.mermaid_link is None or not ui_extra.mermaid_link.link:
-        return None
-    link = _escape_html(ui_extra.mermaid_link.link)
-    lines = ui_extra.mermaid_link.line_count
-
-    copy_btn = ""
+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":
         try:
             args = json.loads(tool_call.arguments)
-            code = args.get("code")
-            if code:
-                escaped_code = _escape_html(code)
-                copy_btn = f'<button type="button" class="copy-mermaid-btn" data-code="{escaped_code}" title="Copy Mermaid Code">Copy Code</button>'
+            code = args.get("code", "")
         except Exception:
-            pass
+            code = ""
+    else:
+        code = ""
 
-    return (
-        '<div style="display: flex; justify-content: space-between; align-items: center; font-family: var(--font-mono);">'
-        f"<span>Lines: {lines}</span>"
-        f"<div>"
-        f"{copy_btn}"
-        f'<a href="{link}" target="_blank" rel="noopener noreferrer" style="color: var(--accent); text-decoration: underline; margin-left: 8px;">View Diagram</a>'
-        f"</div>"
+    if not code and (
+        ui_extra is None or ui_extra.type != model.ToolResultUIExtraType.MERMAID_LINK or not ui_extra.mermaid_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] = []
+
+    if line_count > 0:
+        toolbar_items.append(f"<span>Lines: {line_count}</span>")
+
+    buttons_html: list[str] = []
+    if code:
+        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
+    )
+
+    if link:
+        link_url = _escape_html(link)
+        buttons_html.append(
+            f'<a href="{link_url}" target="_blank" rel="noopener noreferrer" style="color: var(--accent); text-decoration: underline; margin-left: 8px;">View Online</a>'
+        )
+
+    toolbar_items.append(f"<div>{''.join(buttons_html)}</div>")
+
+    toolbar_html = (
+        '<div style="display: flex; justify-content: space-between; align-items: center; font-family: var(--font-mono); margin-top: 8px; padding-top: 8px; border-top: 1px dashed var(--border);">'
+        f"{''.join(toolbar_items)}"
         "</div>"
     )
 
+    # If we have code, render the diagram
+    if code:
+        return (
+            f'<div style="background: white; padding: 16px; border-radius: 4px; margin-top: 8px; border: 1px solid var(--border);">'
+            f'<div class="mermaid">{escaped_code}</div>'
+            f"{toolbar_html}"
+            f"</div>"
+        )
+
+    # Fallback to just link/toolbar if no code available (legacy support behavior)
+    return toolbar_html
+
 
-def _format_tool_call(tool_call: ToolCallItem, result: ToolResultItem | None) -> str:
+def _format_tool_call(tool_call: model.ToolCallItem, result: model.ToolResultItem | None) -> str:
     args_html = None
     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 args_html:
@@ -390,7 +442,21 @@ def _format_tool_call(tool_call: ToolCallItem, result: ToolResultItem | None) ->
     if is_todo_list:
         args_section = f'<div class="tool-args">{args_html}</div>'
     else:
-        open_attr = "" if _should_collapse(args_html) else " open"
+        # Always collapse Mermaid, Edit, Write tools by default
+        always_collapse_tools = {"Mermaid", "Edit", "Write"}
+        force_collapse = tool_call.name in always_collapse_tools
+
+        # Collapse Memory tool for write operations
+        if tool_call.name == "Memory":
+            try:
+                parsed_args = json.loads(tool_call.arguments)
+                if parsed_args.get("command") in {"create", "str_replace", "insert"}:
+                    force_collapse = True
+            except Exception:
+                pass
+
+        should_collapse = force_collapse or _should_collapse(args_html)
+        open_attr = "" if should_collapse else " open"
         args_section = (
             f'<details class="tool-args-collapsible"{open_attr}>'
             "<summary>Arguments</summary>"
@@ -402,7 +468,10 @@ def _format_tool_call(tool_call: ToolCallItem, result: ToolResultItem | None) ->
         '<div class="tool-call">',
         '<div class="tool-header">',
         f'<span class="tool-name">{_escape_html(tool_call.name)}</span>',
+        '<div class="tool-header-right">',
         f'<span class="tool-id">{_escape_html(tool_call.call_id)}</span>',
+        f'<span class="timestamp">{ts_str}</span>',
+        "</div>",
         "</div>",
         args_section,
     ]
@@ -453,46 +522,57 @@ def _format_tool_call(tool_call: ToolCallItem, result: ToolResultItem | None) ->
 
 
 def _build_messages_html(
-    history: list[ConversationItem],
-    tool_results: dict[str, ToolResultItem],
+    history: list[model.ConversationItem],
+    tool_results: dict[str, model.ToolResultItem],
 ) -> str:
     blocks: list[str] = []
     assistant_counter = 0
 
-    renderable_items = [item for item in history if not isinstance(item, (ToolResultItem, ReasoningEncryptedItem))]
+    renderable_items = [
+        item for item in history if not isinstance(item, (model.ToolResultItem, model.ReasoningEncryptedItem))
+    ]
 
     for i, item in enumerate(renderable_items):
-        if isinstance(item, UserMessageItem):
+        if isinstance(item, model.UserMessageItem):
             text = _escape_html(item.content or "")
+            ts_str = _format_msg_timestamp(item.created_at)
             blocks.append(
                 f'<div class="message-group">'
-                f'<div class="role-label user">User</div>'
+                f'<div class="role-label user">'
+                f"User"
+                f'<span class="timestamp">{ts_str}</span>'
+                f"</div>"
                 f'<div class="message-content user" style="white-space: pre-wrap;">{text}</div>'
                 f"</div>"
             )
-        elif isinstance(item, ReasoningTextItem):
+        elif isinstance(item, model.ReasoningTextItem):
             text = _escape_html(item.content.strip())
-            blocks.append(f'<div class="thinking-block">{text.replace(chr(10), "<br>")}</div>')
-        elif isinstance(item, AssistantMessageItem):
+            blocks.append(f'<div class="thinking-block markdown-body markdown-content" data-raw="{text}"></div>')
+        elif isinstance(item, model.AssistantMessageItem):
             assistant_counter += 1
-            blocks.append(_render_assistant_message(assistant_counter, item.content or ""))
-        elif isinstance(item, ResponseMetadataItem):
+            blocks.append(_render_assistant_message(assistant_counter, item.content or "", item.created_at))
+        elif isinstance(item, model.ResponseMetadataItem):
             blocks.append(_render_metadata_item(item))
-        elif isinstance(item, DeveloperMessageItem):
+        elif isinstance(item, model.DeveloperMessageItem):
             content = _escape_html(item.content or "")
+            ts_str = _format_msg_timestamp(item.created_at)
 
             next_item = renderable_items[i + 1] if i + 1 < len(renderable_items) else None
             extra_class = ""
-            if isinstance(next_item, (UserMessageItem, AssistantMessageItem)):
+            if isinstance(next_item, (model.UserMessageItem, model.AssistantMessageItem)):
                 extra_class = " gap-below"
 
             blocks.append(
                 f'<details class="developer-message{extra_class}">'
-                f"<summary>Developer</summary>"
+                f"<summary>"
+                f"Developer"
+                f'<span class="timestamp">{ts_str}</span>'
+                f"</summary>"
                 f'<div class="details-content" style="white-space: pre-wrap;">{content}</div>'
                 f"</details>"
             )
-        elif isinstance(item, ToolCallItem):
+
+        elif isinstance(item, model.ToolCallItem):
             result = tool_results.get(item.call_id)
             blocks.append(_format_tool_call(item, result))
 
@@ -502,7 +582,7 @@ def _build_messages_html(
 def build_export_html(
     session: Session,
     system_prompt: str,
-    tools: list[ToolSchema],
+    tools: list[llm_param.ToolSchema],
     model_name: str,
 ) -> str:
     """Build HTML export for a session.
@@ -517,7 +597,7 @@ def build_export_html(
         Complete HTML document as a string.
     """
     history = session.conversation_history
-    tool_results = {item.call_id: item for item in history if isinstance(item, ToolResultItem)}
+    tool_results = {item.call_id: item for item in history if isinstance(item, model.ToolResultItem)}
     messages_html = _build_messages_html(history, tool_results)
     if not messages_html:
         messages_html = '<div class="text-dim p-4 italic">No messages recorded for this session yet.</div>'
@@ -526,7 +606,7 @@ def build_export_html(
     session_id = session.id
     session_updated = _format_timestamp(session.updated_at)
     work_dir = _shorten_path(str(session.work_dir))
-    total_messages = len([item for item in history if not isinstance(item, ToolResultItem)])
+    total_messages = len([item for item in history if not isinstance(item, model.ToolResultItem)])
     footer_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     first_user_message = get_first_user_message(history)
 

klaude_code/session/session.py

@@ -8,18 +8,17 @@ from typing import ClassVar
 from pydantic import BaseModel, Field
 
 from klaude_code.protocol import events, model
-from klaude_code.protocol.model import ConversationItem, SubAgentState, TodoItem
 
 
 class Session(BaseModel):
     id: str = Field(default_factory=lambda: uuid.uuid4().hex)
     work_dir: Path
-    conversation_history: list[ConversationItem] = Field(default_factory=list)  # pyright: ignore[reportUnknownVariableType]
-    sub_agent_state: SubAgentState | None = None
+    conversation_history: list[model.ConversationItem] = Field(default_factory=list)  # pyright: ignore[reportUnknownVariableType]
+    sub_agent_state: model.SubAgentState | None = None
     # FileTracker: track file path -> last modification time when last read/edited
     file_tracker: dict[str, float] = Field(default_factory=dict)
     # Todo list for the session
-    todos: list[TodoItem] = Field(default_factory=list)  # pyright: ignore[reportUnknownVariableType]
+    todos: list[model.TodoItem] = Field(default_factory=list)  # pyright: ignore[reportUnknownVariableType]
     # Messages count, redundant state for performance optimization to avoid reading entire jsonl file
     messages_count: int = Field(default=0)
     # Model name used for this session
@@ -88,7 +87,11 @@ class Session(BaseModel):
     def load(cls, id: str) -> "Session":
         # Load session metadata
         sessions_dir = cls._sessions_dir()
-        session_candidates = sorted(sessions_dir.glob(f"*-{id}.json"), key=lambda p: p.stat().st_mtime, reverse=True)
+        session_candidates = sorted(
+            sessions_dir.glob(f"*-{id}.json"),
+            key=lambda p: p.stat().st_mtime,
+            reverse=True,
+        )
         if not session_candidates:
             # No existing session; create a new one
             return Session(id=id, work_dir=Path.cwd())
@@ -100,9 +103,9 @@ class Session(BaseModel):
         work_dir_str = raw.get("work_dir", str(Path.cwd()))
 
         sub_agent_state_raw = raw.get("sub_agent_state")
-        sub_agent_state = SubAgentState(**sub_agent_state_raw) if sub_agent_state_raw else None
+        sub_agent_state = model.SubAgentState(**sub_agent_state_raw) if sub_agent_state_raw else None
         file_tracker = dict(raw.get("file_tracker", {}))
-        todos: list[TodoItem] = [TodoItem(**item) for item in raw.get("todos", [])]
+        todos: list[model.TodoItem] = [model.TodoItem(**item) for item in raw.get("todos", [])]
         loaded_memory = list(raw.get("loaded_memory", []))
         created_at = float(raw.get("created_at", time.time()))
         updated_at = float(raw.get("updated_at", created_at))
@@ -125,10 +128,14 @@ class Session(BaseModel):
         # Load conversation history from messages JSONL
         messages_dir = cls._messages_dir()
         # Expect a single messages file per session (prefixed filenames only)
-        msg_candidates = sorted(messages_dir.glob(f"*-{id}.jsonl"), key=lambda p: p.stat().st_mtime, reverse=True)
+        msg_candidates = sorted(
+            messages_dir.glob(f"*-{id}.jsonl"),
+            key=lambda p: p.stat().st_mtime,
+            reverse=True,
+        )
         if msg_candidates:
             messages_path = msg_candidates[0]
-            history: list[ConversationItem] = []
+            history: list[model.ConversationItem] = []
             for line in messages_path.read_text().splitlines():
                 line = line.strip()
                 if not line:
@@ -180,7 +187,7 @@ class Session(BaseModel):
         }
         self._session_file().write_text(json.dumps(payload, ensure_ascii=False, indent=2))
 
-    def append_history(self, items: Sequence[ConversationItem]):
+    def append_history(self, items: Sequence[model.ConversationItem]):
         # Append to in-memory history
         self.conversation_history.extend(items)
         # Update messages count (only UserMessageItem and AssistantMessageItem)
@@ -197,7 +204,7 @@ class Session(BaseModel):
             for it in items:
                 # Serialize with explicit type tag for reliable load
                 t = it.__class__.__name__
-                data = it.model_dump()
+                data = it.model_dump(mode="json")
                 f.write(json.dumps({"type": t, "data": data}, ensure_ascii=False))
                 f.write("\n")
         # Refresh metadata timestamp after history change
@@ -236,7 +243,10 @@ class Session(BaseModel):
             return False
         if prev_item is None:
             return True
-        if isinstance(prev_item, model.UserMessageItem | model.ToolResultItem | model.DeveloperMessageItem):
+        if isinstance(
+            prev_item,
+            model.UserMessageItem | model.ToolResultItem | model.DeveloperMessageItem,
+        ):
             return True
         return False
 
@@ -337,7 +347,9 @@ class Session(BaseModel):
             if not msg_file.exists():
                 # Try to find by pattern if exact file doesn't exist
                 msg_candidates = sorted(
-                    messages_dir.glob(f"*-{session_id}.jsonl"), key=lambda p: p.stat().st_mtime, reverse=True
+                    messages_dir.glob(f"*-{session_id}.jsonl"),
+                    key=lambda p: p.stat().st_mtime,
+                    reverse=True,
                 )
                 if not msg_candidates:
                     return None
@@ -358,7 +370,9 @@ class Session(BaseModel):
                             # Handle structured content - extract text
                             text_parts: list[str] = []
                             for part in content:  # pyright: ignore[reportUnknownVariableType]
-                                if isinstance(part, dict) and part.get("type") == "text":  # pyright: ignore[reportUnknownMemberType]
+                                if (
+                                    isinstance(part, dict) and part.get("type") == "text"  # pyright: ignore[reportUnknownMemberType]
+                                ):
                                     text = part.get("text", "")  # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType]
                                     if isinstance(text, str):
                                         text_parts.append(text)