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

klaude_code/llm/google/input.py

@@ -11,7 +11,7 @@ from typing import Any
 from google.genai import types
 
 from klaude_code.const import EMPTY_TOOL_OUTPUT_MESSAGE
-from klaude_code.llm.image import parse_data_url
+from klaude_code.llm.image import assistant_image_to_data_url, parse_data_url
 from klaude_code.llm.input_common import (
     DeveloperAttachment,
     attach_developer_messages,
@@ -108,51 +108,49 @@ def _tool_messages_to_contents(
     return contents
 
 
+def _decode_thought_signature(sig: str | None) -> bytes | None:
+    """Decode base64 thought signature to bytes."""
+    if not sig:
+        return None
+    try:
+        return b64decode(sig)
+    except (BinasciiError, ValueError):
+        return None
+
+
 def _assistant_message_to_content(msg: message.AssistantMessage, model_name: str | None) -> types.Content | None:
     parts: list[types.Part] = []
     native_thinking_parts, degraded_thinking_texts = split_thinking_parts(msg, model_name)
     native_thinking_ids = {id(part) for part in native_thinking_parts}
-    pending_thought_text: str | None = None
-    pending_thought_signature: str | None = None
-
-    def flush_thought() -> None:
-        nonlocal pending_thought_text, pending_thought_signature
-        if pending_thought_text is None and pending_thought_signature is None:
-            return
-
-        signature_bytes: bytes | None = None
-        if pending_thought_signature:
-            try:
-                signature_bytes = b64decode(pending_thought_signature)
-            except (BinasciiError, ValueError):
-                signature_bytes = None
-
-        parts.append(
-            types.Part(
-                text=pending_thought_text or "",
-                thought=True,
-                thought_signature=signature_bytes,
-            )
-        )
-        pending_thought_text = None
-        pending_thought_signature = None
 
     for part in msg.parts:
         if isinstance(part, message.ThinkingTextPart):
             if id(part) not in native_thinking_ids:
                 continue
-            pending_thought_text = part.text
-            continue
-        if isinstance(part, message.ThinkingSignaturePart):
+            parts.append(types.Part(text=part.text, thought=True))
+
+        elif isinstance(part, message.ThinkingSignaturePart):
             if id(part) not in native_thinking_ids:
                 continue
-            if part.signature and (part.format or "").startswith("google"):
-                pending_thought_signature = part.signature
-            continue
-
-        flush_thought()
-        if isinstance(part, message.TextPart):
+            if not part.signature or part.format != "google":
+                continue
+            # Attach signature to the previous part
+            if parts:
+                sig_bytes = _decode_thought_signature(part.signature)
+                if sig_bytes:
+                    last_part = parts[-1]
+                    parts[-1] = types.Part(
+                        text=last_part.text,
+                        thought=last_part.thought,
+                        function_call=last_part.function_call,
+                        inline_data=last_part.inline_data,
+                        file_data=last_part.file_data,
+                        thought_signature=sig_bytes,
+                    )
+
+        elif isinstance(part, message.TextPart):
             parts.append(types.Part(text=part.text))
+
         elif isinstance(part, message.ToolCallPart):
             args: dict[str, Any]
             if part.arguments_json:
@@ -162,9 +160,19 @@ def _assistant_message_to_content(msg: message.AssistantMessage, model_name: str
                     args = {"_raw": part.arguments_json}
             else:
                 args = {}
-            parts.append(types.Part(function_call=types.FunctionCall(id=part.call_id, name=part.tool_name, args=args)))
+            parts.append(
+                types.Part(
+                    function_call=types.FunctionCall(id=part.call_id, name=part.tool_name, args=args),
+                )
+            )
 
-    flush_thought()
+        elif isinstance(part, message.ImageFilePart):
+            # Convert saved image back to inline_data for multi-turn
+            try:
+                data_url = assistant_image_to_data_url(part)
+                parts.append(_image_part_to_part(message.ImageURLPart(url=data_url)))
+            except (ValueError, FileNotFoundError):
+                pass  # Skip if image cannot be loaded
 
     if degraded_thinking_texts:
         parts.insert(0, types.Part(text="<thinking>\n" + "\n".join(degraded_thinking_texts) + "\n</thinking>"))

klaude_code/llm/openai_compatible/stream.py

@@ -2,8 +2,8 @@
 
 This module provides reusable primitives for OpenAI-compatible providers:
 
-- ``StreamStateManager``: accumulates assistant content and tool calls.
-- ``ReasoningHandlerABC``: provider-specific reasoning extraction + buffering.
+- ``StreamStateManager``: accumulates assistant parts in stream order.
+- ``ReasoningHandlerABC``: provider-specific reasoning extraction.
 - ``OpenAILLMStream``: LLMStream implementation for OpenAI-compatible clients.
 
 OpenRouter uses the same OpenAI Chat Completions API surface but differs in
@@ -15,7 +15,7 @@ from __future__ import annotations
 from abc import ABC, abstractmethod
 from collections.abc import AsyncGenerator, Callable
 from dataclasses import dataclass
-from typing import Any, Literal, cast
+from typing import Any, cast
 
 import httpx
 import openai
@@ -26,91 +26,107 @@ from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
 
 from klaude_code.llm.client import LLMStreamABC
 from klaude_code.llm.image import save_assistant_image
-from klaude_code.llm.openai_compatible.tool_call_accumulator import BasicToolCallAccumulator, ToolCallAccumulatorABC
-from klaude_code.llm.partial_message import degrade_thinking_to_text
+from klaude_code.llm.openai_compatible.tool_call_accumulator import normalize_tool_name
+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, convert_usage
 from klaude_code.protocol import llm_param, message, model
 
-StreamStage = Literal["waiting", "reasoning", "assistant", "tool"]
-
 
 class StreamStateManager:
-    """Manages streaming state and provides flush operations for accumulated content.
+    """Manages streaming state and accumulates parts in stream order.
 
-    This class encapsulates the common state management logic used by both
-    OpenAI-compatible and OpenRouter clients, reducing code duplication.
+    The persisted AssistantMessage is built directly from ``assistant_parts``.
+    ``get_partial_message()`` returns a best-effort message on cancellation.
     """
 
     def __init__(
         self,
         param_model: str,
         response_id: str | None = None,
-        reasoning_flusher: Callable[[], list[message.Part]] | None = None,
     ):
         self.param_model = param_model
         self.response_id = response_id
-        self.stage: StreamStage = "waiting"
-        self.accumulated_content: list[str] = []
-        self.accumulated_images: list[message.ImageFilePart] = []
-        self.accumulated_tool_calls: ToolCallAccumulatorABC = BasicToolCallAccumulator()
-        self.emitted_tool_start_indices: set[int] = set()
-        self._reasoning_flusher = reasoning_flusher
-        self.parts: list[message.Part] = []
+        self.assistant_parts: list[message.Part] = []
+        self._image_index: int = 0
+        self._tool_part_index_by_tc_index: dict[int, int] = {}
+        self._emitted_tool_start_indices: set[int] = set()
         self.stop_reason: model.StopReason | None = None
 
     def set_response_id(self, response_id: str) -> None:
         """Set the response ID once received from the stream."""
         self.response_id = response_id
-        self.accumulated_tool_calls.set_response_id(response_id)
 
-    def flush_reasoning(self) -> None:
-        """Flush accumulated reasoning content into parts."""
-        if self._reasoning_flusher is not None:
-            self.parts.extend(self._reasoning_flusher())
+    def append_thinking_text(self, text: str) -> None:
+        """Append thinking text, merging with the previous ThinkingTextPart when possible."""
+        append_thinking_text_part(self.assistant_parts, text, model_id=self.param_model)
+
+    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 append_image(self, image_part: message.ImageFilePart) -> None:
+        self.assistant_parts.append(image_part)
+        self._image_index += 1
+
+    def upsert_tool_call(self, *, tc_index: int, call_id: str | None, name: str | None, arguments: str | None) -> None:
+        """Insert a ToolCallPart at first sight and keep updating its fields.
+
+        Chat Completions streams tool call fields incrementally (name/id first,
+        then argument fragments). We keep the ToolCallPart in-place to preserve
+        stream order in the persisted AssistantMessage.
+        """
 
-    def flush_assistant(self) -> None:
-        """Flush accumulated assistant content into parts."""
-        if not self.accumulated_content and not self.accumulated_images:
+        part_index = self._tool_part_index_by_tc_index.get(tc_index)
+        if part_index is None:
+            tool_part = message.ToolCallPart(
+                call_id=call_id or "",
+                tool_name=normalize_tool_name(name or ""),
+                arguments_json=arguments or "",
+            )
+            self.assistant_parts.append(tool_part)
+            self._tool_part_index_by_tc_index[tc_index] = len(self.assistant_parts) - 1
             return
-        if self.accumulated_content:
-            self.parts.append(message.TextPart(text="".join(self.accumulated_content)))
-        if self.accumulated_images:
-            self.parts.extend(self.accumulated_images)
-        self.accumulated_content = []
-        self.accumulated_images = []
-        return
-
-    def flush_tool_calls(self) -> None:
-        """Flush accumulated tool calls into parts."""
-        items = self.accumulated_tool_calls.get()
-        if items:
-            self.parts.extend(items)
-            self.accumulated_tool_calls.reset()
-
-    def flush_all(self) -> list[message.Part]:
-        """Flush all accumulated content in order: reasoning, assistant, tool calls."""
-        self.flush_reasoning()
-        self.flush_assistant()
-        if self.stage == "tool":
-            self.flush_tool_calls()
-        return list(self.parts)
+
+        existing = self.assistant_parts[part_index]
+        if not isinstance(existing, message.ToolCallPart):
+            return
+
+        if call_id and not existing.call_id:
+            existing.call_id = call_id
+        if name and not existing.tool_name:
+            existing.tool_name = normalize_tool_name(name)
+        if arguments:
+            existing.arguments_json += arguments
+
+    def mark_tool_start_emitted(self, tc_index: int) -> bool:
+        """Return True if this is the first time we emit ToolCallStartDelta for this index."""
+        if tc_index in self._emitted_tool_start_indices:
+            return False
+        self._emitted_tool_start_indices.add(tc_index)
+        return True
+
+    def next_image_index(self) -> int:
+        return self._image_index
+
+    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.
 
-        Flushes all accumulated content (reasoning, assistant text, tool calls)
-        and returns the message. Returns None if no content has been accumulated.
+        Filters out tool calls and degrades thinking content for safety.
+        Returns None if no content has been accumulated.
         """
-        self.flush_reasoning()
-        self.flush_assistant()
-        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)
 
 
 @dataclass(slots=True)
@@ -148,7 +164,6 @@ class DefaultReasoningHandler(ReasoningHandlerABC):
     ) -> None:
         self._param_model = param_model
         self._response_id = response_id
-        self._accumulated: list[str] = []
 
     def set_response_id(self, response_id: str | None) -> None:
         self._response_id = response_id
@@ -158,18 +173,10 @@ class DefaultReasoningHandler(ReasoningHandlerABC):
         if not reasoning_content:
             return ReasoningDeltaResult(handled=False, outputs=[])
         text = str(reasoning_content)
-        self._accumulated.append(text)
         return ReasoningDeltaResult(handled=True, outputs=[text])
 
     def flush(self) -> list[message.Part]:
-        if not self._accumulated:
-            return []
-        item = message.ThinkingTextPart(
-            text="".join(self._accumulated),
-            model_id=self._param_model,
-        )
-        self._accumulated = []
-        return [item]
+        return []
 
 
 def _map_finish_reason(reason: str) -> model.StopReason | None:
@@ -254,26 +261,21 @@ async def parse_chat_completions_stream(
             # Reasoning
             reasoning_result = reasoning_handler.on_delta(delta)
             if reasoning_result.handled:
-                state.stage = "reasoning"
                 for output in reasoning_result.outputs:
                     if isinstance(output, str):
                         if not output:
                             continue
                         metadata_tracker.record_token()
+                        state.append_thinking_text(output)
                         yield message.ThinkingTextDelta(content=output, response_id=state.response_id)
                     else:
-                        state.parts.append(output)
+                        state.assistant_parts.append(output)
 
             # Assistant
             images = getattr(delta, "images", None)
             if isinstance(images, list) and images:
                 images_list = cast(list[object], images)
                 metadata_tracker.record_token()
-                if state.stage == "reasoning":
-                    state.flush_reasoning()
-                elif state.stage == "tool":
-                    state.flush_tool_calls()
-                state.stage = "assistant"
                 for image_obj in images_list:
                     url = _extract_image_url(image_obj)
                     if not url:
@@ -286,50 +288,59 @@ async def parse_chat_completions_stream(
                             data_url=url,
                             session_id=param.session_id,
                             response_id=state.response_id,
-                            image_index=len(state.accumulated_images),
+                            image_index=state.next_image_index(),
                         )
                     except ValueError as exc:
                         yield message.StreamErrorItem(error=str(exc))
                         return
-                    state.accumulated_images.append(assistant_image)
+                    state.append_image(assistant_image)
                     yield message.AssistantImageDelta(
                         response_id=state.response_id, file_path=assistant_image.file_path
                     )
 
-            if (content := getattr(delta, "content", None)) and (state.stage == "assistant" or str(content).strip()):
+            content_str = str(content) if (content := getattr(delta, "content", None)) is not None else ""
+
+            if content_str and (
+                (state.assistant_parts and isinstance(state.assistant_parts[-1], message.TextPart))
+                or content_str.strip()
+            ):
                 metadata_tracker.record_token()
-                if state.stage == "reasoning":
-                    state.flush_reasoning()
-                elif state.stage == "tool":
-                    state.flush_tool_calls()
-                state.stage = "assistant"
-                state.accumulated_content.append(str(content))
+                state.append_text(content_str)
                 yield message.AssistantTextDelta(
-                    content=str(content),
+                    content=content_str,
                     response_id=state.response_id,
                 )
 
             # Tool
             if (tool_calls := getattr(delta, "tool_calls", None)) and len(tool_calls) > 0:
                 metadata_tracker.record_token()
-                if state.stage == "reasoning":
-                    state.flush_reasoning()
-                elif state.stage == "assistant":
-                    state.flush_assistant()
-                state.stage = "tool"
                 for tc in tool_calls:
-                    if tc.index not in state.emitted_tool_start_indices and tc.function and tc.function.name:
-                        state.emitted_tool_start_indices.add(tc.index)
+                    tc_index = getattr(tc, "index", None)
+                    if not isinstance(tc_index, int):
+                        continue
+                    fn = getattr(tc, "function", None)
+                    fn_name = getattr(fn, "name", None) if fn is not None else None
+                    fn_args = getattr(fn, "arguments", None) if fn is not None else None
+                    tc_id = getattr(tc, "id", None)
+
+                    if fn_name and state.mark_tool_start_emitted(tc_index):
                         yield message.ToolCallStartDelta(
                             response_id=state.response_id,
-                            call_id=tc.id or "",
-                            name=tc.function.name,
+                            call_id=str(tc_id or ""),
+                            name=str(fn_name),
                         )
-                state.accumulated_tool_calls.add(tool_calls)
+                    state.upsert_tool_call(
+                        tc_index=tc_index,
+                        call_id=str(tc_id) if isinstance(tc_id, str) else None,
+                        name=str(fn_name) if isinstance(fn_name, str) else None,
+                        arguments=str(fn_args) if isinstance(fn_args, str) else None,
+                    )
     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()
+    # 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)
     if parts:
         metadata_tracker.record_token()
     metadata_tracker.set_response_id(state.response_id)
@@ -361,7 +372,6 @@ class OpenAILLMStream(LLMStreamABC):
         self._on_event = on_event
         self._state = StreamStateManager(
             param_model=str(param.model_id),
-            reasoning_flusher=reasoning_handler.flush,
         )
         self._completed = False
 

klaude_code/llm/openrouter/reasoning.py

@@ -30,7 +30,6 @@ class ReasoningStreamHandler(ReasoningHandlerABC):
         self._response_id = response_id
 
         self._reasoning_id: str | None = None
-        self._accumulated_reasoning: list[str] = []
 
     def set_response_id(self, response_id: str | None) -> None:
         """Update the response identifier used for emitted items."""
@@ -62,44 +61,20 @@ class ReasoningStreamHandler(ReasoningHandlerABC):
 
         if detail.type == "reasoning.encrypted":
             self._reasoning_id = detail.id
-            # Flush accumulated text before encrypted content
-            items.extend(self._flush_text())
             if signature_part := self._build_signature_part(detail.data, detail):
                 items.append(signature_part)
             return items
 
         if detail.type in ("reasoning.text", "reasoning.summary"):
             self._reasoning_id = detail.id
-            # Accumulate text
-            text = detail.text if detail.type == "reasoning.text" else detail.summary
-            if text:
-                self._accumulated_reasoning.append(text)
-            # Flush on signature (encrypted content)
-            if detail.signature:
-                items.extend(self._flush_text())
-                if signature_part := self._build_signature_part(detail.signature, detail):
-                    items.append(signature_part)
+            # Signature (Anthropic-style) can arrive alongside text/summary.
+            if detail.signature and (signature_part := self._build_signature_part(detail.signature, detail)):
+                items.append(signature_part)
 
         return items
 
     def flush(self) -> list[message.Part]:
-        """Flush buffered reasoning text on finalize."""
-        return self._flush_text()
-
-    def _flush_text(self) -> list[message.Part]:
-        """Flush accumulated reasoning text as a single part."""
-        if not self._accumulated_reasoning:
-            return []
-        item = self._build_text_part("".join(self._accumulated_reasoning))
-        self._accumulated_reasoning = []
-        return [item]
-
-    def _build_text_part(self, content: str) -> message.ThinkingTextPart:
-        return message.ThinkingTextPart(
-            id=self._reasoning_id,
-            text=content,
-            model_id=self._param_model,
-        )
+        return []
 
     def _build_signature_part(
         self,

klaude_code/llm/partial_message.py

@@ -1,35 +1,5 @@
 from __future__ import annotations
 
-from klaude_code.protocol import message
+from klaude_code.llm.stream_parts import degrade_thinking_to_text
 
-
-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]
+__all__ = ["degrade_thinking_to_text"]