klaude-code 1.2.6__py3-none-any.whl → 1.8.0__py3-none-any.whl

klaude_code/core/turn.py

@@ -1,9 +1,14 @@
 from __future__ import annotations
 
-from collections.abc import AsyncGenerator, Callable, MutableMapping, Sequence
-from dataclasses import dataclass
+from collections.abc import AsyncGenerator
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+from klaude_code.core.tool import ToolABC, tool_context
+
+if TYPE_CHECKING:
+    from klaude_code.core.task import SessionContext
 
-from klaude_code.core.tool import TodoContext, ToolABC, tool_context
 from klaude_code.core.tool.tool_runner import (
     ToolExecutionCallStarted,
     ToolExecutionResult,
@@ -12,7 +17,7 @@ from klaude_code.core.tool.tool_runner import (
     ToolExecutorEvent,
 )
 from klaude_code.llm import LLMClientABC
-from klaude_code.protocol import events, llm_param, model
+from klaude_code.protocol import events, llm_param, model, tools
 from klaude_code.trace import DebugType, log_debug
 
 
@@ -26,16 +31,22 @@ class TurnError(Exception):
 class TurnExecutionContext:
     """Execution context required to run a single turn."""
 
-    session_id: str
-    get_conversation_history: Callable[[], list[model.ConversationItem]]
-    append_history: Callable[[Sequence[model.ConversationItem]], None]
+    session_ctx: SessionContext
     llm_client: LLMClientABC
     system_prompt: str | None
     tools: list[llm_param.ToolSchema]
     tool_registry: dict[str, type[ToolABC]]
-    # For tool context
-    file_tracker: MutableMapping[str, float]
-    todo_context: TodoContext
+
+
+@dataclass
+class TurnResult:
+    """Aggregated state produced while executing a turn."""
+
+    reasoning_items: list[model.ReasoningTextItem | model.ReasoningEncryptedItem]
+    assistant_message: model.AssistantMessageItem | None
+    tool_calls: list[model.ToolCallItem]
+    stream_error: model.StreamErrorItem | None
+    report_back_result: str | None = field(default=None)
 
 
 def build_events_from_tool_executor_event(session_id: str, event: ToolExecutorEvent) -> list[events.Event]:
@@ -64,6 +75,7 @@ def build_events_from_tool_executor_event(session_id: str, event: ToolExecutorEv
                     result=tool_result.output or "",
                     ui_extra=tool_result.ui_extra,
                     status=tool_result.status,
+                    task_metadata=tool_result.task_metadata,
                 )
             )
         case ToolExecutionTodoChange(todos=todos):
@@ -87,45 +99,110 @@ class TurnExecutor:
     def __init__(self, context: TurnExecutionContext) -> None:
         self._context = context
         self._tool_executor: ToolExecutor | None = None
-        self._has_tool_call: bool = False
+        self._turn_result: TurnResult | None = None
+        self._assistant_delta_buffer: list[str] = []
+        self._assistant_response_id: str | None = None
+
+    @property
+    def report_back_result(self) -> str | None:
+        return self._turn_result.report_back_result if self._turn_result else None
+
+    @property
+    def task_finished(self) -> bool:
+        """Check if this turn indicates the task should end.
+
+        Task ends when there are no tool calls or report_back was called.
+        """
+        if self._turn_result is None:
+            return True
+        if not self._turn_result.tool_calls:
+            return True
+        return self._turn_result.report_back_result is not None
 
     @property
-    def has_tool_call(self) -> bool:
-        return self._has_tool_call
+    def task_result(self) -> str:
+        """Get the task result from this turn.
+
+        Returns report_back result if available, otherwise returns
+        the assistant message content.
+        """
+        if self._turn_result is not None and self._turn_result.report_back_result is not None:
+            return self._turn_result.report_back_result
+        if self._turn_result is not None and self._turn_result.assistant_message is not None:
+            return self._turn_result.assistant_message.content or ""
+        return ""
+
+    @property
+    def has_structured_output(self) -> bool:
+        """Check if the task result is structured output from report_back."""
+        return bool(self._turn_result and self._turn_result.report_back_result)
 
     def cancel(self) -> list[events.Event]:
         """Cancel running tools and return any resulting events."""
         ui_events: list[events.Event] = []
+        self._persist_partial_assistant_on_cancel()
         if self._tool_executor is not None:
             for exec_event in self._tool_executor.cancel():
-                for ui_event in build_events_from_tool_executor_event(self._context.session_id, exec_event):
+                for ui_event in build_events_from_tool_executor_event(self._context.session_ctx.session_id, exec_event):
                     ui_events.append(ui_event)
             self._tool_executor = None
         return ui_events
 
-    async def run(self) -> AsyncGenerator[events.Event, None]:
+    async def run(self) -> AsyncGenerator[events.Event]:
         """Execute the turn, yielding events as they occur.
 
         Raises:
             TurnError: If the turn fails (stream error or non-completed status).
         """
         ctx = self._context
+        session_ctx = ctx.session_ctx
+
+        yield events.TurnStartEvent(session_id=session_ctx.session_id)
+
+        self._turn_result = TurnResult(
+            reasoning_items=[],
+            assistant_message=None,
+            tool_calls=[],
+            stream_error=None,
+        )
 
-        yield events.TurnStartEvent(session_id=ctx.session_id)
+        async for event in self._consume_llm_stream(self._turn_result):
+            yield event
 
-        turn_reasoning_items: list[model.ReasoningTextItem | model.ReasoningEncryptedItem] = []
-        turn_assistant_message: model.AssistantMessageItem | None = None
-        turn_tool_calls: list[model.ToolCallItem] = []
-        response_failed = False
-        error_message: str | None = None
+        if self._turn_result.stream_error is not None:
+            session_ctx.append_history([self._turn_result.stream_error])
+            yield events.TurnEndEvent(session_id=session_ctx.session_id)
+            raise TurnError(self._turn_result.stream_error.error)
 
+        self._append_success_history(self._turn_result)
+
+        if self._turn_result.tool_calls:
+            # Check for report_back before running tools
+            self._detect_report_back(self._turn_result)
+
+            async for ui_event in self._run_tool_executor(self._turn_result.tool_calls):
+                yield ui_event
+
+        yield events.TurnEndEvent(session_id=session_ctx.session_id)
+
+    def _detect_report_back(self, turn_result: TurnResult) -> None:
+        """Detect report_back tool call and store its arguments as JSON string."""
+        for tool_call in turn_result.tool_calls:
+            if tool_call.name == tools.REPORT_BACK:
+                turn_result.report_back_result = tool_call.arguments
+                break
+
+    async def _consume_llm_stream(self, turn_result: TurnResult) -> AsyncGenerator[events.Event]:
+        """Stream events from LLM and update turn_result in place."""
+
+        ctx = self._context
+        session_ctx = ctx.session_ctx
         async for response_item in ctx.llm_client.call(
             llm_param.LLMCallParameter(
-                input=ctx.get_conversation_history(),
+                input=session_ctx.get_conversation_history(),
                 system=ctx.system_prompt,
                 tools=ctx.tools,
-                store=False,
-                session_id=ctx.session_id,
+                session_id=session_ctx.session_id,
             )
         ):
             log_debug(
@@ -136,41 +213,45 @@ class TurnExecutor:
             )
             match response_item:
                 case model.StartItem():
-                    pass
+                    continue
                 case model.ReasoningTextItem() as item:
-                    turn_reasoning_items.append(item)
+                    turn_result.reasoning_items.append(item)
                     yield events.ThinkingEvent(
                         content=item.content,
                         response_id=item.response_id,
-                        session_id=ctx.session_id,
+                        session_id=session_ctx.session_id,
                     )
                 case model.ReasoningEncryptedItem() as item:
-                    turn_reasoning_items.append(item)
+                    turn_result.reasoning_items.append(item)
+                case model.ReasoningTextDelta() as item:
+                    yield events.ThinkingDeltaEvent(
+                        content=item.content,
+                        response_id=item.response_id,
+                        session_id=session_ctx.session_id,
+                    )
                 case model.AssistantMessageDelta() as item:
+                    if item.response_id:
+                        self._assistant_response_id = item.response_id
+                    self._assistant_delta_buffer.append(item.content)
                     yield events.AssistantMessageDeltaEvent(
                         content=item.content,
                         response_id=item.response_id,
-                        session_id=ctx.session_id,
+                        session_id=session_ctx.session_id,
                     )
                 case model.AssistantMessageItem() as item:
-                    turn_assistant_message = item
+                    turn_result.assistant_message = item
                     yield events.AssistantMessageEvent(
                         content=item.content or "",
                         response_id=item.response_id,
-                        session_id=ctx.session_id,
+                        session_id=session_ctx.session_id,
                     )
                 case model.ResponseMetadataItem() as item:
                     yield events.ResponseMetadataEvent(
-                        session_id=ctx.session_id,
+                        session_id=session_ctx.session_id,
                         metadata=item,
                     )
-                    status = item.status
-                    if status is not None and status != "completed":
-                        response_failed = True
-                        error_message = f"Response status: {status}"
                 case model.StreamErrorItem() as item:
-                    response_failed = True
-                    error_message = item.error
+                    turn_result.stream_error = item
                     log_debug(
                         "[StreamError]",
                         item.error,
@@ -179,42 +260,63 @@ class TurnExecutor:
                     )
                 case model.ToolCallStartItem() as item:
                     yield events.TurnToolCallStartEvent(
-                        session_id=ctx.session_id,
+                        session_id=session_ctx.session_id,
                         response_id=item.response_id,
                         tool_call_id=item.call_id,
                         tool_name=item.name,
                         arguments="",
                     )
                 case model.ToolCallItem() as item:
-                    turn_tool_calls.append(item)
+                    turn_result.tool_calls.append(item)
                 case _:
-                    pass
-
-        if response_failed:
-            yield events.TurnEndEvent(session_id=ctx.session_id)
-            raise TurnError(error_message or "Turn failed")
-
-        # Append to history only on success
-        if turn_reasoning_items:
-            ctx.append_history(turn_reasoning_items)
-        if turn_assistant_message:
-            ctx.append_history([turn_assistant_message])
-        if turn_tool_calls:
-            ctx.append_history(turn_tool_calls)
-            self._has_tool_call = True
-
-        # Execute tools
-        if turn_tool_calls:
-            with tool_context(ctx.file_tracker, ctx.todo_context):
-                executor = ToolExecutor(
-                    registry=ctx.tool_registry,
-                    append_history=ctx.append_history,
-                )
-                self._tool_executor = executor
+                    continue
+
+    def _append_success_history(self, turn_result: TurnResult) -> None:
+        """Persist successful turn artifacts to conversation history."""
+        session_ctx = self._context.session_ctx
+        if turn_result.reasoning_items:
+            session_ctx.append_history(turn_result.reasoning_items)
+        if turn_result.assistant_message:
+            session_ctx.append_history([turn_result.assistant_message])
+        if turn_result.tool_calls:
+            session_ctx.append_history(turn_result.tool_calls)
+        self._assistant_delta_buffer.clear()
+        self._assistant_response_id = None
 
-                async for exec_event in executor.run_tools(turn_tool_calls):
-                    for ui_event in build_events_from_tool_executor_event(ctx.session_id, exec_event):
+    async def _run_tool_executor(self, tool_calls: list[model.ToolCallItem]) -> AsyncGenerator[events.Event]:
+        """Run tools for the turn and translate executor events to UI events."""
+
+        ctx = self._context
+        session_ctx = ctx.session_ctx
+        with tool_context(session_ctx.file_tracker, session_ctx.todo_context):
+            executor = ToolExecutor(
+                registry=ctx.tool_registry,
+                append_history=session_ctx.append_history,
+            )
+            self._tool_executor = executor
+            try:
+                async for exec_event in executor.run_tools(tool_calls):
+                    for ui_event in build_events_from_tool_executor_event(session_ctx.session_id, exec_event):
                         yield ui_event
+            finally:
                 self._tool_executor = None
 
-        yield events.TurnEndEvent(session_id=ctx.session_id)
+    def _persist_partial_assistant_on_cancel(self) -> None:
+        """Persist streamed assistant text when a turn is interrupted.
+
+        Reasoning and tool calls are intentionally discarded on interrupt; only
+        the assistant message text collected so far is saved so it appears in
+        subsequent history/context.
+        """
+
+        if not self._assistant_delta_buffer:
+            return
+        partial_text = "".join(self._assistant_delta_buffer) + "<system interrupted by user>"
+        if not partial_text:
+            return
+        message_item = model.AssistantMessageItem(
+            content=partial_text,
+            response_id=self._assistant_response_id,
+        )
+        self._context.session_ctx.append_history([message_item])
+        self._assistant_delta_buffer.clear()

klaude_code/llm/__init__.py

@@ -1,21 +1,13 @@
 """LLM package init.
 
-Imports built-in LLM clients so their ``@register`` decorators run and they
-become available via the registry.
+LLM clients are lazily loaded to avoid heavy imports at module load time.
+Only LLMClientABC and create_llm_client are exposed.
 """
 
-from .anthropic import AnthropicClient
 from .client import LLMClientABC
-from .openai_compatible import OpenAICompatibleClient
-from .openrouter import OpenRouterClient
 from .registry import create_llm_client
-from .responses import ResponsesClient
 
 __all__ = [
     "LLMClientABC",
-    "ResponsesClient",
-    "OpenAICompatibleClient",
-    "OpenRouterClient",
-    "AnthropicClient",
     "create_llm_client",
 ]