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

klaude_code/llm/responses/client.py

@@ -1,6 +1,6 @@
 import json
 from collections.abc import AsyncGenerator
-from typing import TYPE_CHECKING, Literal, override
+from typing import TYPE_CHECKING, override
 
 import httpx
 import openai
@@ -11,9 +11,14 @@ from openai.types.responses.response_create_params import ResponseCreateParamsSt
 from klaude_code.const import LLM_HTTP_TIMEOUT_CONNECT, LLM_HTTP_TIMEOUT_READ, LLM_HTTP_TIMEOUT_TOTAL
 from klaude_code.llm.client import LLMClientABC, LLMStreamABC
 from klaude_code.llm.input_common import apply_config_defaults
-from klaude_code.llm.partial_message import degrade_thinking_to_text
 from klaude_code.llm.registry import register
 from klaude_code.llm.responses.input import convert_history_to_input, convert_tool_schema
+from klaude_code.llm.stream_parts import (
+    append_text_part,
+    append_thinking_text_part,
+    build_partial_message,
+    build_partial_parts,
+)
 from klaude_code.llm.usage import MetadataTracker, error_llm_stream
 from klaude_code.log import DebugType, log_debug
 from klaude_code.protocol import llm_param, message, model
@@ -58,68 +63,82 @@ def build_payload(param: llm_param.LLMCallParameter) -> ResponseCreateParamsStre
 
 
 class ResponsesStreamStateManager:
-    """Manages streaming state for Responses API and provides partial message access."""
+    """Manages streaming state for Responses API and provides partial message access.
+
+    Accumulates parts directly during streaming to support get_partial_message()
+    for cancellation scenarios. Merges consecutive text parts of the same type.
+    Each reasoning summary is kept as a separate ThinkingTextPart.
+    """
 
     def __init__(self, model_id: str) -> None:
         self.model_id = model_id
         self.response_id: str | None = None
-        self.stage: Literal["waiting", "thinking", "assistant", "tool"] = "waiting"
-        self.accumulated_thinking: list[str] = []
-        self.accumulated_text: list[str] = []
-        self.pending_signature: str | None = None
         self.assistant_parts: list[message.Part] = []
         self.stop_reason: model.StopReason | None = None
-
-    def flush_thinking(self) -> None:
-        """Flush accumulated thinking content into parts."""
-        if self.accumulated_thinking:
-            self.assistant_parts.append(
-                message.ThinkingTextPart(
-                    text="".join(self.accumulated_thinking),
-                    model_id=self.model_id,
-                )
+        self._new_thinking_part: bool = True  # Start fresh for first thinking part
+        self._summary_count: int = 0  # Track number of summary parts seen
+
+    def start_new_thinking_part(self) -> bool:
+        """Mark that the next thinking text should create a new ThinkingTextPart.
+
+        Returns True if this is not the first summary part (needs separator).
+        """
+        self._new_thinking_part = True
+        needs_separator = self._summary_count > 0
+        self._summary_count += 1
+        return needs_separator
+
+    def append_thinking_text(self, text: str) -> None:
+        """Append thinking text, merging with previous ThinkingTextPart if in same summary."""
+        if (
+            append_thinking_text_part(
+                self.assistant_parts,
+                text,
+                model_id=self.model_id,
+                force_new=self._new_thinking_part,
             )
-            self.accumulated_thinking.clear()
-        if self.pending_signature:
-            self.assistant_parts.append(
-                message.ThinkingSignaturePart(
-                    signature=self.pending_signature,
-                    model_id=self.model_id,
-                    format="openai_reasoning",
-                )
+            is not None
+        ):
+            self._new_thinking_part = False
+
+    def append_text(self, text: str) -> None:
+        """Append text, merging with previous TextPart if possible."""
+        append_text_part(self.assistant_parts, text)
+
+    def append_thinking_signature(self, signature: str) -> None:
+        """Append a ThinkingSignaturePart after the current part."""
+        self.assistant_parts.append(
+            message.ThinkingSignaturePart(
+                signature=signature,
+                model_id=self.model_id,
+                format="openai-responses",
             )
-            self.pending_signature = None
+        )
 
-    def flush_text(self) -> None:
-        """Flush accumulated text content into parts."""
-        if not self.accumulated_text:
-            return
-        self.assistant_parts.append(message.TextPart(text="".join(self.accumulated_text)))
-        self.accumulated_text.clear()
+    def append_tool_call(self, call_id: str, item_id: str | None, name: str, arguments_json: str) -> None:
+        """Append a ToolCallPart."""
+        self.assistant_parts.append(
+            message.ToolCallPart(
+                call_id=call_id,
+                id=item_id,
+                tool_name=name,
+                arguments_json=arguments_json,
+            )
+        )
 
-    def flush_all(self) -> list[message.Part]:
-        """Flush all accumulated content and return parts."""
-        self.flush_thinking()
-        self.flush_text()
-        return list(self.assistant_parts)
+    def get_partial_parts(self) -> list[message.Part]:
+        """Get accumulated parts excluding tool calls, with thinking degraded.
+
+        Filters out ToolCallPart and applies degrade_thinking_to_text.
+        """
+        return build_partial_parts(self.assistant_parts)
 
     def get_partial_message(self) -> message.AssistantMessage | None:
-        """Build a partial AssistantMessage from accumulated state."""
-        parts = self.flush_all()
-        filtered_parts: list[message.Part] = []
-        for part in parts:
-            if isinstance(part, message.ToolCallPart):
-                continue
-            filtered_parts.append(part)
-
-        filtered_parts = degrade_thinking_to_text(filtered_parts)
-        if not filtered_parts:
-            return None
-        return message.AssistantMessage(
-            parts=filtered_parts,
-            response_id=self.response_id,
-            stop_reason="aborted",
-        )
+        """Build a partial AssistantMessage from accumulated state.
+
+        Returns None if no content has been accumulated yet.
+        """
+        return build_partial_message(self.assistant_parts, response_id=self.response_id)
 
 
 async def parse_responses_stream(
@@ -157,24 +176,28 @@ async def parse_responses_stream(
             match event:
                 case responses.ResponseCreatedEvent() as event:
                     state.response_id = event.response.id
+                case responses.ResponseReasoningSummaryPartAddedEvent():
+                    # New reasoning summary part started, ensure it becomes a new ThinkingTextPart
+                    needs_separator = state.start_new_thinking_part()
+                    if needs_separator:
+                        # Add blank lines between summary parts for visual separation
+                        yield message.ThinkingTextDelta(content="  \n  \n", response_id=state.response_id)
                 case responses.ResponseReasoningSummaryTextDeltaEvent() as event:
                     if event.delta:
                         metadata_tracker.record_token()
-                        if state.stage == "assistant":
-                            state.flush_text()
-                        state.stage = "thinking"
-                        state.accumulated_thinking.append(event.delta)
+                        state.append_thinking_text(event.delta)
                         yield message.ThinkingTextDelta(content=event.delta, response_id=state.response_id)
                 case responses.ResponseReasoningSummaryTextDoneEvent() as event:
-                    if event.text and not state.accumulated_thinking:
-                        state.accumulated_thinking.append(event.text)
+                    # Fallback: if no delta was received but done has full text, use it
+                    if event.text:
+                        # Check if we already have content for this summary by seeing if last part matches
+                        last_part = state.assistant_parts[-1] if state.assistant_parts else None
+                        if not isinstance(last_part, message.ThinkingTextPart) or not last_part.text:
+                            state.append_thinking_text(event.text)
                 case responses.ResponseTextDeltaEvent() as event:
                     if event.delta:
                         metadata_tracker.record_token()
-                        if state.stage == "thinking":
-                            state.flush_thinking()
-                        state.stage = "assistant"
-                        state.accumulated_text.append(event.delta)
+                        state.append_text(event.delta)
                         yield message.AssistantTextDelta(content=event.delta, response_id=state.response_id)
                 case responses.ResponseOutputItemAddedEvent() as event:
                     if isinstance(event.item, responses.ResponseFunctionToolCall):
@@ -188,30 +211,23 @@ async def parse_responses_stream(
                     match event.item:
                         case responses.ResponseReasoningItem() as item:
                             if item.encrypted_content:
-                                state.pending_signature = item.encrypted_content
+                                state.append_thinking_signature(item.encrypted_content)
                         case responses.ResponseOutputMessage() as item:
-                            if not state.accumulated_text:
+                            # Fallback: if no text delta was received, extract from final message
+                            has_text = any(isinstance(p, message.TextPart) for p in state.assistant_parts)
+                            if not has_text:
                                 text_content = "\n".join(
-                                    [
-                                        part.text
-                                        for part in item.content
-                                        if isinstance(part, responses.ResponseOutputText)
-                                    ]
+                                    part.text for part in item.content if isinstance(part, responses.ResponseOutputText)
                                 )
                                 if text_content:
-                                    state.accumulated_text.append(text_content)
+                                    state.append_text(text_content)
                         case responses.ResponseFunctionToolCall() as item:
                             metadata_tracker.record_token()
-                            state.flush_thinking()
-                            state.flush_text()
-                            state.stage = "tool"
-                            state.assistant_parts.append(
-                                message.ToolCallPart(
-                                    call_id=item.call_id,
-                                    id=item.id,
-                                    tool_name=item.name,
-                                    arguments_json=item.arguments.strip(),
-                                )
+                            state.append_tool_call(
+                                call_id=item.call_id,
+                                item_id=item.id,
+                                name=item.name,
+                                arguments_json=item.arguments.strip(),
                             )
                         case _:
                             pass
@@ -254,10 +270,12 @@ async def parse_responses_stream(
                     )
     except (openai.OpenAIError, httpx.HTTPError) as e:
         yield message.StreamErrorItem(error=f"{e.__class__.__name__} {e!s}")
+        state.stop_reason = "error"
 
-    parts = state.flush_all()
     metadata_tracker.set_response_id(state.response_id)
     metadata = metadata_tracker.finalize()
+    # On error, use partial parts (excluding incomplete tool calls) for potential prefill on retry
+    parts = state.get_partial_parts() if state.stop_reason == "error" else list(state.assistant_parts)
     yield message.AssistantMessage(
         parts=parts,
         response_id=state.response_id,

klaude_code/llm/responses/input.py

@@ -80,8 +80,6 @@ def convert_history_to_input(
     """Convert a list of messages to response input params."""
     items: list[responses.ResponseInputItemParam] = []
 
-    degraded_thinking_texts: list[str] = []
-
     for msg, attachment in attach_developer_messages(history):
         match msg:
             case message.SystemMessage():
@@ -116,12 +114,19 @@ def convert_history_to_input(
             case message.ToolResultMessage():
                 items.append(_build_tool_result_item(msg, attachment))
             case message.AssistantMessage():
-                assistant_text_parts: list[responses.ResponseInputContentParam] = []
+                assistant_text_parts: list[responses.ResponseOutputTextParam] = []
                 pending_thinking_text: str | None = None
                 pending_signature: str | None = None
                 native_thinking_parts, degraded_for_message = split_thinking_parts(msg, model_name)
                 native_thinking_ids = {id(part) for part in native_thinking_parts}
-                degraded_thinking_texts.extend(degraded_for_message)
+                if degraded_for_message:
+                    degraded_text = "<thinking>\n" + "\n".join(degraded_for_message) + "\n</thinking>"
+                    assistant_text_parts.append(
+                        cast(
+                            responses.ResponseOutputTextParam,
+                            {"type": "output_text", "text": degraded_text},
+                        )
+                    )
 
                 def flush_text() -> None:
                     nonlocal assistant_text_parts
@@ -164,8 +169,8 @@ def convert_history_to_input(
                     if isinstance(part, message.TextPart):
                         assistant_text_parts.append(
                             cast(
-                                responses.ResponseInputContentParam,
-                                {"type": "input_text", "text": part.text},
+                                responses.ResponseOutputTextParam,
+                                {"type": "output_text", "text": part.text},
                             )
                         )
                     elif isinstance(part, message.ToolCallPart):
@@ -188,25 +193,6 @@ def convert_history_to_input(
             case _:
                 continue
 
-    if degraded_thinking_texts:
-        degraded_item = cast(
-            responses.ResponseInputItemParam,
-            {
-                "type": "message",
-                "role": "assistant",
-                "content": [
-                    cast(
-                        responses.ResponseInputContentParam,
-                        {
-                            "type": "input_text",
-                            "text": "<thinking>\n" + "\n".join(degraded_thinking_texts) + "\n</thinking>",
-                        },
-                    )
-                ],
-            },
-        )
-        items.insert(0, degraded_item)
-
     return items
 
 

klaude_code/llm/stream_parts.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from collections.abc import MutableSequence
+
+from klaude_code.protocol import message
+
+
+def append_text_part(parts: MutableSequence[message.Part], text: str) -> int | None:
+    if not text:
+        return None
+
+    if parts:
+        last = parts[-1]
+        if isinstance(last, message.TextPart):
+            parts[-1] = message.TextPart(text=last.text + text)
+            return len(parts) - 1
+
+    parts.append(message.TextPart(text=text))
+    return len(parts) - 1
+
+
+def append_thinking_text_part(
+    parts: MutableSequence[message.Part],
+    text: str,
+    *,
+    model_id: str,
+    force_new: bool = False,
+) -> int | None:
+    if not text:
+        return None
+
+    if not force_new and parts:
+        last = parts[-1]
+        if isinstance(last, message.ThinkingTextPart):
+            parts[-1] = message.ThinkingTextPart(
+                text=last.text + text,
+                model_id=model_id,
+            )
+            return len(parts) - 1
+
+    parts.append(message.ThinkingTextPart(text=text, model_id=model_id))
+    return len(parts) - 1
+
+
+def degrade_thinking_to_text(parts: list[message.Part]) -> list[message.Part]:
+    """Degrade thinking parts into a regular TextPart.
+
+    Some providers require thinking signatures/encrypted content to be echoed back
+    for subsequent calls. During interruption we cannot reliably determine whether
+    we have a complete signature, so we persist thinking as plain text instead.
+    """
+
+    thinking_texts: list[str] = []
+    non_thinking_parts: list[message.Part] = []
+
+    for part in parts:
+        if isinstance(part, message.ThinkingTextPart):
+            text = part.text
+            if text and text.strip():
+                thinking_texts.append(text)
+            continue
+        if isinstance(part, message.ThinkingSignaturePart):
+            continue
+        non_thinking_parts.append(part)
+
+    if not thinking_texts:
+        return non_thinking_parts
+
+    joined = "\n".join(thinking_texts).strip()
+    thinking_block = f"<thinking>\n{joined}\n</thinking>"
+    if non_thinking_parts:
+        thinking_block += "\n\n"
+
+    return [message.TextPart(text=thinking_block), *non_thinking_parts]
+
+
+def build_partial_parts(parts: list[message.Part]) -> list[message.Part]:
+    filtered_parts: list[message.Part] = [p for p in parts if not isinstance(p, message.ToolCallPart)]
+    return degrade_thinking_to_text(filtered_parts)
+
+
+def build_partial_message(
+    parts: list[message.Part],
+    *,
+    response_id: str | None,
+) -> message.AssistantMessage | None:
+    partial_parts = build_partial_parts(parts)
+    if not partial_parts:
+        return None
+    return message.AssistantMessage(
+        parts=partial_parts,
+        response_id=response_id,
+        stop_reason="aborted",
+    )

klaude_code/log.py

@@ -1,13 +1,16 @@
 import gzip
+import json
 import logging
 import os
 import shutil
 import subprocess
+from base64 import b64encode
 from collections.abc import Iterable
 from datetime import datetime, timedelta
 from enum import Enum
 from logging.handlers import RotatingFileHandler
 from pathlib import Path
+from typing import cast
 
 from rich.console import Console
 from rich.logging import RichHandler
@@ -316,3 +319,57 @@ def _trash_path(path: Path) -> None:
         )
     except FileNotFoundError:
         path.unlink(missing_ok=True)
+
+
+# Debug JSON serialization utilities
+_DEBUG_TRUNCATE_PREFIX_CHARS = 96
+
+# Keys whose values should be truncated (e.g., signatures, large payloads)
+_TRUNCATE_KEYS = {"thought_signature", "thoughtSignature"}
+
+
+def _truncate_debug_str(value: str, *, prefix_chars: int = _DEBUG_TRUNCATE_PREFIX_CHARS) -> str:
+    if len(value) <= prefix_chars:
+        return value
+    return f"{value[:prefix_chars]}...(truncated,len={len(value)})"
+
+
+def _sanitize_debug_value(value: object) -> object:
+    if isinstance(value, (bytes, bytearray)):
+        encoded = b64encode(bytes(value)).decode("ascii")
+        return _truncate_debug_str(encoded)
+    if isinstance(value, str):
+        return value
+    if isinstance(value, list):
+        return [_sanitize_debug_value(v) for v in cast(list[object], value)]
+    if isinstance(value, dict):
+        return _sanitize_debug_dict(value)  # type: ignore[arg-type]
+    return value
+
+
+def _sanitize_debug_dict(obj: dict[object, object]) -> dict[object, object]:
+    sanitized: dict[object, object] = {}
+    for k, v in obj.items():
+        if k in _TRUNCATE_KEYS:
+            if isinstance(v, str):
+                sanitized[k] = _truncate_debug_str(v)
+            else:
+                sanitized[k] = _sanitize_debug_value(v)
+            continue
+        sanitized[k] = _sanitize_debug_value(v)
+
+    # Truncate inline image payloads (data field with mime_type indicates image blob)
+    if "data" in sanitized and ("mime_type" in sanitized or "mimeType" in sanitized):
+        data = sanitized.get("data")
+        if isinstance(data, str):
+            sanitized["data"] = _truncate_debug_str(data)
+        elif isinstance(data, (bytes, bytearray)):
+            encoded = b64encode(bytes(data)).decode("ascii")
+            sanitized["data"] = _truncate_debug_str(encoded)
+
+    return sanitized
+
+
+def debug_json(value: object) -> str:
+    """Serialize a value to JSON for debug logging, truncating large payloads."""
+    return json.dumps(_sanitize_debug_value(value), ensure_ascii=False)

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