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

klaude_code/llm/anthropic/client.py

@@ -33,8 +33,13 @@ from klaude_code.const import (
 from klaude_code.llm.anthropic.input import convert_history_to_input, convert_system_to_input, convert_tool_schema
 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.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
@@ -64,11 +69,10 @@ class AnthropicStreamStateManager:
 
     def __init__(self, model_id: str) -> None:
         self.model_id = model_id
-        self.accumulated_thinking: list[str] = []
-        self.accumulated_content: list[str] = []
-        self.parts: list[message.Part] = []
+        self.assistant_parts: list[message.Part] = []
         self.response_id: str | None = None
-        self.pending_signature: str | None = None
+        self._pending_signature: str | None = None
+        self._pending_signature_thinking_index: int | None = None
         self.stop_reason: model.StopReason | None = None
 
         # Tool call state
@@ -80,34 +84,59 @@ class AnthropicStreamStateManager:
         self.input_token: int = 0
         self.cached_token: int = 0
 
-    def flush_thinking(self) -> None:
-        """Flush accumulated thinking content into parts."""
-        if not self.accumulated_thinking:
+    def append_thinking_text(self, text: str) -> None:
+        """Append thinking text, merging with the previous ThinkingTextPart when possible."""
+        index = append_thinking_text_part(self.assistant_parts, text, model_id=self.model_id)
+        if index is not None:
+            self._pending_signature_thinking_index = index
+
+    def append_text(self, text: str) -> None:
+        """Append assistant text, merging with the previous TextPart when possible."""
+        append_text_part(self.assistant_parts, text)
+
+    def set_pending_signature(self, signature: str) -> None:
+        if signature:
+            self._pending_signature = signature
+
+    def flush_pending_signature(self) -> None:
+        """Attach any pending signature to the most recent thinking segment.
+
+        Anthropic's signature is semantically tied to its thinking content. The
+        signature delta may arrive slightly after thinking text, so we insert the
+        signature part adjacent to the thinking part it signs.
+        """
+
+        if not self._pending_signature:
+            return
+        if self._pending_signature_thinking_index is None:
+            # No thinking part seen for this signature; drop it.
+            self._pending_signature = None
             return
-        full_thinking = "".join(self.accumulated_thinking)
-        self.parts.append(message.ThinkingTextPart(text=full_thinking, model_id=self.model_id))
-        if self.pending_signature:
-            self.parts.append(
-                message.ThinkingSignaturePart(
-                    signature=self.pending_signature,
-                    model_id=self.model_id,
-                    format="anthropic",
-                )
-            )
-        self.accumulated_thinking.clear()
-        self.pending_signature = None
 
-    def flush_content(self) -> None:
-        """Flush accumulated content into parts."""
-        if not self.accumulated_content:
+        insert_at = self._pending_signature_thinking_index + 1
+        # Avoid inserting duplicates if flush is called multiple times.
+        if insert_at < len(self.assistant_parts) and isinstance(
+            self.assistant_parts[insert_at], message.ThinkingSignaturePart
+        ):
+            self._pending_signature = None
             return
-        self.parts.append(message.TextPart(text="".join(self.accumulated_content)))
-        self.accumulated_content.clear()
+
+        self.assistant_parts.insert(
+            insert_at,
+            message.ThinkingSignaturePart(
+                signature=self._pending_signature,
+                model_id=self.model_id,
+                format="anthropic",
+            ),
+        )
+
+        self._pending_signature = None
+        self._pending_signature_thinking_index = None
 
     def flush_tool_call(self) -> None:
         """Flush current tool call into parts."""
         if self.current_tool_name and self.current_tool_call_id:
-            self.parts.append(
+            self.assistant_parts.append(
                 message.ToolCallPart(
                     call_id=self.current_tool_call_id,
                     tool_name=self.current_tool_name,
@@ -119,11 +148,17 @@ class AnthropicStreamStateManager:
         self.current_tool_inputs = None
 
     def flush_all(self) -> list[message.Part]:
-        """Flush all accumulated content in order and return parts."""
-        self.flush_thinking()
-        self.flush_content()
+        """Flush all pending content in order and return parts."""
+        self.flush_pending_signature()
         self.flush_tool_call()
-        return list(self.parts)
+        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.
@@ -131,16 +166,7 @@ class AnthropicStreamStateManager:
         Flushes all accumulated content and returns the message with
         stop_reason="aborted". Returns None if no content has been accumulated.
         """
-        self.flush_thinking()
-        self.flush_content()
-        parts = degrade_thinking_to_text(list(self.parts))
-        if not parts:
-            return None
-        return message.AssistantMessage(
-            parts=parts,
-            response_id=self.response_id,
-            stop_reason="aborted",
-        )
+        return build_partial_message(self.assistant_parts, response_id=self.response_id)
 
 
 def build_payload(
@@ -226,17 +252,18 @@ async def parse_anthropic_stream(
                     case BetaThinkingDelta() as delta:
                         if delta.thinking:
                             metadata_tracker.record_token()
-                            state.accumulated_thinking.append(delta.thinking)
+                            state.append_thinking_text(delta.thinking)
                             yield message.ThinkingTextDelta(
                                 content=delta.thinking,
                                 response_id=state.response_id,
                             )
                     case BetaSignatureDelta() as delta:
-                        state.pending_signature = delta.signature
+                        state.set_pending_signature(delta.signature)
                     case BetaTextDelta() as delta:
                         if delta.text:
                             metadata_tracker.record_token()
-                            state.accumulated_content.append(delta.text)
+                            state.flush_pending_signature()
+                            state.append_text(delta.text)
                             yield message.AssistantTextDelta(
                                 content=delta.text,
                                 response_id=state.response_id,
@@ -251,6 +278,7 @@ async def parse_anthropic_stream(
                 match event.content_block:
                     case BetaToolUseBlock() as block:
                         metadata_tracker.record_token()
+                        state.flush_pending_signature()
                         yield message.ToolCallStartDelta(
                             response_id=state.response_id,
                             call_id=block.id,
@@ -262,12 +290,7 @@ async def parse_anthropic_stream(
                     case _:
                         pass
             case BetaRawContentBlockStopEvent():
-                if state.accumulated_thinking:
-                    metadata_tracker.record_token()
-                    state.flush_thinking()
-                if state.accumulated_content:
-                    metadata_tracker.record_token()
-                    state.flush_content()
+                state.flush_pending_signature()
                 if state.current_tool_name and state.current_tool_call_id:
                     metadata_tracker.record_token()
                     state.flush_tool_call()
@@ -337,8 +360,10 @@ class AnthropicLLMStream(LLMStreamABC):
             self._metadata_tracker.set_model_name(str(self._param.model_id))
             self._metadata_tracker.set_response_id(self._state.response_id)
             metadata = self._metadata_tracker.finalize()
+            # Use accumulated parts for potential prefill on retry
+            self._state.flush_all()
             yield message.AssistantMessage(
-                parts=[],
+                parts=self._state.get_partial_parts(),
                 response_id=self._state.response_id,
                 usage=metadata,
                 stop_reason="error",

klaude_code/llm/anthropic/input.py

@@ -109,19 +109,34 @@ def _tool_blocks_to_message(blocks: list[BetaToolResultBlockParam]) -> BetaMessa
 def _assistant_message_to_message(msg: message.AssistantMessage, model_name: str | None) -> BetaMessageParam:
     content: list[BetaContentBlockParam] = []
     current_thinking_content: str | None = None
-    native_thinking_parts, degraded_thinking_texts = split_thinking_parts(msg, model_name)
+    native_thinking_parts, _ = split_thinking_parts(msg, model_name)
     native_thinking_ids = {id(part) for part in native_thinking_parts}
 
-    def _flush_thinking() -> None:
+    def _degraded_thinking_block(text: str) -> BetaTextBlockParam | None:
+        stripped = text.strip()
+        if not stripped:
+            return None
+        return cast(
+            BetaTextBlockParam,
+            {
+                "type": "text",
+                "text": f"<thinking>\n{stripped}\n</thinking>",
+            },
+        )
+
+    def _flush_thinking_as_text_block() -> None:
         nonlocal current_thinking_content
         if current_thinking_content is None:
             return
-        degraded_thinking_texts.append(current_thinking_content)
+        if block := _degraded_thinking_block(current_thinking_content):
+            content.append(block)
         current_thinking_content = None
 
     for part in msg.parts:
         if isinstance(part, message.ThinkingTextPart):
             if id(part) not in native_thinking_ids:
+                if block := _degraded_thinking_block(part.text):
+                    content.append(block)
                 continue
             current_thinking_content = part.text
             continue
@@ -142,7 +157,7 @@ def _assistant_message_to_message(msg: message.AssistantMessage, model_name: str
                 current_thinking_content = None
             continue
 
-        _flush_thinking()
+        _flush_thinking_as_text_block()
         if isinstance(part, message.TextPart):
             content.append(cast(BetaTextBlockParam, {"type": "text", "text": part.text}))
         elif isinstance(part, message.ToolCallPart):
@@ -166,11 +181,7 @@ def _assistant_message_to_message(msg: message.AssistantMessage, model_name: str
                 )
             )
 
-    _flush_thinking()
-
-    if degraded_thinking_texts:
-        degraded_text = "<thinking>\n" + "\n".join(degraded_thinking_texts) + "\n</thinking>"
-        content.insert(0, cast(BetaTextBlockParam, {"type": "text", "text": degraded_text}))
+    _flush_thinking_as_text_block()
 
     return {"role": "assistant", "content": content}