klaude-code 2.4.1__py3-none-any.whl → 2.5.0__py3-none-any.whl

klaude_code/llm/google/client.py

@@ -21,11 +21,12 @@ from google.genai.types import (
     UsageMetadata,
 )
 
-from klaude_code.llm.client import LLMClientABC
+from klaude_code.llm.client import LLMClientABC, LLMStreamABC
 from klaude_code.llm.google.input import convert_history_to_contents, convert_tool_schema
 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.usage import MetadataTracker
+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
 
@@ -135,19 +136,83 @@ def _map_finish_reason(reason: str) -> model.StopReason | None:
     return mapping.get(normalized)
 
 
+class GoogleStreamStateManager:
+    """Manages streaming state for Google LLM responses.
+
+    Accumulates thinking content, assistant text, and tool calls during streaming
+    to support get_partial_message() for cancellation scenarios.
+    """
+
+    def __init__(self, param_model: str) -> None:
+        self.param_model = param_model
+        self.accumulated_thoughts: list[str] = []
+        self.accumulated_text: list[str] = []
+        self.thought_signature: str | None = None
+        self.assistant_parts: list[message.Part] = []
+        self.response_id: str | None = None
+        self.stop_reason: model.StopReason | None = None
+
+    def flush_thinking(self) -> None:
+        """Flush accumulated thinking content into assistant_parts."""
+        if self.accumulated_thoughts:
+            self.assistant_parts.append(
+                message.ThinkingTextPart(
+                    text="".join(self.accumulated_thoughts),
+                    model_id=self.param_model,
+                )
+            )
+            self.accumulated_thoughts.clear()
+        if self.thought_signature:
+            self.assistant_parts.append(
+                message.ThinkingSignaturePart(
+                    signature=self.thought_signature,
+                    model_id=self.param_model,
+                    format="google_thought_signature",
+                )
+            )
+            self.thought_signature = None
+
+    def flush_text(self) -> None:
+        """Flush accumulated text content into assistant_parts."""
+        if not self.accumulated_text:
+            return
+        self.assistant_parts.append(message.TextPart(text="".join(self.accumulated_text)))
+        self.accumulated_text.clear()
+
+    def get_partial_message(self) -> message.AssistantMessage | None:
+        """Build a partial AssistantMessage from accumulated state.
+
+        Flushes all accumulated content and returns the message.
+        Returns None if no content has been accumulated yet.
+        """
+        self.flush_thinking()
+        self.flush_text()
+
+        filtered_parts: list[message.Part] = []
+        for part in self.assistant_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",
+        )
+
+
 async def parse_google_stream(
     stream: AsyncIterator[Any],
     param: llm_param.LLMCallParameter,
     metadata_tracker: MetadataTracker,
+    state: GoogleStreamStateManager,
 ) -> AsyncGenerator[message.LLMStreamItem]:
-    response_id: str | None = None
     stage: Literal["waiting", "thinking", "assistant", "tool"] = "waiting"
 
-    accumulated_text: list[str] = []
-    accumulated_thoughts: list[str] = []
-    thought_signature: str | None = None
-    assistant_parts: list[message.Part] = []
-
     # Track tool calls where args arrive as partial updates.
     partial_args_by_call: dict[str, dict[str, Any]] = {}
     started_tool_calls: dict[str, str] = {}  # call_id -> name
@@ -155,33 +220,6 @@ async def parse_google_stream(
     completed_tool_items: set[str] = set()
 
     last_usage_metadata: UsageMetadata | None = None
-    stop_reason: model.StopReason | None = None
-
-    def flush_thinking() -> None:
-        nonlocal thought_signature
-        if accumulated_thoughts:
-            assistant_parts.append(
-                message.ThinkingTextPart(
-                    text="".join(accumulated_thoughts),
-                    model_id=str(param.model_id),
-                )
-            )
-            accumulated_thoughts.clear()
-        if thought_signature:
-            assistant_parts.append(
-                message.ThinkingSignaturePart(
-                    signature=thought_signature,
-                    model_id=str(param.model_id),
-                    format="google_thought_signature",
-                )
-            )
-            thought_signature = None
-
-    def flush_text() -> None:
-        if not accumulated_text:
-            return
-        assistant_parts.append(message.TextPart(text="".join(accumulated_text)))
-        accumulated_text.clear()
 
     async for chunk in stream:
         log_debug(
@@ -190,8 +228,8 @@ async def parse_google_stream(
             debug_type=DebugType.LLM_STREAM,
         )
 
-        if response_id is None:
-            response_id = getattr(chunk, "response_id", None) or uuid4().hex
+        if state.response_id is None:
+            state.response_id = getattr(chunk, "response_id", None) or uuid4().hex
 
         if getattr(chunk, "usage_metadata", None) is not None:
             last_usage_metadata = chunk.usage_metadata
@@ -204,7 +242,7 @@ async def parse_google_stream(
                 reason_value = finish_reason
             else:
                 reason_value = getattr(finish_reason, "name", None) or str(finish_reason)
-            stop_reason = _map_finish_reason(reason_value)
+            state.stop_reason = _map_finish_reason(reason_value)
         content = getattr(candidate0, "content", None) if candidate0 else None
         content_parts = getattr(content, "parts", None) if content else None
         if not content_parts:
@@ -218,18 +256,18 @@ async def parse_google_stream(
                 metadata_tracker.record_token()
                 if getattr(part, "thought", False) is True:
                     if stage == "assistant":
-                        flush_text()
+                        state.flush_text()
                     stage = "thinking"
-                    accumulated_thoughts.append(text)
+                    state.accumulated_thoughts.append(text)
                     if getattr(part, "thought_signature", None):
-                        thought_signature = part.thought_signature
-                    yield message.ThinkingTextDelta(content=text, response_id=response_id)
+                        state.thought_signature = part.thought_signature
+                    yield message.ThinkingTextDelta(content=text, response_id=state.response_id)
                 else:
                     if stage == "thinking":
-                        flush_thinking()
+                        state.flush_thinking()
                     stage = "assistant"
-                    accumulated_text.append(text)
-                    yield message.AssistantTextDelta(content=text, response_id=response_id)
+                    state.accumulated_text.append(text)
+                    yield message.AssistantTextDelta(content=text, response_id=state.response_id)
 
             function_call = getattr(part, "function_call", None)
             if function_call is None:
@@ -242,16 +280,16 @@ async def parse_google_stream(
 
             if call_id not in started_tool_items:
                 started_tool_items.add(call_id)
-                yield message.ToolCallStartDelta(response_id=response_id, call_id=call_id, name=name)
+                yield message.ToolCallStartDelta(response_id=state.response_id, call_id=call_id, name=name)
 
             args_obj = getattr(function_call, "args", None)
             if args_obj is not None:
                 if stage == "thinking":
-                    flush_thinking()
+                    state.flush_thinking()
                 if stage == "assistant":
-                    flush_text()
+                    state.flush_text()
                 stage = "tool"
-                assistant_parts.append(
+                state.assistant_parts.append(
                     message.ToolCallPart(
                         call_id=call_id,
                         tool_name=name,
@@ -269,11 +307,11 @@ async def parse_google_stream(
             will_continue = getattr(function_call, "will_continue", None)
             if will_continue is False and call_id in partial_args_by_call and call_id not in completed_tool_items:
                 if stage == "thinking":
-                    flush_thinking()
+                    state.flush_thinking()
                 if stage == "assistant":
-                    flush_text()
+                    state.flush_text()
                 stage = "tool"
-                assistant_parts.append(
+                state.assistant_parts.append(
                     message.ToolCallPart(
                         call_id=call_id,
                         tool_name=name,
@@ -287,7 +325,7 @@ async def parse_google_stream(
         if call_id in completed_tool_items:
             continue
         args = partial_args_by_call.get(call_id, {})
-        assistant_parts.append(
+        state.assistant_parts.append(
             message.ToolCallPart(
                 call_id=call_id,
                 tool_name=name,
@@ -295,23 +333,64 @@ async def parse_google_stream(
             )
         )
 
-    flush_thinking()
-    flush_text()
+    state.flush_thinking()
+    state.flush_text()
 
     usage = _usage_from_metadata(last_usage_metadata, context_limit=param.context_limit, max_tokens=param.max_tokens)
     if usage is not None:
         metadata_tracker.set_usage(usage)
     metadata_tracker.set_model_name(str(param.model_id))
-    metadata_tracker.set_response_id(response_id)
+    metadata_tracker.set_response_id(state.response_id)
     metadata = metadata_tracker.finalize()
     yield message.AssistantMessage(
-        parts=assistant_parts,
-        response_id=response_id,
+        parts=state.assistant_parts,
+        response_id=state.response_id,
         usage=metadata,
-        stop_reason=stop_reason,
+        stop_reason=state.stop_reason,
     )
 
 
+class GoogleLLMStream(LLMStreamABC):
+    """LLMStream implementation for Google LLM clients."""
+
+    def __init__(
+        self,
+        stream: AsyncIterator[Any],
+        *,
+        param: llm_param.LLMCallParameter,
+        metadata_tracker: MetadataTracker,
+        state: GoogleStreamStateManager,
+    ) -> None:
+        self._stream = stream
+        self._param = param
+        self._metadata_tracker = metadata_tracker
+        self._state = state
+        self._completed = False
+
+    def __aiter__(self) -> AsyncGenerator[message.LLMStreamItem]:
+        return self._iterate()
+
+    async def _iterate(self) -> AsyncGenerator[message.LLMStreamItem]:
+        try:
+            async for item in parse_google_stream(
+                self._stream,
+                param=self._param,
+                metadata_tracker=self._metadata_tracker,
+                state=self._state,
+            ):
+                if isinstance(item, message.AssistantMessage):
+                    self._completed = True
+                yield item
+        except (APIError, ClientError, ServerError, httpx.HTTPError) as e:
+            yield message.StreamErrorItem(error=f"{e.__class__.__name__} {e!s}")
+            yield message.AssistantMessage(parts=[], response_id=None, usage=self._metadata_tracker.finalize())
+
+    def get_partial_message(self) -> message.AssistantMessage | None:
+        if self._completed:
+            return None
+        return self._state.get_partial_message()
+
+
 @register(llm_param.LLMClientProtocol.GOOGLE)
 class GoogleClient(LLMClientABC):
     def __init__(self, config: llm_param.LLMConfigParameter):
@@ -332,7 +411,7 @@ class GoogleClient(LLMClientABC):
         return cls(config)
 
     @override
-    async def call(self, param: llm_param.LLMCallParameter) -> AsyncGenerator[message.LLMStreamItem]:
+    async def call(self, param: llm_param.LLMCallParameter) -> LLMStreamABC:
         param = apply_config_defaults(param, self.get_llm_config())
         metadata_tracker = MetadataTracker(cost_config=self.get_llm_config().cost)
 
@@ -359,13 +438,15 @@ class GoogleClient(LLMClientABC):
                 config=config,
             )
         except (APIError, ClientError, ServerError, httpx.HTTPError) as e:
-            yield message.StreamErrorItem(error=f"{e.__class__.__name__} {e!s}")
-            yield message.AssistantMessage(parts=[], response_id=None, usage=metadata_tracker.finalize())
-            return
+            return error_llm_stream(
+                metadata_tracker,
+                error=f"{e.__class__.__name__} {e!s}",
+            )
 
-        try:
-            async for item in parse_google_stream(stream, param=param, metadata_tracker=metadata_tracker):
-                yield item
-        except (APIError, ClientError, ServerError, httpx.HTTPError) as e:
-            yield message.StreamErrorItem(error=f"{e.__class__.__name__} {e!s}")
-            yield message.AssistantMessage(parts=[], response_id=None, usage=metadata_tracker.finalize())
+        state = GoogleStreamStateManager(param_model=str(param.model_id))
+        return GoogleLLMStream(
+            stream,
+            param=param,
+            metadata_tracker=metadata_tracker,
+            state=state,
+        )

klaude_code/llm/openai_compatible/client.py

@@ -1,5 +1,4 @@
 import json
-from collections.abc import AsyncGenerator
 from typing import Any, override
 
 import httpx
@@ -7,14 +6,14 @@ import openai
 from openai.types.chat.completion_create_params import CompletionCreateParamsStreaming
 
 from klaude_code.const import LLM_HTTP_TIMEOUT_CONNECT, LLM_HTTP_TIMEOUT_READ, LLM_HTTP_TIMEOUT_TOTAL
-from klaude_code.llm.client import LLMClientABC
+from klaude_code.llm.client import LLMClientABC, LLMStreamABC
 from klaude_code.llm.input_common import apply_config_defaults
 from klaude_code.llm.openai_compatible.input import convert_history_to_input, convert_tool_schema
-from klaude_code.llm.openai_compatible.stream import DefaultReasoningHandler, parse_chat_completions_stream
+from klaude_code.llm.openai_compatible.stream import DefaultReasoningHandler, OpenAILLMStream
 from klaude_code.llm.registry import register
-from klaude_code.llm.usage import MetadataTracker
+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
+from klaude_code.protocol import llm_param
 
 
 def build_payload(param: llm_param.LLMCallParameter) -> tuple[CompletionCreateParamsStreaming, dict[str, object]]:
@@ -77,7 +76,7 @@ class OpenAICompatibleClient(LLMClientABC):
         return cls(config)
 
     @override
-    async def call(self, param: llm_param.LLMCallParameter) -> AsyncGenerator[message.LLMStreamItem]:
+    async def call(self, param: llm_param.LLMCallParameter) -> LLMStreamABC:
         param = apply_config_defaults(param, self.get_llm_config())
 
         metadata_tracker = MetadataTracker(cost_config=self.get_llm_config().cost)
@@ -85,9 +84,8 @@ class OpenAICompatibleClient(LLMClientABC):
         try:
             payload, extra_body = build_payload(param)
         except (ValueError, OSError) as e:
-            yield message.StreamErrorItem(error=f"{e.__class__.__name__} {e!s}")
-            yield message.AssistantMessage(parts=[], response_id=None, usage=metadata_tracker.finalize())
-            return
+            return error_llm_stream(metadata_tracker, error=f"{e.__class__.__name__} {e!s}")
+
         extra_headers: dict[str, str] = {"extra": json.dumps({"session_id": param.session_id}, sort_keys=True)}
 
         log_debug(
@@ -103,9 +101,7 @@ class OpenAICompatibleClient(LLMClientABC):
                 extra_headers=extra_headers,
             )
         except (openai.OpenAIError, httpx.HTTPError) as e:
-            yield message.StreamErrorItem(error=f"{e.__class__.__name__} {e!s}")
-            yield message.AssistantMessage(parts=[], response_id=None, usage=metadata_tracker.finalize())
-            return
+            return error_llm_stream(metadata_tracker, error=f"{e.__class__.__name__} {e!s}")
 
         reasoning_handler = DefaultReasoningHandler(
             param_model=str(param.model_id),
@@ -119,11 +115,10 @@ class OpenAICompatibleClient(LLMClientABC):
                 debug_type=DebugType.LLM_STREAM,
             )
 
-        async for item in parse_chat_completions_stream(
+        return OpenAILLMStream(
             stream,
             param=param,
             metadata_tracker=metadata_tracker,
             reasoning_handler=reasoning_handler,
             on_event=on_event,
-        ):
-            yield item
+        )

klaude_code/llm/openai_compatible/stream.py

@@ -4,7 +4,7 @@ This module provides reusable primitives for OpenAI-compatible providers:
 
 - ``StreamStateManager``: accumulates assistant content and tool calls.
 - ``ReasoningHandlerABC``: provider-specific reasoning extraction + buffering.
-- ``parse_chat_completions_stream``: shared stream loop that emits stream/history items.
+- ``OpenAILLMStream``: LLMStream implementation for OpenAI-compatible clients.
 
 OpenRouter uses the same OpenAI Chat Completions API surface but differs in
 how reasoning is represented (``reasoning_details`` vs ``reasoning_content``).
@@ -24,8 +24,10 @@ import pydantic
 from openai import AsyncStream
 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.usage import MetadataTracker, convert_usage
 from klaude_code.protocol import llm_param, message, model
 
@@ -93,6 +95,23 @@ class StreamStateManager:
             self.flush_tool_calls()
         return list(self.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.
+        """
+        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",
+        )
+
 
 @dataclass(slots=True)
 class ReasoningDeltaResult:
@@ -168,6 +187,7 @@ def _map_finish_reason(reason: str) -> model.StopReason | None:
 async def parse_chat_completions_stream(
     stream: AsyncStream[ChatCompletionChunk],
     *,
+    state: StreamStateManager,
     param: llm_param.LLMCallParameter,
     metadata_tracker: MetadataTracker,
     reasoning_handler: ReasoningHandlerABC,
@@ -176,13 +196,10 @@ async def parse_chat_completions_stream(
     """Parse OpenAI Chat Completions stream into stream items.
 
     This is shared by OpenAI-compatible and OpenRouter clients.
+    The state parameter allows external access to accumulated content
+    for cancellation scenarios.
     """
 
-    state = StreamStateManager(
-        param_model=str(param.model_id),
-        reasoning_flusher=reasoning_handler.flush,
-    )
-
     def _extract_image_url(image_obj: object) -> str | None:
         image_url = getattr(image_obj, "image_url", None)
         if image_url is not None:
@@ -323,3 +340,48 @@ async def parse_chat_completions_stream(
         usage=metadata,
         stop_reason=state.stop_reason,
     )
+
+
+class OpenAILLMStream(LLMStreamABC):
+    """LLMStream implementation for OpenAI-compatible clients."""
+
+    def __init__(
+        self,
+        stream: AsyncStream[ChatCompletionChunk],
+        *,
+        param: llm_param.LLMCallParameter,
+        metadata_tracker: MetadataTracker,
+        reasoning_handler: ReasoningHandlerABC,
+        on_event: Callable[[object], None] | None = None,
+    ) -> None:
+        self._stream = stream
+        self._param = param
+        self._metadata_tracker = metadata_tracker
+        self._reasoning_handler = reasoning_handler
+        self._on_event = on_event
+        self._state = StreamStateManager(
+            param_model=str(param.model_id),
+            reasoning_flusher=reasoning_handler.flush,
+        )
+        self._completed = False
+
+    def __aiter__(self) -> AsyncGenerator[message.LLMStreamItem]:
+        return self._iterate()
+
+    async def _iterate(self) -> AsyncGenerator[message.LLMStreamItem]:
+        async for item in parse_chat_completions_stream(
+            self._stream,
+            state=self._state,
+            param=self._param,
+            metadata_tracker=self._metadata_tracker,
+            reasoning_handler=self._reasoning_handler,
+            on_event=self._on_event,
+        ):
+            if isinstance(item, message.AssistantMessage):
+                self._completed = True
+            yield item
+
+    def get_partial_message(self) -> message.AssistantMessage | None:
+        if self._completed:
+            return None
+        return self._state.get_partial_message()

klaude_code/llm/openrouter/client.py

@@ -1,5 +1,4 @@
 import json
-from collections.abc import AsyncGenerator
 from typing import Any, cast, override
 
 import httpx
@@ -14,16 +13,16 @@ from klaude_code.const import (
     LLM_HTTP_TIMEOUT_TOTAL,
     OPENROUTER_BASE_URL,
 )
-from klaude_code.llm.client import LLMClientABC
+from klaude_code.llm.client import LLMClientABC, LLMStreamABC
 from klaude_code.llm.input_common import apply_config_defaults
 from klaude_code.llm.openai_compatible.input import convert_tool_schema
-from klaude_code.llm.openai_compatible.stream import parse_chat_completions_stream
+from klaude_code.llm.openai_compatible.stream import OpenAILLMStream
 from klaude_code.llm.openrouter.input import convert_history_to_input, is_claude_model
 from klaude_code.llm.openrouter.reasoning import ReasoningStreamHandler
 from klaude_code.llm.registry import register
-from klaude_code.llm.usage import MetadataTracker
+from klaude_code.llm.usage import MetadataTracker, error_llm_stream
 from klaude_code.log import DebugType, is_debug_enabled, log_debug
-from klaude_code.protocol import llm_param, message
+from klaude_code.protocol import llm_param
 
 
 def build_payload(
@@ -103,7 +102,7 @@ class OpenRouterClient(LLMClientABC):
         return cls(config)
 
     @override
-    async def call(self, param: llm_param.LLMCallParameter) -> AsyncGenerator[message.LLMStreamItem]:
+    async def call(self, param: llm_param.LLMCallParameter) -> LLMStreamABC:
         param = apply_config_defaults(param, self.get_llm_config())
 
         metadata_tracker = MetadataTracker(cost_config=self.get_llm_config().cost)
@@ -111,9 +110,7 @@ class OpenRouterClient(LLMClientABC):
         try:
             payload, extra_body, extra_headers = build_payload(param)
         except (ValueError, OSError) as e:
-            yield message.StreamErrorItem(error=f"{e.__class__.__name__} {e!s}")
-            yield message.AssistantMessage(parts=[], response_id=None, usage=metadata_tracker.finalize())
-            return
+            return error_llm_stream(metadata_tracker, error=f"{e.__class__.__name__} {e!s}")
 
         log_debug(
             json.dumps({**payload, **extra_body}, ensure_ascii=False, default=str),
@@ -128,9 +125,7 @@ class OpenRouterClient(LLMClientABC):
                 extra_headers=extra_headers,
             )
         except (openai.OpenAIError, httpx.HTTPError) as e:
-            yield message.StreamErrorItem(error=f"{e.__class__.__name__} {e!s}")
-            yield message.AssistantMessage(parts=[], response_id=None, usage=metadata_tracker.finalize())
-            return
+            return error_llm_stream(metadata_tracker, error=f"{e.__class__.__name__} {e!s}")
 
         reasoning_handler = ReasoningStreamHandler(
             param_model=str(param.model_id),
@@ -144,11 +139,10 @@ class OpenRouterClient(LLMClientABC):
                 debug_type=DebugType.LLM_STREAM,
             )
 
-        async for item in parse_chat_completions_stream(
+        return OpenAILLMStream(
             stream,
             param=param,
             metadata_tracker=metadata_tracker,
             reasoning_handler=reasoning_handler,
             on_event=on_event,
-        ):
-            yield item
+        )

klaude_code/llm/partial_message.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from klaude_code.protocol import message
+
+
+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]