klaude-code 1.2.8__py3-none-any.whl → 1.2.9__py3-none-any.whl

klaude_code/llm/anthropic/client.py

@@ -1,5 +1,4 @@
 import json
-import time
 from collections.abc import AsyncGenerator
 from typing import override
 
@@ -22,7 +21,7 @@ from klaude_code.llm.anthropic.input import convert_history_to_input, convert_sy
 from klaude_code.llm.client import LLMClientABC, call_with_logged_payload
 from klaude_code.llm.input_common import apply_config_defaults
 from klaude_code.llm.registry import register
-from klaude_code.llm.usage import calculate_cost
+from klaude_code.llm.usage import MetadataTracker, convert_anthropic_usage
 from klaude_code.protocol import llm_param, model
 from klaude_code.trace import DebugType, log_debug
 
@@ -47,9 +46,7 @@ class AnthropicClient(LLMClientABC):
     async def call(self, param: llm_param.LLMCallParameter) -> AsyncGenerator[model.ConversationItem, None]:
         param = apply_config_defaults(param, self.get_llm_config())
 
-        request_start_time = time.time()
-        first_token_time: float | None = None
-        last_token_time: float | None = None
+        metadata_tracker = MetadataTracker(cost_config=self._config.cost)
 
         messages = convert_history_to_input(param.input, param.model)
         tools = convert_tool_schema(param.tools)
@@ -77,7 +74,7 @@ class AnthropicClient(LLMClientABC):
             else anthropic.types.ThinkingConfigDisabledParam(
                 type="disabled",
             ),
-            extra_headers={"extra": json.dumps({"session_id": param.session_id})},
+            extra_headers={"extra": json.dumps({"session_id": param.session_id}, sort_keys=True)},
         )
 
         accumulated_thinking: list[str] = []
@@ -112,32 +109,24 @@ class AnthropicClient(LLMClientABC):
                     case BetaRawContentBlockDeltaEvent() as event:
                         match event.delta:
                             case BetaThinkingDelta() as delta:
-                                if first_token_time is None:
-                                    first_token_time = time.time()
-                                last_token_time = time.time()
+                                metadata_tracker.record_token()
                                 accumulated_thinking.append(delta.thinking)
                             case BetaSignatureDelta() as delta:
-                                if first_token_time is None:
-                                    first_token_time = time.time()
-                                last_token_time = time.time()
+                                metadata_tracker.record_token()
                                 yield model.ReasoningEncryptedItem(
                                     encrypted_content=delta.signature,
                                     response_id=response_id,
                                     model=str(param.model),
                                 )
                             case BetaTextDelta() as delta:
-                                if first_token_time is None:
-                                    first_token_time = time.time()
-                                last_token_time = time.time()
+                                metadata_tracker.record_token()
                                 accumulated_content.append(delta.text)
                                 yield model.AssistantMessageDelta(
                                     content=delta.text,
                                     response_id=response_id,
                                 )
                             case BetaInputJSONDelta() as delta:
-                                if first_token_time is None:
-                                    first_token_time = time.time()
-                                last_token_time = time.time()
+                                metadata_tracker.record_token()
                                 if current_tool_inputs is not None:
                                     current_tool_inputs.append(delta.partial_json)
                             case _:
@@ -184,37 +173,17 @@ class AnthropicClient(LLMClientABC):
                         input_tokens += (event.usage.input_tokens or 0) + (event.usage.cache_creation_input_tokens or 0)
                         output_tokens += event.usage.output_tokens or 0
                         cached_tokens += event.usage.cache_read_input_tokens or 0
-                        total_tokens = input_tokens + cached_tokens + output_tokens
-                        context_usage_percent = (
-                            (total_tokens / param.context_limit) * 100 if param.context_limit else None
-                        )
-
-                        throughput_tps: float | None = None
-                        first_token_latency_ms: float | None = None
-
-                        if first_token_time is not None:
-                            first_token_latency_ms = (first_token_time - request_start_time) * 1000
 
-                        if first_token_time is not None and last_token_time is not None and output_tokens > 0:
-                            time_duration = last_token_time - first_token_time
-                            if time_duration >= 0.15:
-                                throughput_tps = output_tokens / time_duration
-
-                        usage = model.Usage(
+                        usage = convert_anthropic_usage(
                             input_tokens=input_tokens,
                             output_tokens=output_tokens,
                             cached_tokens=cached_tokens,
-                            total_tokens=total_tokens,
-                            context_usage_percent=context_usage_percent,
-                            throughput_tps=throughput_tps,
-                            first_token_latency_ms=first_token_latency_ms,
-                        )
-                        calculate_cost(usage, self._config.cost)
-                        yield model.ResponseMetadataItem(
-                            usage=usage,
-                            response_id=response_id,
-                            model_name=str(param.model),
+                            context_limit=param.context_limit,
                         )
+                        metadata_tracker.set_usage(usage)
+                        metadata_tracker.set_model_name(str(param.model))
+                        metadata_tracker.set_response_id(response_id)
+                        yield metadata_tracker.finalize()
                     case _:
                         pass
         except (APIError, httpx.HTTPError) as e:

klaude_code/llm/client.py

@@ -42,7 +42,7 @@ def call_with_logged_payload(func: Callable[P, R], *args: P.args, **kwargs: P.kw
 
     payload = {k: v for k, v in kwargs.items() if v is not None}
     log_debug(
-        json.dumps(payload, ensure_ascii=False, default=str),
+        json.dumps(payload, ensure_ascii=False, default=str, sort_keys=True),
         style="yellow",
         debug_type=DebugType.LLM_PAYLOAD,
     )

klaude_code/llm/codex/client.py

@@ -1,6 +1,5 @@
 """Codex LLM client using ChatGPT subscription via OAuth."""
 
-import time
 from collections.abc import AsyncGenerator
 from typing import override
 
@@ -16,6 +15,7 @@ from klaude_code.llm.input_common import apply_config_defaults
 from klaude_code.llm.registry import register
 from klaude_code.llm.responses.client import parse_responses_stream
 from klaude_code.llm.responses.input import convert_history_to_input, convert_tool_schema
+from klaude_code.llm.usage import MetadataTracker
 from klaude_code.protocol import llm_param, model
 
 # Codex API configuration
@@ -24,6 +24,7 @@ CODEX_HEADERS = {
     "originator": "codex_cli_rs",
     # Mocked Codex-style user agent string
     "User-Agent": "codex_cli_rs/0.0.0-klaude",
+    "OpenAI-Beta": "responses=experimental",
 }
 
 
@@ -83,7 +84,7 @@ class CodexClient(LLMClientABC):
         # Codex API requires store=False
         param.store = False
 
-        request_start_time = time.time()
+        metadata_tracker = MetadataTracker(cost_config=self._config.cost)
 
         inputs = convert_history_to_input(param.input, param.model)
         tools = convert_tool_schema(param.tools)
@@ -125,5 +126,5 @@ class CodexClient(LLMClientABC):
             yield model.StreamErrorItem(error=f"{e.__class__.__name__} {str(e)}")
             return
 
-        async for item in parse_responses_stream(stream, param, self._config.cost, request_start_time):
+        async for item in parse_responses_stream(stream, param, metadata_tracker):
             yield item

klaude_code/llm/input_common.py

@@ -49,10 +49,6 @@ class AssistantGroup:
 
     text_content: str | None = None
     tool_calls: list[model.ToolCallItem] = field(default_factory=lambda: [])
-    reasoning_text: list[model.ReasoningTextItem] = field(default_factory=lambda: [])
-    reasoning_encrypted: list[model.ReasoningEncryptedItem] = field(default_factory=lambda: [])
-    # Preserve original ordering of reasoning items for providers that
-    # need to emit them as an ordered stream (e.g. OpenRouter).
     reasoning_items: list[model.ReasoningTextItem | model.ReasoningEncryptedItem] = field(default_factory=lambda: [])
 
 
@@ -184,10 +180,8 @@ def parse_message_groups(history: list[model.ConversationItem]) -> list[MessageG
                         case model.ToolCallItem():
                             group.tool_calls.append(item)
                         case model.ReasoningTextItem():
-                            group.reasoning_text.append(item)
                             group.reasoning_items.append(item)
                         case model.ReasoningEncryptedItem():
-                            group.reasoning_encrypted.append(item)
                             group.reasoning_items.append(item)
                         case _:
                             pass

klaude_code/llm/openai_compatible/client.py

@@ -1,15 +1,14 @@
 import json
 from collections.abc import AsyncGenerator
-from typing import Literal, override
+from typing import override
 
 import httpx
 import openai
 
-
 from klaude_code.llm.client import LLMClientABC, call_with_logged_payload
 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.tool_call_accumulator import BasicToolCallAccumulator, ToolCallAccumulatorABC
+from klaude_code.llm.openai_compatible.stream_processor import StreamStateManager
 from klaude_code.llm.registry import register
 from klaude_code.llm.usage import MetadataTracker, convert_usage
 from klaude_code.protocol import llm_param, model
@@ -51,7 +50,7 @@ class OpenAICompatibleClient(LLMClientABC):
         metadata_tracker = MetadataTracker(cost_config=self._config.cost)
 
         extra_body = {}
-        extra_headers = {"extra": json.dumps({"session_id": param.session_id})}
+        extra_headers = {"extra": json.dumps({"session_id": param.session_id}, sort_keys=True)}
 
         if param.thinking:
             extra_body["thinking"] = {
@@ -74,42 +73,7 @@ class OpenAICompatibleClient(LLMClientABC):
             extra_headers=extra_headers,
         )
 
-        stage: Literal["waiting", "reasoning", "assistant", "tool", "done"] = "waiting"
-        accumulated_reasoning: list[str] = []
-        accumulated_content: list[str] = []
-        accumulated_tool_calls: ToolCallAccumulatorABC = BasicToolCallAccumulator()
-        emitted_tool_start_indices: set[int] = set()
-        response_id: str | None = None
-
-        def flush_reasoning_items() -> list[model.ConversationItem]:
-            nonlocal accumulated_reasoning
-            if not accumulated_reasoning:
-                return []
-            item = model.ReasoningTextItem(
-                content="".join(accumulated_reasoning),
-                response_id=response_id,
-                model=str(param.model),
-            )
-            accumulated_reasoning = []
-            return [item]
-
-        def flush_assistant_items() -> list[model.ConversationItem]:
-            nonlocal accumulated_content
-            if len(accumulated_content) == 0:
-                return []
-            item = model.AssistantMessageItem(
-                content="".join(accumulated_content),
-                response_id=response_id,
-            )
-            accumulated_content = []
-            return [item]
-
-        def flush_tool_call_items() -> list[model.ToolCallItem]:
-            nonlocal accumulated_tool_calls
-            items: list[model.ToolCallItem] = accumulated_tool_calls.get()
-            if items:
-                accumulated_tool_calls.chunks_by_step = []  # pyright: ignore[reportAttributeAccessIssue]
-            return items
+        state = StreamStateManager(param_model=str(param.model))
 
         try:
             async for event in await stream:
@@ -118,10 +82,9 @@ class OpenAICompatibleClient(LLMClientABC):
                     style="blue",
                     debug_type=DebugType.LLM_STREAM,
                 )
-                if not response_id and event.id:
-                    response_id = event.id
-                    accumulated_tool_calls.response_id = response_id
-                    yield model.StartItem(response_id=response_id)
+                if not state.response_id and event.id:
+                    state.set_response_id(event.id)
+                    yield model.StartItem(response_id=event.id)
                 if (
                     event.usage is not None and event.usage.completion_tokens is not None  # pyright: ignore[reportUnnecessaryComparison] gcp gemini will return None usage field
                 ):
@@ -152,60 +115,53 @@ class OpenAICompatibleClient(LLMClientABC):
                     reasoning_content = getattr(delta, "reasoning_content")
                 if reasoning_content:
                     metadata_tracker.record_token()
-                    stage = "reasoning"
-                    accumulated_reasoning.append(reasoning_content)
+                    state.stage = "reasoning"
+                    state.accumulated_reasoning.append(reasoning_content)
 
                 # Assistant
                 if delta.content and (
-                    stage == "assistant" or delta.content.strip()
+                    state.stage == "assistant" or delta.content.strip()
                 ):  # Process all content in assistant stage, filter empty content in reasoning stage
                     metadata_tracker.record_token()
-                    if stage == "reasoning":
-                        for item in flush_reasoning_items():
+                    if state.stage == "reasoning":
+                        for item in state.flush_reasoning():
                             yield item
-                    elif stage == "tool":
-                        for item in flush_tool_call_items():
+                    elif state.stage == "tool":
+                        for item in state.flush_tool_calls():
                             yield item
-                    stage = "assistant"
-                    accumulated_content.append(delta.content)
+                    state.stage = "assistant"
+                    state.accumulated_content.append(delta.content)
                     yield model.AssistantMessageDelta(
                         content=delta.content,
-                        response_id=response_id,
+                        response_id=state.response_id,
                     )
 
                 # Tool
                 if delta.tool_calls and len(delta.tool_calls) > 0:
                     metadata_tracker.record_token()
-                    if stage == "reasoning":
-                        for item in flush_reasoning_items():
+                    if state.stage == "reasoning":
+                        for item in state.flush_reasoning():
                             yield item
-                    elif stage == "assistant":
-                        for item in flush_assistant_items():
+                    elif state.stage == "assistant":
+                        for item in state.flush_assistant():
                             yield item
-                    stage = "tool"
+                    state.stage = "tool"
                     # Emit ToolCallStartItem for new tool calls
                     for tc in delta.tool_calls:
-                        if tc.index not in emitted_tool_start_indices and tc.function and tc.function.name:
-                            emitted_tool_start_indices.add(tc.index)
+                        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=response_id,
+                                response_id=state.response_id,
                                 call_id=tc.id or "",
                                 name=tc.function.name,
                             )
-                    accumulated_tool_calls.add(delta.tool_calls)
+                    state.accumulated_tool_calls.add(delta.tool_calls)
         except (openai.OpenAIError, httpx.HTTPError) as e:
             yield model.StreamErrorItem(error=f"{e.__class__.__name__} {str(e)}")
 
         # Finalize
-        for item in flush_reasoning_items():
+        for item in state.flush_all():
             yield item
 
-        for item in flush_assistant_items():
-            yield item
-
-        if stage == "tool":
-            for tool_call_item in flush_tool_call_items():
-                yield tool_call_item
-
-        metadata_tracker.set_response_id(response_id)
+        metadata_tracker.set_response_id(state.response_id)
         yield metadata_tracker.finalize()

klaude_code/llm/openai_compatible/input.py

@@ -10,7 +10,8 @@ from klaude_code.llm.input_common import AssistantGroup, ToolGroup, UserGroup, m
 from klaude_code.protocol import llm_param, model
 
 
-def _user_group_to_message(group: UserGroup) -> chat.ChatCompletionMessageParam:
+def user_group_to_openai_message(group: UserGroup) -> chat.ChatCompletionMessageParam:
+    """Convert a UserGroup to an OpenAI-compatible chat message."""
     parts: list[ChatCompletionContentPartParam] = []
     for text in group.text_parts:
         parts.append({"type": "text", "text": text + "\n"})
@@ -21,7 +22,8 @@ def _user_group_to_message(group: UserGroup) -> chat.ChatCompletionMessageParam:
     return {"role": "user", "content": parts}
 
 
-def _tool_group_to_message(group: ToolGroup) -> chat.ChatCompletionMessageParam:
+def tool_group_to_openai_message(group: ToolGroup) -> chat.ChatCompletionMessageParam:
+    """Convert a ToolGroup to an OpenAI-compatible chat message."""
     merged_text = merge_reminder_text(
         group.tool_result.output or "<system-reminder>Tool ran without output or errors</system-reminder>",
         group.reminder_texts,
@@ -82,9 +84,9 @@ def convert_history_to_input(
     for group in parse_message_groups(history):
         match group:
             case UserGroup():
-                messages.append(_user_group_to_message(group))
+                messages.append(user_group_to_openai_message(group))
             case ToolGroup():
-                messages.append(_tool_group_to_message(group))
+                messages.append(tool_group_to_openai_message(group))
             case AssistantGroup():
                 messages.append(_assistant_group_to_message(group))
 

klaude_code/llm/openai_compatible/stream_processor.py

@@ -0,0 +1,82 @@
+"""Shared stream processing utilities for OpenAI-compatible clients.
+
+This module provides a reusable stream state manager that handles the common
+logic for accumulating and flushing reasoning, assistant content, and tool calls
+across different LLM providers (OpenAI-compatible, OpenRouter).
+"""
+
+from typing import Callable, Literal
+
+from klaude_code.llm.openai_compatible.tool_call_accumulator import BasicToolCallAccumulator, ToolCallAccumulatorABC
+from klaude_code.protocol import 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

klaude_code/llm/openrouter/client.py

@@ -1,5 +1,5 @@
 from collections.abc import AsyncGenerator
-from typing import Literal, override
+from typing import override
 
 import httpx
 import openai
@@ -7,7 +7,7 @@ import openai
 from klaude_code.llm.client import LLMClientABC, call_with_logged_payload
 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.tool_call_accumulator import BasicToolCallAccumulator, ToolCallAccumulatorABC
+from klaude_code.llm.openai_compatible.stream_processor import StreamStateManager
 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.registry import register
@@ -77,36 +77,14 @@ class OpenRouterClient(LLMClientABC):
             extra_headers=extra_headers,  # pyright: ignore[reportUnknownArgumentType]
         )
 
-        stage: Literal["waiting", "reasoning", "assistant", "tool", "done"] = "waiting"
-        response_id: str | None = None
-        accumulated_content: list[str] = []
-        accumulated_tool_calls: ToolCallAccumulatorABC = BasicToolCallAccumulator()
-        emitted_tool_start_indices: set[int] = set()
         reasoning_handler = ReasoningStreamHandler(
             param_model=str(param.model),
-            response_id=response_id,
+            response_id=None,
+        )
+        state = StreamStateManager(
+            param_model=str(param.model),
+            reasoning_flusher=reasoning_handler.flush,
         )
-
-        def flush_reasoning_items() -> list[model.ConversationItem]:
-            return reasoning_handler.flush()
-
-        def flush_assistant_items() -> list[model.ConversationItem]:
-            nonlocal accumulated_content
-            if len(accumulated_content) == 0:
-                return []
-            item = model.AssistantMessageItem(
-                content="".join(accumulated_content),
-                response_id=response_id,
-            )
-            accumulated_content = []
-            return [item]
-
-        def flush_tool_call_items() -> list[model.ToolCallItem]:
-            nonlocal accumulated_tool_calls
-            items: list[model.ToolCallItem] = accumulated_tool_calls.get()
-            if items:
-                accumulated_tool_calls.chunks_by_step = []  # pyright: ignore[reportAttributeAccessIssue]
-            return items
 
         try:
             async for event in await stream:
@@ -115,11 +93,10 @@ class OpenRouterClient(LLMClientABC):
                     style="blue",
                     debug_type=DebugType.LLM_STREAM,
                 )
-                if not response_id and event.id:
-                    response_id = event.id
-                    reasoning_handler.set_response_id(response_id)
-                    accumulated_tool_calls.response_id = response_id
-                    yield model.StartItem(response_id=response_id)
+                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 and event.usage.completion_tokens is not None  # pyright: ignore[reportUnnecessaryComparison]
                 ):  # gcp gemini will return None usage field
@@ -140,7 +117,7 @@ class OpenRouterClient(LLMClientABC):
                         try:
                             reasoning_detail = ReasoningDetail.model_validate(item)
                             metadata_tracker.record_token()
-                            stage = "reasoning"
+                            state.stage = "reasoning"
                             for conversation_item in reasoning_handler.on_detail(reasoning_detail):
                                 yield conversation_item
                         except Exception as e:
@@ -148,53 +125,46 @@ class OpenRouterClient(LLMClientABC):
 
                 # Assistant
                 if delta.content and (
-                    stage == "assistant" or delta.content.strip()
+                    state.stage == "assistant" or delta.content.strip()
                 ):  # Process all content in assistant stage, filter empty content in reasoning stage
                     metadata_tracker.record_token()
-                    if stage == "reasoning":
-                        for item in flush_reasoning_items():
+                    if state.stage == "reasoning":
+                        for item in state.flush_reasoning():
                             yield item
-                    stage = "assistant"
-                    accumulated_content.append(delta.content)
+                    state.stage = "assistant"
+                    state.accumulated_content.append(delta.content)
                     yield model.AssistantMessageDelta(
                         content=delta.content,
-                        response_id=response_id,
+                        response_id=state.response_id,
                     )
 
                 # Tool
                 if delta.tool_calls and len(delta.tool_calls) > 0:
                     metadata_tracker.record_token()
-                    if stage == "reasoning":
-                        for item in flush_reasoning_items():
+                    if state.stage == "reasoning":
+                        for item in state.flush_reasoning():
                             yield item
-                    elif stage == "assistant":
-                        for item in flush_assistant_items():
+                    elif state.stage == "assistant":
+                        for item in state.flush_assistant():
                             yield item
-                    stage = "tool"
+                    state.stage = "tool"
                     # Emit ToolCallStartItem for new tool calls
                     for tc in delta.tool_calls:
-                        if tc.index not in emitted_tool_start_indices and tc.function and tc.function.name:
-                            emitted_tool_start_indices.add(tc.index)
+                        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=response_id,
+                                response_id=state.response_id,
                                 call_id=tc.id or "",
                                 name=tc.function.name,
                             )
-                    accumulated_tool_calls.add(delta.tool_calls)
+                    state.accumulated_tool_calls.add(delta.tool_calls)
 
         except (openai.OpenAIError, httpx.HTTPError) as e:
             yield model.StreamErrorItem(error=f"{e.__class__.__name__} {str(e)}")
 
         # Finalize
-        for item in flush_reasoning_items():
-            yield item
-
-        for item in flush_assistant_items():
+        for item in state.flush_all():
             yield item
 
-        if stage == "tool":
-            for tool_call_item in flush_tool_call_items():
-                yield tool_call_item
-
-        metadata_tracker.set_response_id(response_id)
+        metadata_tracker.set_response_id(state.response_id)
         yield metadata_tracker.finalize()