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

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/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/metadata.py

@@ -30,13 +30,12 @@ def _render_task_metadata_block(
     currency = metadata.usage.currency if metadata.usage else "USD"
     currency_symbol = "¥" if currency == "CNY" else "$"
 
-    # 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_ITALIC)

klaude_code/tui/components/sub_agent.py

@@ -135,6 +135,7 @@ def build_sub_agent_state_from_tool_call(e: events.ToolCallEvent) -> model.SubAg
     description = profile.name
     prompt = ""
     output_schema: dict[str, Any] | None = None
+    generation: dict[str, Any] | None = None
     resume: str | None = None
     if e.arguments:
         try:
@@ -155,10 +156,15 @@ def build_sub_agent_state_from_tool_call(e: events.ToolCallEvent) -> model.SubAg
             schema_value = payload.get(profile.output_schema_arg)
             if isinstance(schema_value, dict):
                 output_schema = cast(dict[str, Any], schema_value)
+        # Extract generation config for ImageGen
+        generation_value = payload.get("generation")
+        if isinstance(generation_value, dict):
+            generation = cast(dict[str, Any], generation_value)
     return model.SubAgentState(
         sub_agent_type=profile.name,
         sub_agent_desc=description,
         sub_agent_prompt=prompt,
         resume=resume,
         output_schema=output_schema,
+        generation=generation,
     )

klaude_code/tui/display.py

@@ -30,8 +30,18 @@ class TUIDisplay(DisplayABC):
 
     @override
     async def consume_event(self, event: events.Event) -> None:
+        if isinstance(event, events.ReplayHistoryEvent):
+            await self._renderer.execute(self._machine.begin_replay())
+            for item in event.events:
+                commands = self._machine.transition_replay(item)
+                if commands:
+                    await self._renderer.execute(commands)
+            await self._renderer.execute(self._machine.end_replay())
+            return
+
         commands = self._machine.transition(event)
-        await self._renderer.execute(commands)
+        if commands:
+            await self._renderer.execute(commands)
 
     @override
     async def start(self) -> None:

klaude_code/tui/input/completers.py

@@ -313,11 +313,13 @@ class _AtFilesCompleter(Completer):
             if not suggestions:
                 return []  # type: ignore[reportUnknownVariableType]
             start_position = token_start_in_input - len(text_before)
-            for s in suggestions[: self._max_results]:
+            suggestions_to_show = suggestions[: self._max_results]
+            align_width = self._display_align_width(suggestions_to_show)
+            for s in suggestions_to_show:
                 yield Completion(
                     text=self._format_completion_text(s, is_quoted=is_quoted),
                     start_position=start_position,
-                    display=self._format_display_label(s, 0),
+                    display=self._format_display_label(s, align_width),
                     display_meta=s,
                 )
             return []  # type: ignore[reportUnknownVariableType]
@@ -329,12 +331,14 @@ class _AtFilesCompleter(Completer):
 
         # Prepare Completion objects. Replace from the '@' character.
         start_position = token_start_in_input - len(text_before)  # negative
-        for s in suggestions[: self._max_results]:
+        suggestions_to_show = suggestions[: self._max_results]
+        align_width = self._display_align_width(suggestions_to_show)
+        for s in suggestions_to_show:
             # Insert formatted text (with quoting when needed) so that subsequent typing does not keep triggering
             yield Completion(
                 text=self._format_completion_text(s, is_quoted=is_quoted),
                 start_position=start_position,
-                display=self._format_display_label(s, 0),
+                display=self._format_display_label(s, align_width),
                 display_meta=s,
             )
 
@@ -543,9 +547,9 @@ class _AtFilesCompleter(Completer):
         Keep this unstyled so that the completion menu's selection style can
         fully override the selected row.
         """
-
-        _ = align_width
-        return self._display_name(suggestion)
+        name = self._display_name(suggestion)
+        # Pad to align_width + extra padding for visual separation from meta
+        return name.ljust(align_width + 6)
 
     def _display_align_width(self, suggestions: list[str]) -> int:
         """Calculate alignment width for display labels."""

klaude_code/tui/machine.py

@@ -19,6 +19,7 @@ from klaude_code.tui.commands import (
     EmitTmuxSignal,
     EndAssistantStream,
     EndThinkingStream,
+    PrintBlankLine,
     PrintRuleLine,
     RenderAssistantImage,
     RenderCommand,
@@ -26,7 +27,6 @@ from klaude_code.tui.commands import (
     RenderDeveloperMessage,
     RenderError,
     RenderInterrupt,
-    RenderReplayHistory,
     RenderTaskFinish,
     RenderTaskMetadata,
     RenderTaskStart,
@@ -382,17 +382,25 @@ class DisplayStateMachine:
         self._spinner.set_toast_status(None)
         return self._spinner_update_commands()
 
+    def begin_replay(self) -> list[RenderCommand]:
+        self._spinner.reset()
+        return [SpinnerStop(), PrintBlankLine()]
+
+    def end_replay(self) -> list[RenderCommand]:
+        return [SpinnerStop()]
+
+    def transition_replay(self, event: events.Event) -> list[RenderCommand]:
+        return self._transition(event, is_replay=True)
+
     def transition(self, event: events.Event) -> list[RenderCommand]:
+        return self._transition(event, is_replay=False)
+
+    def _transition(self, event: events.Event, *, is_replay: bool) -> list[RenderCommand]:
         session_id = getattr(event, "session_id", "__app__")
         s = self._session(session_id)
         cmds: list[RenderCommand] = []
 
         match event:
-            case events.ReplayHistoryEvent() as e:
-                cmds.append(RenderReplayHistory(e))
-                cmds.append(SpinnerStop())
-                return cmds
-
             case events.WelcomeEvent() as e:
                 cmds.append(RenderWelcome(e))
                 return cmds
@@ -408,13 +416,16 @@ class DisplayStateMachine:
                 s.model_id = e.model_id
                 if not s.is_sub_agent:
                     self._set_primary_if_needed(e.session_id)
-                    cmds.append(TaskClockStart())
+                    if not is_replay:
+                        cmds.append(TaskClockStart())
                 else:
                     s.sub_agent_thinking_header = SubAgentThinkingHeaderState()
 
-                cmds.append(SpinnerStart())
+                if not is_replay:
+                    cmds.append(SpinnerStart())
                 cmds.append(RenderTaskStart(e))
-                cmds.extend(self._spinner_update_commands())
+                if not is_replay:
+                    cmds.extend(self._spinner_update_commands())
                 return cmds
 
             case events.DeveloperMessageEvent() as e:
@@ -427,9 +438,10 @@ class DisplayStateMachine:
 
             case events.TurnStartEvent() as e:
                 cmds.append(RenderTurnStart(e))
-                self._spinner.clear_for_new_turn()
-                self._spinner.set_reasoning_status(None)
-                cmds.extend(self._spinner_update_commands())
+                if not is_replay:
+                    self._spinner.clear_for_new_turn()
+                    self._spinner.set_reasoning_status(None)
+                    cmds.extend(self._spinner_update_commands())
                 return cmds
 
             case events.ThinkingStartEvent() as e:
@@ -441,9 +453,11 @@ class DisplayStateMachine:
                 s.thinking_tail = ""
                 # Ensure the status reflects that reasoning has started even
                 # before we receive any deltas (or a bold header).
-                self._spinner.set_reasoning_status(STATUS_THINKING_TEXT)
+                if not is_replay:
+                    self._spinner.set_reasoning_status(STATUS_THINKING_TEXT)
                 cmds.append(StartThinkingStream(session_id=e.session_id))
-                cmds.extend(self._spinner_update_commands())
+                if not is_replay:
+                    cmds.extend(self._spinner_update_commands())
                 return cmds
 
             case events.ThinkingDeltaEvent() as e:
@@ -463,7 +477,7 @@ class DisplayStateMachine:
 
                 # Update reasoning status for spinner (based on bounded tail).
                 # Only extract headers for models that use markdown bold headers in thinking.
-                if s.should_extract_reasoning_header:
+                if not is_replay and s.should_extract_reasoning_header:
                     s.thinking_tail = (s.thinking_tail + e.content)[-8192:]
                     header = extract_last_bold_header(normalize_thinking_content(s.thinking_tail))
                     if header:
@@ -478,26 +492,31 @@ class DisplayStateMachine:
                 if not self._is_primary(e.session_id):
                     return []
                 s.thinking_stream_active = False
-                self._spinner.clear_default_reasoning_status()
+                if not is_replay:
+                    self._spinner.clear_default_reasoning_status()
                 cmds.append(EndThinkingStream(session_id=e.session_id))
-                cmds.append(SpinnerStart())
-                cmds.extend(self._spinner_update_commands())
+                if not is_replay:
+                    cmds.append(SpinnerStart())
+                    cmds.extend(self._spinner_update_commands())
                 return cmds
 
             case events.AssistantTextStartEvent() as e:
                 if s.is_sub_agent:
-                    self._spinner.set_composing(True)
-                    cmds.extend(self._spinner_update_commands())
+                    if not is_replay:
+                        self._spinner.set_composing(True)
+                        cmds.extend(self._spinner_update_commands())
                     return cmds
                 if not self._is_primary(e.session_id):
                     return []
 
                 s.assistant_stream_active = True
                 s.assistant_char_count = 0
-                self._spinner.set_composing(True)
-                self._spinner.clear_tool_calls()
+                if not is_replay:
+                    self._spinner.set_composing(True)
+                    self._spinner.clear_tool_calls()
                 cmds.append(StartAssistantStream(session_id=e.session_id))
-                cmds.extend(self._spinner_update_commands())
+                if not is_replay:
+                    cmds.extend(self._spinner_update_commands())
                 return cmds
 
             case events.AssistantTextDeltaEvent() as e:
@@ -507,24 +526,29 @@ class DisplayStateMachine:
                     return []
 
                 s.assistant_char_count += len(e.content)
-                self._spinner.set_buffer_length(s.assistant_char_count)
+                if not is_replay:
+                    self._spinner.set_buffer_length(s.assistant_char_count)
                 cmds.append(AppendAssistant(session_id=e.session_id, content=e.content))
-                cmds.extend(self._spinner_update_commands())
+                if not is_replay:
+                    cmds.extend(self._spinner_update_commands())
                 return cmds
 
             case events.AssistantTextEndEvent() as e:
                 if s.is_sub_agent:
-                    self._spinner.set_composing(False)
-                    cmds.extend(self._spinner_update_commands())
+                    if not is_replay:
+                        self._spinner.set_composing(False)
+                        cmds.extend(self._spinner_update_commands())
                     return cmds
                 if not self._is_primary(e.session_id):
                     return []
 
                 s.assistant_stream_active = False
-                self._spinner.set_composing(False)
+                if not is_replay:
+                    self._spinner.set_composing(False)
                 cmds.append(EndAssistantStream(session_id=e.session_id))
-                cmds.append(SpinnerStart())
-                cmds.extend(self._spinner_update_commands())
+                if not is_replay:
+                    cmds.append(SpinnerStart())
+                    cmds.extend(self._spinner_update_commands())
                 return cmds
 
             case events.AssistantImageDeltaEvent() as e:
@@ -536,9 +560,10 @@ class DisplayStateMachine:
                     return []
                 if not self._is_primary(e.session_id):
                     return []
-                self._spinner.set_composing(False)
-                cmds.append(SpinnerStart())
-                cmds.extend(self._spinner_update_commands())
+                if not is_replay:
+                    self._spinner.set_composing(False)
+                    cmds.append(SpinnerStart())
+                    cmds.extend(self._spinner_update_commands())
                 return cmds
 
             case events.ToolCallStartEvent() as e:
@@ -553,18 +578,20 @@ class DisplayStateMachine:
                         primary.thinking_stream_active = False
                         cmds.append(EndThinkingStream(session_id=primary.session_id))
 
-                self._spinner.set_composing(False)
+                if not is_replay:
+                    self._spinner.set_composing(False)
 
                 # Skip activity state for fast tools on non-streaming models (e.g., Gemini)
                 # to avoid flash-and-disappear effect
-                if not s.should_skip_tool_activity(e.tool_name):
+                if not is_replay and not s.should_skip_tool_activity(e.tool_name):
                     tool_active_form = get_tool_active_form(e.tool_name)
                     if is_sub_agent_tool(e.tool_name):
                         self._spinner.add_sub_agent_tool_call(e.tool_call_id, tool_active_form)
                     else:
                         self._spinner.add_tool_call(tool_active_form)
 
-                cmds.extend(self._spinner_update_commands())
+                if not is_replay:
+                    cmds.extend(self._spinner_update_commands())
                 return cmds
 
             case events.ToolCallEvent() as e:
@@ -583,7 +610,7 @@ class DisplayStateMachine:
                 return cmds
 
             case events.ToolResultEvent() as e:
-                if is_sub_agent_tool(e.tool_name):
+                if not is_replay and is_sub_agent_tool(e.tool_name):
                     self._spinner.finish_sub_agent_tool_call(e.tool_call_id, get_tool_active_form(e.tool_name))
                     cmds.extend(self._spinner_update_commands())
 
@@ -601,9 +628,10 @@ class DisplayStateMachine:
 
             case events.TodoChangeEvent() as e:
                 todo_text = _extract_active_form_text(e)
-                self._spinner.set_todo_status(todo_text)
-                self._spinner.clear_for_new_turn()
-                cmds.extend(self._spinner_update_commands())
+                if not is_replay:
+                    self._spinner.set_todo_status(todo_text)
+                    self._spinner.clear_for_new_turn()
+                    cmds.extend(self._spinner_update_commands())
                 return cmds
 
             case events.UsageEvent() as e:
@@ -613,7 +641,7 @@ class DisplayStateMachine:
                 if not self._is_primary(e.session_id):
                     return []
                 context_percent = e.usage.context_usage_percent
-                if context_percent is not None:
+                if not is_replay and context_percent is not None:
                     self._spinner.set_context_percent(context_percent)
                     cmds.extend(self._spinner_update_commands())
                 return cmds
@@ -624,37 +652,43 @@ class DisplayStateMachine:
             case events.TaskFinishEvent() as e:
                 cmds.append(RenderTaskFinish(e))
                 if not s.is_sub_agent:
-                    cmds.append(TaskClockClear())
-                    self._spinner.reset()
-                    cmds.append(SpinnerStop())
-                    cmds.append(PrintRuleLine())
-                    cmds.append(EmitTmuxSignal())
+                    if not is_replay:
+                        cmds.append(TaskClockClear())
+                        self._spinner.reset()
+                        cmds.append(SpinnerStop())
+                        cmds.append(PrintRuleLine())
+                        cmds.append(EmitTmuxSignal())
                 else:
                     s.sub_agent_thinking_header = None
                 return cmds
 
             case events.InterruptEvent() as e:
-                self._spinner.reset()
-                cmds.append(SpinnerStop())
+                if not is_replay:
+                    self._spinner.reset()
+                    cmds.append(SpinnerStop())
                 cmds.append(EndThinkingStream(session_id=e.session_id))
                 cmds.append(EndAssistantStream(session_id=e.session_id))
-                cmds.append(TaskClockClear())
+                if not is_replay:
+                    cmds.append(TaskClockClear())
                 cmds.append(RenderInterrupt(session_id=e.session_id))
                 return cmds
 
             case events.ErrorEvent() as e:
-                cmds.append(EmitOsc94Error())
+                if not is_replay:
+                    cmds.append(EmitOsc94Error())
                 cmds.append(RenderError(e))
-                if not e.can_retry:
+                if not is_replay and not e.can_retry:
                     self._spinner.reset()
                     cmds.append(SpinnerStop())
-                cmds.extend(self._spinner_update_commands())
+                if not is_replay:
+                    cmds.extend(self._spinner_update_commands())
                 return cmds
 
             case events.EndEvent():
-                self._spinner.reset()
-                cmds.append(SpinnerStop())
-                cmds.append(TaskClockClear())
+                if not is_replay:
+                    self._spinner.reset()
+                    cmds.append(SpinnerStop())
+                    cmds.append(TaskClockClear())
                 return cmds
 
             case _:

klaude_code/tui/renderer.py

@@ -35,7 +35,6 @@ from klaude_code.tui.commands import (
     RenderDeveloperMessage,
     RenderError,
     RenderInterrupt,
-    RenderReplayHistory,
     RenderTaskFinish,
     RenderTaskMetadata,
     RenderTaskStart,
@@ -437,67 +436,10 @@ class TUICommandRenderer:
             Text.assemble(
                 (c_thinking.THINKING_MESSAGE_MARK, ThemeKey.THINKING),
                 " ",
-                (stripped, ThemeKey.THINKING_BOLD),
+                (stripped, ThemeKey.THINKING),
             )
         )
 
-    async def replay_history(self, history_events: events.ReplayHistoryEvent) -> None:
-        tool_call_dict: dict[str, events.ToolCallEvent] = {}
-        self.print()
-        for event in history_events.events:
-            event_session_id = getattr(event, "session_id", history_events.session_id)
-            is_sub_agent = self.is_sub_agent_session(event_session_id)
-
-            with self.session_print_context(event_session_id):
-                match event:
-                    case events.TaskStartEvent() as e:
-                        self.display_task_start(e)
-                    case events.TurnStartEvent():
-                        self.print()
-                    case events.AssistantImageDeltaEvent() as e:
-                        self.display_image(e.file_path)
-                    case events.ResponseCompleteEvent() as e:
-                        if is_sub_agent:
-                            if self._should_display_sub_agent_thinking_header(event_session_id) and e.thinking_text:
-                                header = c_thinking.extract_last_bold_header(
-                                    c_thinking.normalize_thinking_content(e.thinking_text)
-                                )
-                                if header:
-                                    self.display_thinking_header(header)
-                            continue
-                        if e.thinking_text:
-                            self.display_thinking(e.thinking_text)
-                        renderable = c_assistant.render_assistant_message(e.content, code_theme=self.themes.code_theme)
-                        if renderable is not None:
-                            self.print(renderable)
-                            self.print()
-                    case events.DeveloperMessageEvent() as e:
-                        self.display_developer_message(e)
-                    case events.UserMessageEvent() as e:
-                        if is_sub_agent:
-                            continue
-                        self.print(c_user_input.render_user_input(e.content))
-                    case events.ToolCallEvent() as e:
-                        tool_call_dict[e.tool_call_id] = e
-                    case events.ToolResultEvent() as e:
-                        tool_call_event = tool_call_dict.get(e.tool_call_id)
-                        if tool_call_event is not None:
-                            self.display_tool_call(tool_call_event)
-                        tool_call_dict.pop(e.tool_call_id, None)
-                        if is_sub_agent:
-                            continue
-                        self.display_tool_call_result(e)
-                    case events.TaskMetadataEvent() as e:
-                        self.print(c_metadata.render_task_metadata(e))
-                        self.print()
-                    case events.InterruptEvent():
-                        self.print()
-                        self.print(c_user_input.render_interrupt())
-                    case events.ErrorEvent() as e:
-                        self.display_error(e)
-                    case events.TaskFinishEvent() as e:
-                        self.display_task_finish(e)
-
     def display_developer_message(self, e: events.DeveloperMessageEvent) -> None:
         if not c_developer.need_render_developer_message(e):
             return
@@ -615,9 +557,6 @@ class TUICommandRenderer:
     async def execute(self, commands: list[RenderCommand]) -> None:
         for cmd in commands:
             match cmd:
-                case RenderReplayHistory(event=event):
-                    await self.replay_history(event)
-                    self.spinner_stop()
                 case RenderWelcome(event=event):
                     self.display_welcome(event)
                 case RenderUserMessage(event=event):