klaude-code 1.2.7__py3-none-any.whl → 1.2.9__py3-none-any.whl

klaude_code/protocol/model.py

@@ -1,8 +1,8 @@
 from datetime import datetime
 from enum import Enum
-from typing import Literal
+from typing import Annotated, Literal
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field, computed_field
 
 from klaude_code.protocol.commands import CommandName
 from klaude_code.protocol.tools import SubAgentType
@@ -12,12 +12,16 @@ TodoStatusType = Literal["pending", "in_progress", "completed"]
 
 
 class Usage(BaseModel):
+    # Token Usage (primary state)
     input_tokens: int = 0
     cached_tokens: int = 0
     reasoning_tokens: int = 0
     output_tokens: int = 0
-    total_tokens: int = 0
-    context_usage_percent: float | None = None
+
+    # Context window tracking
+    context_window_size: int | None = None  # Peak total_tokens seen (for context usage display)
+    context_limit: int | None = None  # Model's context limit
+
     throughput_tps: float | None = None
     first_token_latency_ms: float | None = None
 
@@ -25,14 +29,39 @@ class Usage(BaseModel):
     input_cost: float | None = None  # Cost for non-cached input tokens
     output_cost: float | None = None  # Cost for output tokens (including reasoning)
     cache_read_cost: float | None = None  # Cost for cached tokens
-    total_cost: float | None = None  # Total cost (input + output + cache_read)
     currency: str = "USD"  # Currency for cost display (USD or CNY)
 
+    @computed_field  # type: ignore[prop-decorator]
+    @property
+    def total_tokens(self) -> int:
+        """Total tokens computed from input + output tokens."""
+        return self.input_tokens + self.output_tokens
+
+    @computed_field  # type: ignore[prop-decorator]
+    @property
+    def total_cost(self) -> float | None:
+        """Total cost computed from input + output + cache_read costs."""
+        costs = [self.input_cost, self.output_cost, self.cache_read_cost]
+        non_none = [c for c in costs if c is not None]
+        return sum(non_none) if non_none else None
+
+    @computed_field  # type: ignore[prop-decorator]
+    @property
+    def context_usage_percent(self) -> float | None:
+        """Context usage percentage computed from context_window_size / context_limit."""
+        if self.context_limit is None or self.context_limit <= 0:
+            return None
+        if self.context_window_size is None:
+            return None
+        return (self.context_window_size / self.context_limit) * 100
+
 
 class TodoItem(BaseModel):
+    model_config = ConfigDict(populate_by_name=True)
+
     content: str
     status: TodoStatusType
-    activeForm: str = ""
+    active_form: str = Field(default="", alias="activeForm")
 
 
 class TodoUIExtra(BaseModel):
@@ -40,43 +69,55 @@ class TodoUIExtra(BaseModel):
     new_completed: list[str]
 
 
-class ToolResultUIExtraType(str, Enum):
-    DIFF_TEXT = "diff_text"
-    TODO_LIST = "todo_list"
-    SESSION_ID = "session_id"
-    MERMAID_LINK = "mermaid_link"
-    TRUNCATION = "truncation"
-    SESSION_STATUS = "session_status"
-
-
 class ToolSideEffect(str, Enum):
     TODO_CHANGE = "todo_change"
 
 
+# Discriminated union types for ToolResultUIExtra
+class DiffTextUIExtra(BaseModel):
+    type: Literal["diff_text"] = "diff_text"
+    diff_text: str
+
+
+class TodoListUIExtra(BaseModel):
+    type: Literal["todo_list"] = "todo_list"
+    todo_list: TodoUIExtra
+
+
+class SessionIdUIExtra(BaseModel):
+    type: Literal["session_id"] = "session_id"
+    session_id: str
+
+
 class MermaidLinkUIExtra(BaseModel):
+    type: Literal["mermaid_link"] = "mermaid_link"
     link: str
     line_count: int
 
 
 class TruncationUIExtra(BaseModel):
+    type: Literal["truncation"] = "truncation"
     saved_file_path: str
     original_length: int
     truncated_length: int
 
 
 class SessionStatusUIExtra(BaseModel):
+    type: Literal["session_status"] = "session_status"
     usage: "Usage"
     task_count: int
+    by_model: list["TaskMetadata"] = []
 
 
-class ToolResultUIExtra(BaseModel):
-    type: ToolResultUIExtraType
-    diff_text: str | None = None
-    todo_list: TodoUIExtra | None = None
-    session_id: str | None = None
-    mermaid_link: MermaidLinkUIExtra | None = None
-    truncation: TruncationUIExtra | None = None
-    session_status: SessionStatusUIExtra | None = None
+ToolResultUIExtra = Annotated[
+    DiffTextUIExtra
+    | TodoListUIExtra
+    | SessionIdUIExtra
+    | MermaidLinkUIExtra
+    | TruncationUIExtra
+    | SessionStatusUIExtra,
+    Field(discriminator="type"),
+]
 
 
 class AtPatternParseResult(BaseModel):
@@ -240,6 +281,7 @@ class ToolResultItem(BaseModel):
     ui_extra: ToolResultUIExtra | None = None  # Extra data for UI display, e.g. diff render
     images: list[ImageURLPart] | None = None
     side_effects: list[ToolSideEffect] | None = None
+    task_metadata: "TaskMetadata | None" = None  # Sub-agent task metadata for propagation to main agent
     created_at: datetime = Field(default_factory=datetime.now)
 
 
@@ -255,13 +297,80 @@ class StreamErrorItem(BaseModel):
 
 
 class ResponseMetadataItem(BaseModel):
+    """Metadata for a single LLM response (turn-level)."""
+
     response_id: str | None = None
     usage: Usage | None = None
     model_name: str = ""
     provider: str | None = None  # OpenRouter's provider name
     task_duration_s: float | None = None
-    status: str | None = None
-    error_reason: str | None = None
+    created_at: datetime = Field(default_factory=datetime.now)
+
+
+class TaskMetadata(BaseModel):
+    """Base metadata for a task execution (used by both main and sub-agents)."""
+
+    usage: Usage | None = None
+    model_name: str = ""
+    provider: str | None = None
+    task_duration_s: float | None = None
+
+    @staticmethod
+    def aggregate_by_model(metadata_list: list["TaskMetadata"]) -> list["TaskMetadata"]:
+        """Aggregate multiple TaskMetadata by (model_name, provider).
+
+        Returns a list sorted by total_cost descending.
+
+        Note: total_tokens and total_cost are now computed fields,
+        so we only accumulate the primary state fields here.
+        """
+        aggregated: dict[tuple[str, str | None], TaskMetadata] = {}
+
+        for meta in metadata_list:
+            if not meta.usage:
+                continue
+
+            key = (meta.model_name, meta.provider)
+            usage = meta.usage
+
+            if key not in aggregated:
+                aggregated[key] = TaskMetadata(
+                    model_name=meta.model_name,
+                    provider=meta.provider,
+                    usage=Usage(currency=usage.currency),
+                )
+
+            agg = aggregated[key]
+            if agg.usage is None:
+                continue
+
+            # Accumulate primary token fields (total_tokens is computed)
+            agg.usage.input_tokens += usage.input_tokens
+            agg.usage.cached_tokens += usage.cached_tokens
+            agg.usage.reasoning_tokens += usage.reasoning_tokens
+            agg.usage.output_tokens += usage.output_tokens
+
+            # Accumulate cost components (total_cost is computed)
+            if usage.input_cost is not None:
+                agg.usage.input_cost = (agg.usage.input_cost or 0.0) + usage.input_cost
+            if usage.output_cost is not None:
+                agg.usage.output_cost = (agg.usage.output_cost or 0.0) + usage.output_cost
+            if usage.cache_read_cost is not None:
+                agg.usage.cache_read_cost = (agg.usage.cache_read_cost or 0.0) + usage.cache_read_cost
+
+        # Sort by total_cost descending
+        return sorted(
+            aggregated.values(),
+            key=lambda m: m.usage.total_cost if m.usage and m.usage.total_cost else 0.0,
+            reverse=True,
+        )
+
+
+class TaskMetadataItem(BaseModel):
+    """Aggregated metadata for a complete task, stored in conversation history."""
+
+    main: TaskMetadata = Field(default_factory=TaskMetadata)
+    sub_agent_task_metadata: list[TaskMetadata] = Field(default_factory=lambda: list[TaskMetadata]())
     created_at: datetime = Field(default_factory=datetime.now)
 
 
@@ -280,7 +389,14 @@ MessageItem = (
 StreamItem = AssistantMessageDelta
 
 ConversationItem = (
-    StartItem | InterruptItem | StreamErrorItem | StreamItem | MessageItem | ResponseMetadataItem | ToolCallStartItem
+    StartItem
+    | InterruptItem
+    | StreamErrorItem
+    | StreamItem
+    | MessageItem
+    | ResponseMetadataItem
+    | TaskMetadataItem
+    | ToolCallStartItem
 )
 
 

klaude_code/protocol/sub_agent.py

@@ -1,10 +1,13 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Any, Callable
+from typing import TYPE_CHECKING, Any, Callable
 
 from klaude_code.protocol import tools
 
+if TYPE_CHECKING:
+    from klaude_code.protocol import model
+
 AvailabilityPredicate = Callable[[str], bool]
 PromptBuilder = Callable[[dict[str, Any]], str]
 
@@ -14,6 +17,7 @@ class SubAgentResult:
     task_result: str
     session_id: str
     error: bool = False
+    task_metadata: model.TaskMetadata | None = None
 
 
 def _default_prompt_builder(args: dict[str, Any]) -> str:

klaude_code/session/export.py

@@ -159,20 +159,35 @@ def _format_cost(cost: float, currency: str = "USD") -> str:
     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:
@@ -194,22 +209,39 @@ def _render_metadata_item(item: model.ResponseMetadataItem) -> str:
 
         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:
+        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, item.usage.currency)}</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"><div class="metadata-line">{joined_html}</div></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, 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:
@@ -336,11 +368,9 @@ def _render_diff_block(diff: str) -> str:
 
 
 def _get_diff_text(ui_extra: model.ToolResultUIExtra | None) -> str | None:
-    if ui_extra is None:
-        return None
-    if ui_extra.type != model.ToolResultUIExtraType.DIFF_TEXT:
-        return None
-    return ui_extra.diff_text
+    if isinstance(ui_extra, model.DiffTextUIExtra):
+        return ui_extra.diff_text
+    return None
 
 
 def _get_mermaid_link_html(
@@ -355,9 +385,7 @@ def _get_mermaid_link_html(
     else:
         code = ""
 
-    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 isinstance(ui_extra, model.MermaidLinkUIExtra):
         return None
 
     # Prepare code for rendering and copy
@@ -376,11 +404,7 @@ def _get_mermaid_link_html(
             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
-    )
+    link = ui_extra.link if isinstance(ui_extra, model.MermaidLinkUIExtra) else None
 
     if link:
         link_url = _escape_html(link)
@@ -544,7 +568,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 "")

klaude_code/session/session.py

@@ -5,7 +5,7 @@ from collections.abc import Iterable, Sequence
 from pathlib import Path
 from typing import ClassVar
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, PrivateAttr
 
 from klaude_code.protocol import events, model
 
@@ -19,8 +19,6 @@ class Session(BaseModel):
     file_tracker: dict[str, float] = Field(default_factory=dict)
     # Todo list for the session
     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
     # Used in list method SessionMetaBrief
     model_name: str | None = None
@@ -33,6 +31,27 @@ class Session(BaseModel):
     need_todo_empty_cooldown_counter: int = Field(exclude=True, default=0)
     need_todo_not_used_cooldown_counter: int = Field(exclude=True, default=0)
 
+    # Cached messages count (computed property)
+    _messages_count_cache: int | None = PrivateAttr(default=None)
+
+    @property
+    def messages_count(self) -> int:
+        """Count of user and assistant messages in conversation history.
+
+        This is a cached property that is invalidated when append_history is called.
+        """
+        if self._messages_count_cache is None:
+            self._messages_count_cache = sum(
+                1
+                for it in self.conversation_history
+                if isinstance(it, (model.UserMessageItem, model.AssistantMessageItem))
+            )
+        return self._messages_count_cache
+
+    def _invalidate_messages_count_cache(self) -> None:
+        """Invalidate the cached messages count."""
+        self._messages_count_cache = None
+
     # Internal: mapping for (de)serialization of conversation items
     _TypeMap: ClassVar[dict[str, type[BaseModel]]] = {
         # Messages
@@ -50,7 +69,7 @@ class Session(BaseModel):
         "AssistantMessageDelta": model.AssistantMessageDelta,
         "StartItem": model.StartItem,
         "StreamErrorItem": model.StreamErrorItem,
-        "ResponseMetadataItem": model.ResponseMetadataItem,
+        "TaskMetadataItem": model.TaskMetadataItem,
         "InterruptItem": model.InterruptItem,
     }
 
@@ -109,7 +128,6 @@ class Session(BaseModel):
         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))
-        messages_count = int(raw.get("messages_count", 0))
         model_name = raw.get("model_name")
 
         sess = Session(
@@ -121,7 +139,6 @@ class Session(BaseModel):
             loaded_memory=loaded_memory,
             created_at=created_at,
             updated_at=updated_at,
-            messages_count=messages_count,
             model_name=model_name,
         )
 
@@ -154,10 +171,7 @@ class Session(BaseModel):
                     # Best-effort load; skip malformed lines
                     continue
             sess.conversation_history = history
-            # Update messages count based on loaded history (only UserMessageItem and AssistantMessageItem)
-            sess.messages_count = sum(
-                1 for it in history if isinstance(it, (model.UserMessageItem, model.AssistantMessageItem))
-            )
+            # messages_count is now a computed property, no need to set it
 
         return sess
 
@@ -190,10 +204,8 @@ class Session(BaseModel):
     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)
-        self.messages_count += sum(
-            1 for it in items if isinstance(it, (model.UserMessageItem, model.AssistantMessageItem))
-        )
+        # Invalidate messages count cache
+        self._invalidate_messages_count_cache()
 
         # Incrementally persist to JSONL under messages directory
         messages_dir = self._messages_dir()
@@ -295,8 +307,8 @@ class Session(BaseModel):
                         content=ri.content,
                         session_id=self.id,
                     )
-                case model.ResponseMetadataItem() as mt:
-                    yield events.ResponseMetadataEvent(
+                case model.TaskMetadataItem() as mt:
+                    yield events.TaskMetadataEvent(
                         session_id=self.id,
                         metadata=mt,
                     )
@@ -309,6 +321,11 @@ class Session(BaseModel):
                         session_id=self.id,
                         item=dm,
                     )
+                case model.StreamErrorItem() as se:
+                    yield events.ErrorEvent(
+                        error_message=se.error,
+                        can_retry=False,
+                    )
                 case _:
                     continue
             prev_item = it

klaude_code/session/templates/export_session.html

@@ -21,7 +21,7 @@
       rel="stylesheet"
     />
     <link
-      href="https://fonts.googleapis.com/css2?family=IBM+Plex+Mono:wght@400;500;700&family=IBM+Plex+Sans:wght@400;500;700&display=swap"
+      href="https://fonts.googleapis.com/css2?family=IBM+Plex+Mono:ital,wght@0,400;0,500;0,700;1,400;1,500;1,700&family=IBM+Plex+Sans:ital,wght@0,400;0,500;0,700;1,400;1,500;1,700&display=swap"
       rel="stylesheet"
     />
     <style>
@@ -411,6 +411,9 @@
         font-size: var(--font-size-xs);
         color: var(--text-dim);
         border-left: 2px solid transparent;
+        display: flex;
+        flex-direction: column;
+        gap: 8px;
       }
       .metadata-line {
         display: flex;

klaude_code/ui/modes/repl/__init__.py

@@ -22,11 +22,7 @@ def build_repl_status_snapshot(agent: "Agent | None", update_message: str | None
     tool_calls = 0
 
     if agent is not None:
-        profile = agent.profile
-        if profile is not None:
-            model_name = profile.llm_client.model_name or ""
-        else:
-            model_name = "N/A"
+        model_name = agent.profile.llm_client.model_name or ""
 
         history = agent.session.conversation_history
         for item in history: