klaude-code 1.2.23__py3-none-any.whl → 1.2.25__py3-none-any.whl

klaude_code/llm/openai_compatible/stream.py

@@ -0,0 +1,272 @@
+"""Shared stream processing utilities for Chat Completions streaming.
+
+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 ConversationItems.
+
+OpenRouter uses the same OpenAI Chat Completions API surface but differs in
+how reasoning is represented (``reasoning_details`` vs ``reasoning_content``).
+"""
+
+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
+
+import httpx
+import openai
+import openai.types
+from openai import AsyncStream
+from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
+
+from klaude_code.llm.openai_compatible.tool_call_accumulator import BasicToolCallAccumulator, ToolCallAccumulatorABC
+from klaude_code.llm.usage import MetadataTracker, convert_usage
+from klaude_code.protocol import llm_param, model
+
+StreamStage = Literal["waiting", "reasoning", "assistant", "tool"]
+
+
+class StreamStateManager:
+    """Manages streaming state and provides flush operations for accumulated content.
+
+    This class encapsulates the common state management logic used by both
+    OpenAI-compatible and OpenRouter clients, reducing code duplication.
+    """
+
+    def __init__(
+        self,
+        param_model: str,
+        response_id: str | None = None,
+        reasoning_flusher: Callable[[], list[model.ConversationItem]] | None = None,
+    ):
+        self.param_model = param_model
+        self.response_id = response_id
+        self.stage: StreamStage = "waiting"
+        self.accumulated_reasoning: list[str] = []
+        self.accumulated_content: list[str] = []
+        self.accumulated_tool_calls: ToolCallAccumulatorABC = BasicToolCallAccumulator()
+        self.emitted_tool_start_indices: set[int] = set()
+        self._reasoning_flusher = reasoning_flusher
+
+    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.response_id = response_id  # pyright: ignore[reportAttributeAccessIssue]
+
+    def flush_reasoning(self) -> list[model.ConversationItem]:
+        """Flush accumulated reasoning content and return items."""
+        if self._reasoning_flusher is not None:
+            return self._reasoning_flusher()
+        if not self.accumulated_reasoning:
+            return []
+        item = model.ReasoningTextItem(
+            content="".join(self.accumulated_reasoning),
+            response_id=self.response_id,
+            model=self.param_model,
+        )
+        self.accumulated_reasoning = []
+        return [item]
+
+    def flush_assistant(self) -> list[model.ConversationItem]:
+        """Flush accumulated assistant content and return items."""
+        if not self.accumulated_content:
+            return []
+        item = model.AssistantMessageItem(
+            content="".join(self.accumulated_content),
+            response_id=self.response_id,
+        )
+        self.accumulated_content = []
+        return [item]
+
+    def flush_tool_calls(self) -> list[model.ToolCallItem]:
+        """Flush accumulated tool calls and return items."""
+        items: list[model.ToolCallItem] = self.accumulated_tool_calls.get()
+        if items:
+            self.accumulated_tool_calls.chunks_by_step = []  # pyright: ignore[reportAttributeAccessIssue]
+        return items
+
+    def flush_all(self) -> list[model.ConversationItem]:
+        """Flush all accumulated content in order: reasoning, assistant, tool calls."""
+        items: list[model.ConversationItem] = []
+        items.extend(self.flush_reasoning())
+        items.extend(self.flush_assistant())
+        if self.stage == "tool":
+            items.extend(self.flush_tool_calls())
+        return items
+
+
+@dataclass(slots=True)
+class ReasoningDeltaResult:
+    """Result of processing a single provider delta for reasoning signals."""
+
+    handled: bool
+    outputs: list[str | model.ConversationItem]
+
+
+class ReasoningHandlerABC(ABC):
+    """Provider-specific reasoning handler for Chat Completions streaming."""
+
+    @abstractmethod
+    def set_response_id(self, response_id: str | None) -> None:
+        """Update the response identifier used for emitted items."""
+
+    @abstractmethod
+    def on_delta(self, delta: object) -> ReasoningDeltaResult:
+        """Process a single delta and return ordered reasoning outputs."""
+
+    @abstractmethod
+    def flush(self) -> list[model.ConversationItem]:
+        """Flush buffered reasoning content (usually at stage transition/finalize)."""
+
+
+class DefaultReasoningHandler(ReasoningHandlerABC):
+    """Handles OpenAI-compatible reasoning fields (reasoning_content / reasoning)."""
+
+    def __init__(
+        self,
+        *,
+        param_model: str,
+        response_id: str | None,
+    ) -> 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
+
+    def on_delta(self, delta: object) -> ReasoningDeltaResult:
+        reasoning_content = getattr(delta, "reasoning_content", None) or getattr(delta, "reasoning", None) or ""
+        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[model.ConversationItem]:
+        if not self._accumulated:
+            return []
+        item = model.ReasoningTextItem(
+            content="".join(self._accumulated),
+            response_id=self._response_id,
+            model=self._param_model,
+        )
+        self._accumulated = []
+        return [item]
+
+
+async def parse_chat_completions_stream(
+    stream: AsyncStream[ChatCompletionChunk],
+    *,
+    param: llm_param.LLMCallParameter,
+    metadata_tracker: MetadataTracker,
+    reasoning_handler: ReasoningHandlerABC,
+    on_event: Callable[[object], None] | None = None,
+) -> AsyncGenerator[model.ConversationItem]:
+    """Parse OpenAI Chat Completions stream into ConversationItems.
+
+    This is shared by OpenAI-compatible and OpenRouter clients.
+    """
+
+    state = StreamStateManager(
+        param_model=str(param.model),
+        reasoning_flusher=reasoning_handler.flush,
+    )
+
+    try:
+        async for event in stream:
+            if on_event is not None:
+                on_event(event)
+
+            if not state.response_id and (event_id := getattr(event, "id", None)):
+                state.set_response_id(str(event_id))
+                reasoning_handler.set_response_id(str(event_id))
+                yield model.StartItem(response_id=str(event_id))
+
+            if (event_usage := getattr(event, "usage", None)) is not None:
+                metadata_tracker.set_usage(convert_usage(event_usage, param.context_limit, param.max_tokens))
+            if event_model := getattr(event, "model", None):
+                metadata_tracker.set_model_name(str(event_model))
+            if provider := getattr(event, "provider", None):
+                metadata_tracker.set_provider(str(provider))
+
+            choices = cast(Any, getattr(event, "choices", None))
+            if not choices:
+                continue
+
+            # Support Moonshot Kimi K2's usage field in choice
+            choice0 = choices[0]
+            if choice_usage := getattr(choice0, "usage", None):
+                try:
+                    usage = openai.types.CompletionUsage.model_validate(choice_usage)
+                    metadata_tracker.set_usage(convert_usage(usage, param.context_limit, param.max_tokens))
+                except Exception:
+                    pass
+
+            delta = cast(Any, getattr(choice0, "delta", None))
+            if delta is None:
+                continue
+
+            # 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()
+                        yield model.ReasoningTextDelta(content=output, response_id=state.response_id)
+                    else:
+                        yield output
+
+            # Assistant
+            if (content := getattr(delta, "content", None)) and (state.stage == "assistant" or str(content).strip()):
+                metadata_tracker.record_token()
+                if state.stage == "reasoning":
+                    for item in state.flush_reasoning():
+                        yield item
+                elif state.stage == "tool":
+                    for item in state.flush_tool_calls():
+                        yield item
+                state.stage = "assistant"
+                state.accumulated_content.append(str(content))
+                yield model.AssistantMessageDelta(
+                    content=str(content),
+                    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":
+                    for item in state.flush_reasoning():
+                        yield item
+                elif state.stage == "assistant":
+                    for item in state.flush_assistant():
+                        yield item
+                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)
+                        yield model.ToolCallStartItem(
+                            response_id=state.response_id,
+                            call_id=tc.id or "",
+                            name=tc.function.name,
+                        )
+                state.accumulated_tool_calls.add(tool_calls)
+    except (openai.OpenAIError, httpx.HTTPError) as e:
+        yield model.StreamErrorItem(error=f"{e.__class__.__name__} {e!s}")
+
+    flushed_items = state.flush_all()
+    if flushed_items:
+        metadata_tracker.record_token()
+    for item in flushed_items:
+        yield item
+
+    metadata_tracker.set_response_id(state.response_id)
+    yield metadata_tracker.finalize()

klaude_code/llm/openrouter/client.py

@@ -1,6 +1,6 @@
 import json
 from collections.abc import AsyncGenerator
-from typing import override
+from typing import Any, override
 
 import httpx
 import openai
@@ -9,13 +9,13 @@ from openai.types.chat.completion_create_params import CompletionCreateParamsStr
 from klaude_code.llm.client import LLMClientABC
 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_processor import StreamStateManager
+from klaude_code.llm.openai_compatible.stream import parse_chat_completions_stream
 from klaude_code.llm.openrouter.input import convert_history_to_input, is_claude_model
-from klaude_code.llm.openrouter.reasoning_handler import ReasoningDetail, ReasoningStreamHandler
+from klaude_code.llm.openrouter.reasoning import ReasoningStreamHandler
 from klaude_code.llm.registry import register
-from klaude_code.llm.usage import MetadataTracker, convert_usage
+from klaude_code.llm.usage import MetadataTracker
 from klaude_code.protocol import llm_param, model
-from klaude_code.trace import DebugType, is_debug_enabled, log, log_debug
+from klaude_code.trace import DebugType, is_debug_enabled, log_debug
 
 
 def build_payload(
@@ -96,114 +96,34 @@ class OpenRouterClient(LLMClientABC):
             debug_type=DebugType.LLM_PAYLOAD,
         )
 
-        stream = self.client.chat.completions.create(
-            **payload,
-            extra_body=extra_body,
-            extra_headers=extra_headers,
-        )
+        try:
+            stream = await self.client.chat.completions.create(
+                **payload,
+                extra_body=extra_body,
+                extra_headers=extra_headers,
+            )
+        except (openai.OpenAIError, httpx.HTTPError) as e:
+            yield model.StreamErrorItem(error=f"{e.__class__.__name__} {e!s}")
+            yield metadata_tracker.finalize()
+            return
 
         reasoning_handler = ReasoningStreamHandler(
             param_model=str(param.model),
             response_id=None,
         )
 
-        state = StreamStateManager(
-            param_model=str(param.model),
-            reasoning_flusher=reasoning_handler.flush,
-        )
-
-        try:
-            async for event in await stream:
-                log_debug(
-                    event.model_dump_json(exclude_none=True),
-                    style="blue",
-                    debug_type=DebugType.LLM_STREAM,
-                )
-
-                if not state.response_id and event.id:
-                    state.set_response_id(event.id)
-                    reasoning_handler.set_response_id(event.id)
-                    yield model.StartItem(response_id=event.id)
-                if event.usage is not None:
-                    metadata_tracker.set_usage(convert_usage(event.usage, param.context_limit, param.max_tokens))
-                if event.model:
-                    metadata_tracker.set_model_name(event.model)
-                if provider := getattr(event, "provider", None):
-                    metadata_tracker.set_provider(str(provider))
-                if len(event.choices) == 0:
-                    continue
-                delta = event.choices[0].delta
-
-                # Reasoning
-                if reasoning_details := getattr(delta, "reasoning_details", None):
-                    for item in reasoning_details:
-                        try:
-                            reasoning_detail = ReasoningDetail.model_validate(item)
-                            if reasoning_detail.text or reasoning_detail.summary:
-                                metadata_tracker.record_token()
-                            state.stage = "reasoning"
-                            # Yield delta immediately for streaming
-                            if reasoning_detail.text:
-                                yield model.ReasoningTextDelta(
-                                    content=reasoning_detail.text,
-                                    response_id=state.response_id,
-                                )
-                            if reasoning_detail.summary:
-                                yield model.ReasoningTextDelta(
-                                    content=reasoning_detail.summary,
-                                    response_id=state.response_id,
-                                )
-                            # Keep existing handler logic for final items
-                            for conversation_item in reasoning_handler.on_detail(reasoning_detail):
-                                yield conversation_item
-                        except Exception as e:
-                            log("reasoning_details error", str(e), style="red")
-
-                # Assistant
-                if delta.content and (
-                    state.stage == "assistant" or delta.content.strip()
-                ):  # Process all content in assistant stage, filter empty content in reasoning stage
-                    metadata_tracker.record_token()
-                    if state.stage == "reasoning":
-                        for item in state.flush_reasoning():
-                            yield item
-                    state.stage = "assistant"
-                    state.accumulated_content.append(delta.content)
-                    yield model.AssistantMessageDelta(
-                        content=delta.content,
-                        response_id=state.response_id,
-                    )
-
-                # Tool
-                if delta.tool_calls and len(delta.tool_calls) > 0:
-                    metadata_tracker.record_token()
-                    if state.stage == "reasoning":
-                        for item in state.flush_reasoning():
-                            yield item
-                    elif state.stage == "assistant":
-                        for item in state.flush_assistant():
-                            yield item
-                    state.stage = "tool"
-                    # Emit ToolCallStartItem for new tool calls
-                    for tc in delta.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)
-                            yield model.ToolCallStartItem(
-                                response_id=state.response_id,
-                                call_id=tc.id or "",
-                                name=tc.function.name,
-                            )
-                    state.accumulated_tool_calls.add(delta.tool_calls)
-
-        except (openai.OpenAIError, httpx.HTTPError) as e:
-            yield model.StreamErrorItem(error=f"{e.__class__.__name__} {e!s}")
-
-        # Finalize
-        flushed_items = state.flush_all()
-        if flushed_items:
-            metadata_tracker.record_token()
-        for item in flushed_items:
+        def on_event(event: Any) -> None:
+            log_debug(
+                event.model_dump_json(exclude_none=True),
+                style="blue",
+                debug_type=DebugType.LLM_STREAM,
+            )
+
+        async for item in parse_chat_completions_stream(
+            stream,
+            param=param,
+            metadata_tracker=metadata_tracker,
+            reasoning_handler=reasoning_handler,
+            on_event=on_event,
+        ):
             yield item
-
-        metadata_tracker.set_response_id(state.response_id)
-        yield metadata_tracker.finalize()

klaude_code/llm/openrouter/{reasoning_handler.py → reasoning.py}

@@ -1,6 +1,8 @@
 from pydantic import BaseModel
 
+from klaude_code.llm.openai_compatible.stream import ReasoningDeltaResult, ReasoningHandlerABC
 from klaude_code.protocol import model
+from klaude_code.trace import log
 
 
 class ReasoningDetail(BaseModel):
@@ -16,8 +18,8 @@ class ReasoningDetail(BaseModel):
     signature: str | None = None  # Claude's signature
 
 
-class ReasoningStreamHandler:
-    """Accumulates reasoning text and flushes on encrypted content or finalize."""
+class ReasoningStreamHandler(ReasoningHandlerABC):
+    """Accumulates OpenRouter reasoning details and emits ordered outputs."""
 
     def __init__(
         self,
@@ -34,6 +36,26 @@ class ReasoningStreamHandler:
         """Update the response identifier used for emitted items."""
         self._response_id = response_id
 
+    def on_delta(self, delta: object) -> ReasoningDeltaResult:
+        """Parse OpenRouter's reasoning_details and return ordered stream outputs."""
+        reasoning_details = getattr(delta, "reasoning_details", None)
+        if not reasoning_details:
+            return ReasoningDeltaResult(handled=False, outputs=[])
+
+        outputs: list[str | model.ConversationItem] = []
+        for item in reasoning_details:
+            try:
+                reasoning_detail = ReasoningDetail.model_validate(item)
+                if reasoning_detail.text:
+                    outputs.append(reasoning_detail.text)
+                if reasoning_detail.summary:
+                    outputs.append(reasoning_detail.summary)
+                outputs.extend(self.on_detail(reasoning_detail))
+            except Exception as e:
+                log("reasoning_details error", str(e), style="red")
+
+        return ReasoningDeltaResult(handled=True, outputs=outputs)
+
     def on_detail(self, detail: ReasoningDetail) -> list[model.ConversationItem]:
         """Process a single reasoning detail and return streamable items."""
         items: list[model.ConversationItem] = []

klaude_code/protocol/model.py

@@ -138,6 +138,12 @@ class TruncationUIExtra(BaseModel):
     truncated_length: int
 
 
+class MarkdownDocUIExtra(BaseModel):
+    type: Literal["markdown_doc"] = "markdown_doc"
+    file_path: str
+    content: str
+
+
 class SessionStatusUIExtra(BaseModel):
     type: Literal["session_status"] = "session_status"
     usage: "Usage"
@@ -146,7 +152,13 @@ class SessionStatusUIExtra(BaseModel):
 
 
 ToolResultUIExtra = Annotated[
-    DiffUIExtra | TodoListUIExtra | SessionIdUIExtra | MermaidLinkUIExtra | TruncationUIExtra | SessionStatusUIExtra,
+    DiffUIExtra
+    | TodoListUIExtra
+    | SessionIdUIExtra
+    | MermaidLinkUIExtra
+    | TruncationUIExtra
+    | MarkdownDocUIExtra
+    | SessionStatusUIExtra,
     Field(discriminator="type"),
 ]
 

klaude_code/protocol/op.py

@@ -25,6 +25,7 @@ class OperationType(Enum):
     USER_INPUT = "user_input"
     RUN_AGENT = "run_agent"
     CHANGE_MODEL = "change_model"
+    CHANGE_THINKING = "change_thinking"
     CLEAR_SESSION = "clear_session"
     EXPORT_SESSION = "export_session"
     INTERRUPT = "interrupt"
@@ -77,6 +78,16 @@ class ChangeModelOperation(Operation):
         await handler.handle_change_model(self)
 
 
+class ChangeThinkingOperation(Operation):
+    """Operation for changing the thinking/reasoning configuration."""
+
+    type: OperationType = OperationType.CHANGE_THINKING
+    session_id: str
+
+    async def execute(self, handler: OperationHandler) -> None:
+        await handler.handle_change_thinking(self)
+
+
 class ClearSessionOperation(Operation):
     """Operation for clearing the active session and starting a new one."""
 

klaude_code/protocol/op_handler.py

@@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Protocol
 if TYPE_CHECKING:
     from klaude_code.protocol.op import (
         ChangeModelOperation,
+        ChangeThinkingOperation,
         ClearSessionOperation,
         ExportSessionOperation,
         InitAgentOperation,
@@ -35,6 +36,10 @@ class OperationHandler(Protocol):
         """Handle a change model operation."""
         ...
 
+    async def handle_change_thinking(self, operation: ChangeThinkingOperation) -> None:
+        """Handle a change thinking operation."""
+        ...
+
     async def handle_clear_session(self, operation: ClearSessionOperation) -> None:
         """Handle a clear session operation."""
         ...

klaude_code/ui/core/stage_manager.py

@@ -20,12 +20,10 @@ class StageManager:
         *,
         finish_assistant: Callable[[], Awaitable[None]],
         finish_thinking: Callable[[], Awaitable[None]],
-        on_enter_thinking: Callable[[], None],
     ):
         self._stage = Stage.WAITING
         self._finish_assistant = finish_assistant
         self._finish_thinking = finish_thinking
-        self._on_enter_thinking = on_enter_thinking
 
     @property
     def current_stage(self) -> Stage:
@@ -41,7 +39,6 @@ class StageManager:
         if self._stage == Stage.THINKING:
             return
         await self.transition_to(Stage.THINKING)
-        self._on_enter_thinking()
 
     async def finish_assistant(self) -> None:
         if self._stage != Stage.ASSISTANT:

klaude_code/ui/modes/repl/display.py

@@ -57,3 +57,5 @@ class REPLDisplay(DisplayABC):
         # Spinner may already be stopped or not started; ignore.
         with contextlib.suppress(Exception):
             self.renderer.spinner_stop()
+        with contextlib.suppress(Exception):
+            self.renderer.stop_bottom_live()