klaude-code 2.5.2__py3-none-any.whl → 2.6.0__py3-none-any.whl

klaude_code/protocol/events.py

@@ -0,0 +1,214 @@
+from __future__ import annotations
+
+import time
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+from klaude_code.protocol import llm_param, message, model
+from klaude_code.protocol.commands import CommandName
+
+__all__ = [
+    "AssistantImageDeltaEvent",
+    "AssistantTextDeltaEvent",
+    "AssistantTextEndEvent",
+    "AssistantTextStartEvent",
+    "CommandOutputEvent",
+    "DeveloperMessageEvent",
+    "EndEvent",
+    "ErrorEvent",
+    "Event",
+    "InterruptEvent",
+    "ReplayEventUnion",
+    "ReplayHistoryEvent",
+    "ResponseCompleteEvent",
+    "ResponseEvent",
+    "TaskFinishEvent",
+    "TaskMetadataEvent",
+    "TaskStartEvent",
+    "ThinkingDeltaEvent",
+    "ThinkingEndEvent",
+    "ThinkingStartEvent",
+    "TodoChangeEvent",
+    "ToolCallEvent",
+    "ToolCallStartEvent",
+    "ToolResultEvent",
+    "TurnEndEvent",
+    "TurnStartEvent",
+    "UsageEvent",
+    "UserMessageEvent",
+    "WelcomeEvent",
+]
+
+
+class Event(BaseModel):
+    """Base event."""
+
+    session_id: str
+    timestamp: float = Field(default_factory=time.time)
+
+
+class ResponseEvent(Event):
+    """Event associated with a single model response."""
+
+    response_id: str | None = None
+
+
+class UserMessageEvent(Event):
+    content: str
+    images: list[message.ImageURLPart] | None = None
+
+
+class DeveloperMessageEvent(Event):
+    """DeveloperMessages are reminders in user messages or tool results."""
+
+    item: message.DeveloperMessage
+
+
+class TodoChangeEvent(Event):
+    todos: list[model.TodoItem]
+
+
+class CommandOutputEvent(Event):
+    """Event for command output display. Not persisted to session history."""
+
+    command_name: CommandName | str
+    content: str = ""
+    ui_extra: model.ToolResultUIExtra | None = None
+    is_error: bool = False
+
+
+class TaskStartEvent(Event):
+    sub_agent_state: model.SubAgentState | None = None
+    model_id: str | None = None
+
+
+class TaskFinishEvent(Event):
+    task_result: str
+    has_structured_output: bool = False
+
+
+class TurnStartEvent(Event):
+    pass
+
+
+class TurnEndEvent(Event):
+    pass
+
+
+class UsageEvent(ResponseEvent):
+    usage: model.Usage
+
+
+class TaskMetadataEvent(Event):
+    metadata: model.TaskMetadataItem
+    cancelled: bool = False
+
+
+class ThinkingStartEvent(ResponseEvent):
+    pass
+
+
+class ThinkingDeltaEvent(ResponseEvent):
+    content: str
+
+
+class ThinkingEndEvent(ResponseEvent):
+    pass
+
+
+class AssistantTextStartEvent(ResponseEvent):
+    pass
+
+
+class AssistantTextDeltaEvent(ResponseEvent):
+    content: str
+
+
+class AssistantTextEndEvent(ResponseEvent):
+    pass
+
+
+class AssistantImageDeltaEvent(ResponseEvent):
+    file_path: str
+
+
+class ToolCallStartEvent(ResponseEvent):
+    tool_call_id: str
+    tool_name: str
+
+
+class ResponseCompleteEvent(ResponseEvent):
+    """Final snapshot of the model response."""
+
+    content: str
+    thinking_text: str | None = None
+
+
+class WelcomeEvent(Event):
+    work_dir: str
+    llm_config: llm_param.LLMConfigParameter
+    show_klaude_code_info: bool = True
+    loaded_skills: dict[str, list[str]] = Field(default_factory=dict)
+
+
+class ErrorEvent(Event):
+    error_message: str
+    can_retry: bool = False
+
+
+class InterruptEvent(Event):
+    pass
+
+
+class EndEvent(Event):
+    """Global display shutdown."""
+
+    session_id: str = "__app__"
+
+
+type ReplayEventUnion = (
+    TaskStartEvent
+    | TaskFinishEvent
+    | TurnStartEvent
+    | ThinkingStartEvent
+    | ThinkingDeltaEvent
+    | ThinkingEndEvent
+    | AssistantTextStartEvent
+    | AssistantTextDeltaEvent
+    | AssistantTextEndEvent
+    | AssistantImageDeltaEvent
+    | ToolCallEvent
+    | ToolResultEvent
+    | UserMessageEvent
+    | TaskMetadataEvent
+    | InterruptEvent
+    | DeveloperMessageEvent
+    | ErrorEvent
+)
+
+
+class ReplayHistoryEvent(Event):
+    events: list[ReplayEventUnion]
+    updated_at: float
+    is_load: bool = True
+
+
+class ToolCallEvent(ResponseEvent):
+    tool_call_id: str
+    tool_name: str
+    arguments: str
+
+
+class ToolResultEvent(ResponseEvent):
+    tool_call_id: str
+    tool_name: str
+    result: str
+    ui_extra: model.ToolResultUIExtra | None = None
+    status: Literal["success", "error", "aborted"]
+    task_metadata: model.TaskMetadata | None = None
+    is_last_in_turn: bool = True
+
+    @property
+    def is_error(self) -> bool:
+        return self.status in ("error", "aborted")

klaude_code/protocol/sub_agent/image_gen.py

@@ -66,10 +66,6 @@ IMAGE_GEN_PARAMETERS: dict[str, Any] = {
                     "enum": ["1K", "2K", "4K"],
                     "description": "Output size for Nano Banana Pro (must use uppercase K).",
                 },
-                "extra": {
-                    "type": "object",
-                    "description": "Provider/model-specific extra parameters (future-proofing).",
-                },
             },
             "additionalProperties": False,
         },

klaude_code/session/session.py

@@ -304,24 +304,57 @@ class Session(BaseModel):
                 yield events.TurnStartEvent(session_id=self.id)
             match it:
                 case message.AssistantMessage() as am:
-                    content = message.join_text_parts(am.parts)
-                    images = [part for part in am.parts if isinstance(part, message.ImageFilePart)]
-                    last_assistant_content = message.format_saved_images(images, content)
-                    thinking_text = "".join(
-                        part.text for part in am.parts if isinstance(part, message.ThinkingTextPart)
-                    )
-                    for image in images:
-                        yield events.AssistantImageDeltaEvent(
-                            file_path=image.file_path,
-                            response_id=am.response_id,
-                            session_id=self.id,
-                        )
-                    yield events.ResponseCompleteEvent(
-                        thinking_text=thinking_text or None,
-                        content=content,
-                        response_id=am.response_id,
-                        session_id=self.id,
-                    )
+                    all_images = [part for part in am.parts if isinstance(part, message.ImageFilePart)]
+                    full_content = message.join_text_parts(am.parts)
+                    last_assistant_content = message.format_saved_images(all_images, full_content)
+
+                    # Reconstruct streaming boundaries from saved parts.
+                    # This allows replay to reuse the same TUI state machine as live events.
+                    thinking_open = False
+                    assistant_open = False
+
+                    for part in am.parts:
+                        if isinstance(part, message.ThinkingTextPart):
+                            if assistant_open:
+                                assistant_open = False
+                                yield events.AssistantTextEndEvent(response_id=am.response_id, session_id=self.id)
+                            if not thinking_open:
+                                thinking_open = True
+                                yield events.ThinkingStartEvent(response_id=am.response_id, session_id=self.id)
+                            if part.text:
+                                yield events.ThinkingDeltaEvent(
+                                    content=part.text,
+                                    response_id=am.response_id,
+                                    session_id=self.id,
+                                )
+                            continue
+
+                        if thinking_open:
+                            thinking_open = False
+                            yield events.ThinkingEndEvent(response_id=am.response_id, session_id=self.id)
+
+                        if isinstance(part, message.TextPart):
+                            if not assistant_open:
+                                assistant_open = True
+                                yield events.AssistantTextStartEvent(response_id=am.response_id, session_id=self.id)
+                            if part.text:
+                                yield events.AssistantTextDeltaEvent(
+                                    content=part.text,
+                                    response_id=am.response_id,
+                                    session_id=self.id,
+                                )
+                        elif isinstance(part, message.ImageFilePart):
+                            yield events.AssistantImageDeltaEvent(
+                                file_path=part.file_path,
+                                response_id=am.response_id,
+                                session_id=self.id,
+                            )
+
+                    if thinking_open:
+                        yield events.ThinkingEndEvent(response_id=am.response_id, session_id=self.id)
+                    if assistant_open:
+                        yield events.AssistantTextEndEvent(response_id=am.response_id, session_id=self.id)
+
                     for part in am.parts:
                         if not isinstance(part, message.ToolCallPart):
                             continue

klaude_code/tui/command/fork_session_cmd.py

@@ -3,26 +3,23 @@ import sys
 from dataclasses import dataclass
 from typing import Literal
 
-from prompt_toolkit.styles import Style
+from prompt_toolkit.styles import Style, merge_styles
 
 from klaude_code.protocol import commands, events, message, model
 from klaude_code.tui.input.clipboard import copy_to_clipboard
-from klaude_code.tui.terminal.selector import SelectItem, select_one
+from klaude_code.tui.terminal.selector import DEFAULT_PICKER_STYLE, SelectItem, select_one
 
 from .command_abc import Agent, CommandABC, CommandResult
 
-FORK_SELECT_STYLE = Style(
+FORK_SELECT_STYLE = merge_styles(
     [
-        ("msg", ""),
-        ("meta", "fg:ansibrightblack"),
-        ("separator", "fg:ansibrightblack"),
-        ("assistant", "fg:ansiblue"),
-        ("pointer", "bold fg:ansigreen"),
-        ("search_prefix", "fg:ansibrightblack"),
-        ("search_success", "noinherit fg:ansigreen"),
-        ("search_none", "noinherit fg:ansired"),
-        ("question", "bold"),
-        ("text", ""),
+        DEFAULT_PICKER_STYLE,
+        Style(
+            [
+                ("separator", "fg:ansibrightblack"),
+                ("assistant", "fg:ansiblue"),
+            ]
+        ),
     ]
 )
 
@@ -144,6 +141,7 @@ def _build_select_items(fork_points: list[ForkPoint]) -> list[SelectItem[int]]:
                 title=title_parts,
                 value=fp.history_index,
                 search_text=fp.user_message if not is_last else "fork entire conversation",
+                selectable=not is_first,
             )
         )
 
@@ -163,6 +161,9 @@ def _select_fork_point_sync(fork_points: list[ForkPoint]) -> int | Literal["canc
 
     # Default to the last option (fork entire conversation)
     last_value = items[-1].value
+    if last_value is None:
+        # Should not happen as we populate all items with int values
+        return -1
 
     # Non-interactive environments default to forking entire conversation
     if not sys.stdin.isatty() or not sys.stdout.isatty():
@@ -241,16 +242,6 @@ class ForkSessionCommand(CommandABC):
             )
             return CommandResult(events=[event])
 
-        # First option (empty session) is just for UI display, not a valid fork point
-        if selected == fork_points[0].history_index:
-            event = events.CommandOutputEvent(
-                session_id=agent.session.id,
-                command_name=self.name,
-                content="(cannot fork to empty session)",
-                is_error=True,
-            )
-            return CommandResult(events=[event])
-
         # Perform the fork
         new_session = agent.session.fork(until_index=selected)
         await new_session.wait_for_flush()

klaude_code/tui/command/model_picker.py

@@ -72,9 +72,7 @@ def select_model_interactive(
         return ModelSelectResult(status=ModelSelectStatus.NON_TTY)
 
     # Interactive selection
-    from prompt_toolkit.styles import Style
-
-    from klaude_code.tui.terminal.selector import build_model_select_items, select_one
+    from klaude_code.tui.terminal.selector import DEFAULT_PICKER_STYLE, build_model_select_items, select_one
 
     names = [m.selector for m in result.filtered_models]
 
@@ -100,20 +98,7 @@ def select_model_interactive(
             pointer="→",
             use_search_filter=True,
             initial_value=initial_value,
-            style=Style(
-                [
-                    ("pointer", "ansigreen"),
-                    ("highlighted", "ansigreen"),
-                    ("msg", ""),
-                    ("meta", "fg:ansibrightblack"),
-                    ("text", "ansibrightblack"),
-                    ("question", "bold"),
-                    ("search_prefix", "ansibrightblack"),
-                    # search filter colors at the bottom
-                    ("search_success", "noinherit fg:ansigreen"),
-                    ("search_none", "noinherit fg:ansired"),
-                ]
-            ),
+            style=DEFAULT_PICKER_STYLE,
         )
         if isinstance(selected, str) and selected in names:
             return ModelSelectResult(status=ModelSelectStatus.SELECTED, model=selected)

klaude_code/tui/command/resume_cmd.py

@@ -1,28 +1,12 @@
 import asyncio
 
-from prompt_toolkit.styles import Style
-
 from klaude_code.log import log
 from klaude_code.protocol import commands, events, message, op
 from klaude_code.session.selector import build_session_select_options, format_user_messages_display
-from klaude_code.tui.terminal.selector import SelectItem, select_one
+from klaude_code.tui.terminal.selector import DEFAULT_PICKER_STYLE, SelectItem, select_one
 
 from .command_abc import Agent, CommandABC, CommandResult
 
-SESSION_SELECT_STYLE = Style(
-    [
-        ("msg", "fg:ansibrightblack"),
-        ("meta", ""),
-        ("pointer", "bold fg:ansigreen"),
-        ("highlighted", "fg:ansigreen"),
-        ("search_prefix", "fg:ansibrightblack"),
-        ("search_success", "noinherit fg:ansigreen"),
-        ("search_none", "noinherit fg:ansired"),
-        ("question", "bold"),
-        ("text", ""),
-    ]
-)
-
 
 def select_session_sync() -> str | None:
     """Interactive session selection (sync version for asyncio.to_thread)."""
@@ -62,7 +46,7 @@ def select_session_sync() -> str | None:
             message="Select a session to resume:",
             items=items,
             pointer="→",
-            style=SESSION_SELECT_STYLE,
+            style=DEFAULT_PICKER_STYLE,
         )
     except KeyboardInterrupt:
         return None

klaude_code/tui/command/sub_agent_model_cmd.py

@@ -4,27 +4,13 @@ from __future__ import annotations
 
 import asyncio
 
-from prompt_toolkit.styles import Style
-
 from klaude_code.config.config import load_config
 from klaude_code.config.sub_agent_model_helper import SubAgentModelHelper, SubAgentModelInfo
 from klaude_code.protocol import commands, events, message, op
-from klaude_code.tui.terminal.selector import SelectItem, build_model_select_items, select_one
+from klaude_code.tui.terminal.selector import DEFAULT_PICKER_STYLE, SelectItem, build_model_select_items, select_one
 
 from .command_abc import Agent, CommandABC, CommandResult
 
-SELECT_STYLE = Style(
-    [
-        ("instruction", "ansibrightblack"),
-        ("pointer", "ansigreen"),
-        ("highlighted", "ansigreen"),
-        ("text", "ansibrightblack"),
-        ("question", "bold"),
-        ("meta", "fg:ansibrightblack"),
-        ("msg", ""),
-    ]
-)
-
 USE_DEFAULT_BEHAVIOR = "__default__"
 
 
@@ -69,8 +55,8 @@ def _select_sub_agent_sync(
         result = select_one(
             message="Select sub-agent to configure:",
             items=items,
-            pointer="->",
-            style=SELECT_STYLE,
+            pointer="→",
+            style=DEFAULT_PICKER_STYLE,
             use_search_filter=False,
         )
         return result if isinstance(result, str) else None
@@ -103,8 +89,8 @@ def _select_model_for_sub_agent_sync(
         result = select_one(
             message=f"Select model for {sub_agent_type}:",
             items=all_items,
-            pointer="->",
-            style=SELECT_STYLE,
+            pointer="→",
+            style=DEFAULT_PICKER_STYLE,
             use_search_filter=True,
         )
         return result if isinstance(result, str) else None

klaude_code/tui/command/thinking_cmd.py

@@ -1,23 +1,11 @@
 import asyncio
 
-from prompt_toolkit.styles import Style
-
 from klaude_code.config.thinking import get_thinking_picker_data, parse_thinking_value
 from klaude_code.protocol import commands, events, llm_param, message, op
-from klaude_code.tui.terminal.selector import SelectItem, select_one
+from klaude_code.tui.terminal.selector import DEFAULT_PICKER_STYLE, SelectItem, select_one
 
 from .command_abc import Agent, CommandABC, CommandResult
 
-SELECT_STYLE = Style(
-    [
-        ("instruction", "ansibrightblack"),
-        ("pointer", "ansigreen"),
-        ("highlighted", "ansigreen"),
-        ("text", "ansibrightblack"),
-        ("question", "bold"),
-    ]
-)
-
 
 def _select_thinking_sync(config: llm_param.LLMConfigParameter) -> llm_param.Thinking | None:
     """Select thinking level (sync version)."""
@@ -35,7 +23,7 @@ def _select_thinking_sync(config: llm_param.LLMConfigParameter) -> llm_param.Thi
             message=data.message,
             items=items,
             pointer="→",
-            style=SELECT_STYLE,
+            style=DEFAULT_PICKER_STYLE,
             use_search_filter=False,
         )
         if result is None:

klaude_code/tui/commands.py

@@ -13,11 +13,6 @@ class RenderCommand:
     pass
 
 
-@dataclass(frozen=True, slots=True)
-class RenderReplayHistory(RenderCommand):
-    event: events.ReplayHistoryEvent
-
-
 @dataclass(frozen=True, slots=True)
 class RenderWelcome(RenderCommand):
     event: events.WelcomeEvent

klaude_code/tui/components/common.py

@@ -40,7 +40,7 @@ def truncate_middle(
         remaining = max(0, len(truncated_lines))
         return Text(f"  … (more {remaining} lines)", style=ThemeKey.TOOL_RESULT_TRUNCATED)
 
-    lines = text.split("\n")
+    lines = [line for line in text.split("\n") if line.strip()]
     truncated_lines = 0
     head_lines: list[str] = []
     tail_lines: list[str] = []

klaude_code/tui/components/metadata.py

@@ -1,5 +1,4 @@
 from rich.console import Group, RenderableType
-from rich.padding import Padding
 from rich.text import Text
 
 from klaude_code.const import DEFAULT_MAX_TOKENS
@@ -12,14 +11,14 @@ from klaude_code.ui.common import format_number
 def _render_task_metadata_block(
     metadata: model.TaskMetadata,
     *,
-    is_sub_agent: bool = False,
+    mark: Text,
     show_context_and_time: bool = True,
 ) -> RenderableType:
     """Render a single TaskMetadata block.
 
     Args:
         metadata: The TaskMetadata to render.
-        is_sub_agent: Whether this is a sub-agent block.
+        mark: The mark to display in the first column.
         show_context_and_time: Whether to show context usage percent and time.
 
     Returns:
@@ -31,19 +30,15 @@ def _render_task_metadata_block(
     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 description / tokens / cost / …
+    # Second column: provider/model description / tokens / cost / …
     content = Text()
-    content.append_text(Text(metadata.model_name, style=ThemeKey.METADATA_BOLD))
     if metadata.provider is not None:
-        content.append_text(Text("@", style=ThemeKey.METADATA)).append_text(
-            Text(metadata.provider.lower().replace(" ", "-"), style=ThemeKey.METADATA)
-        )
+        content.append_text(Text(metadata.provider.lower().replace(" ", "-"), style=ThemeKey.METADATA))
+        content.append_text(Text("/", style=ThemeKey.METADATA_DIM))
+    content.append_text(Text(metadata.model_name, style=ThemeKey.METADATA_BOLD))
     if metadata.description:
         content.append_text(Text(" ", style=ThemeKey.METADATA)).append_text(
-            Text(metadata.description, style=ThemeKey.METADATA_DIM)
+            Text(metadata.description, style=ThemeKey.METADATA_ITALIC)
         )
 
     # All info parts (tokens, cost, context, etc.)
@@ -63,7 +58,7 @@ def _render_task_metadata_block(
             token_text.append(" ∿", style=ThemeKey.METADATA_DIM)
             token_text.append(format_number(metadata.usage.reasoning_tokens), style=ThemeKey.METADATA)
         if metadata.usage.image_tokens > 0:
-            token_text.append(" ⌗ ", style=ThemeKey.METADATA_DIM)
+            token_text.append(" ⊡", style=ThemeKey.METADATA_DIM)
             token_text.append(format_number(metadata.usage.image_tokens), style=ThemeKey.METADATA)
         parts.append(token_text)
 
@@ -134,7 +129,7 @@ def _render_task_metadata_block(
         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))
+    return grid
 
 
 def render_task_metadata(e: events.TaskMetadataEvent) -> RenderableType:
@@ -144,16 +139,20 @@ def render_task_metadata(e: events.TaskMetadataEvent) -> RenderableType:
     if e.cancelled:
         renderables.append(Text())
 
-    renderables.append(
-        _render_task_metadata_block(e.metadata.main_agent, is_sub_agent=False, show_context_and_time=True)
-    )
+    has_sub_agents = len(e.metadata.sub_agent_task_metadata) > 0
+    # Use an extra space for the main agent mark to align with two-character marks (├─, └─)
+    main_mark_text = "✓"
+    main_mark = Text(main_mark_text, style=ThemeKey.METADATA)
+
+    renderables.append(_render_task_metadata_block(e.metadata.main_agent, mark=main_mark, 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=True))
+        sub_mark = Text("  └", style=ThemeKey.METADATA_DIM)
+        renderables.append(_render_task_metadata_block(meta, mark=sub_mark, show_context_and_time=True))
 
     # Add total cost line when there are sub-agents
-    if e.metadata.sub_agent_task_metadata:
+    if has_sub_agents:
         total_cost = 0.0
         currency = "USD"
         # Sum up costs from main agent and all sub-agents
@@ -166,12 +165,13 @@ def render_task_metadata(e: events.TaskMetadataEvent) -> RenderableType:
 
         currency_symbol = "¥" if currency == "CNY" else "$"
         total_line = Text.assemble(
-            ("Σ ", ThemeKey.METADATA_DIM),
+            ("  └", ThemeKey.METADATA_DIM),
+            (" Σ ", ThemeKey.METADATA_DIM),
             ("total ", ThemeKey.METADATA_DIM),
             (currency_symbol, ThemeKey.METADATA_DIM),
             (f"{total_cost:.4f}", ThemeKey.METADATA_DIM),
         )
 
-        renderables.append(Padding(total_line, (0, 0, 0, 2)))
+        renderables.append(total_line)
 
     return Group(*renderables)