klaude-code 2.0.2__py3-none-any.whl → 2.1.1__py3-none-any.whl

klaude_code/core/tool/web/web_search_tool.py

@@ -5,6 +5,7 @@ from pathlib import Path
 from pydantic import BaseModel
 
 from klaude_code.const import WEB_SEARCH_DEFAULT_MAX_RESULTS, WEB_SEARCH_MAX_RESULTS_LIMIT
+from klaude_code.core.tool.context import ToolContext
 from klaude_code.core.tool.tool_abc import ToolABC, ToolConcurrencyPolicy, ToolMetadata, load_desc
 from klaude_code.core.tool.tool_registry import register
 from klaude_code.protocol import llm_param, message, tools
@@ -91,7 +92,7 @@ class WebSearchTool(ToolABC):
         max_results: int = WEB_SEARCH_DEFAULT_MAX_RESULTS
 
     @classmethod
-    async def call(cls, arguments: str) -> message.ToolResultMessage:
+    async def call(cls, arguments: str, context: ToolContext) -> message.ToolResultMessage:
         try:
             args = WebSearchTool.WebSearchArguments.model_validate_json(arguments)
         except ValueError as e:
@@ -99,10 +100,11 @@ class WebSearchTool(ToolABC):
                 status="error",
                 output_text=f"Invalid arguments: {e}",
             )
-        return await cls.call_with_args(args)
+        return await cls.call_with_args(args, context)
 
     @classmethod
-    async def call_with_args(cls, args: WebSearchArguments) -> message.ToolResultMessage:
+    async def call_with_args(cls, args: WebSearchArguments, context: ToolContext) -> message.ToolResultMessage:
+        del context
         query = args.query.strip()
         if not query:
             return message.ToolResultMessage(

klaude_code/core/turn.py

@@ -5,8 +5,8 @@ from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
 from klaude_code.const import INTERRUPT_MARKER, SUPPORTED_IMAGE_SIZES
-from klaude_code.core.tool import ToolABC, tool_context
-from klaude_code.core.tool.tool_context import current_sub_agent_resume_claims
+from klaude_code.core.tool import ToolABC
+from klaude_code.core.tool.context import SubAgentResumeClaims, ToolContext
 
 if TYPE_CHECKING:
     from klaude_code.core.task import SessionContext
@@ -20,8 +20,8 @@ from klaude_code.core.tool.tool_runner import (
     ToolExecutorEvent,
 )
 from klaude_code.llm import LLMClientABC
+from klaude_code.log import DebugType, log_debug
 from klaude_code.protocol import events, llm_param, message, model, tools
-from klaude_code.trace import DebugType, log_debug
 
 
 class TurnError(Exception):
@@ -69,7 +69,6 @@ def build_events_from_tool_executor_event(session_id: str, event: ToolExecutorEv
                 )
             )
         case ToolExecutionResult(tool_call=tool_call, tool_result=tool_result, is_last_in_turn=is_last_in_turn):
-            status = "success" if tool_result.status == "success" else "error"
             ui_events.append(
                 events.ToolResultEvent(
                     session_id=session_id,
@@ -78,13 +77,11 @@ def build_events_from_tool_executor_event(session_id: str, event: ToolExecutorEv
                     tool_name=tool_call.tool_name,
                     result=tool_result.output_text,
                     ui_extra=tool_result.ui_extra,
-                    status=status,
+                    status=tool_result.status,
                     task_metadata=tool_result.task_metadata,
                     is_last_in_turn=is_last_in_turn,
                 )
             )
-            if tool_result.status == "aborted":
-                ui_events.append(events.InterruptEvent(session_id=session_id))
         case ToolExecutionTodoChange(todos=todos):
             ui_events.append(
                 events.TodoChangeEvent(
@@ -206,6 +203,8 @@ class TurnExecutor:
 
         ctx = self._context
         session_ctx = ctx.session_ctx
+        thinking_active = False
+        assistant_text_active = False
         message_types = (
             message.SystemMessage,
             message.DeveloperMessage,
@@ -247,12 +246,30 @@ class TurnExecutor:
             )
             match delta:
                 case message.ThinkingTextDelta() as delta:
+                    if not thinking_active:
+                        thinking_active = True
+                        yield events.ThinkingStartEvent(
+                            response_id=delta.response_id,
+                            session_id=session_ctx.session_id,
+                        )
                     yield events.ThinkingDeltaEvent(
                         content=delta.content,
                         response_id=delta.response_id,
                         session_id=session_ctx.session_id,
                     )
                 case message.AssistantTextDelta() as delta:
+                    if thinking_active:
+                        thinking_active = False
+                        yield events.ThinkingEndEvent(
+                            response_id=delta.response_id,
+                            session_id=session_ctx.session_id,
+                        )
+                    if not assistant_text_active:
+                        assistant_text_active = True
+                        yield events.AssistantTextStartEvent(
+                            response_id=delta.response_id,
+                            session_id=session_ctx.session_id,
+                        )
                     if delta.response_id:
                         self._assistant_response_id = delta.response_id
                     self._assistant_delta_buffer.append(delta.content)
@@ -262,6 +279,12 @@ class TurnExecutor:
                         session_id=session_ctx.session_id,
                     )
                 case message.AssistantImageDelta() as delta:
+                    if thinking_active:
+                        thinking_active = False
+                        yield events.ThinkingEndEvent(
+                            response_id=delta.response_id,
+                            session_id=session_ctx.session_id,
+                        )
                     yield events.AssistantImageDeltaEvent(
                         file_path=delta.file_path,
                         response_id=delta.response_id,
@@ -270,6 +293,18 @@ class TurnExecutor:
                 case message.AssistantMessage() as msg:
                     if msg.response_id is None and self._assistant_response_id:
                         msg.response_id = self._assistant_response_id
+                    if thinking_active:
+                        thinking_active = False
+                        yield events.ThinkingEndEvent(
+                            response_id=msg.response_id,
+                            session_id=session_ctx.session_id,
+                        )
+                    if assistant_text_active:
+                        assistant_text_active = False
+                        yield events.AssistantTextEndEvent(
+                            response_id=msg.response_id,
+                            session_id=session_ctx.session_id,
+                        )
                     turn_result.assistant_message = msg
                     for part in msg.parts:
                         if isinstance(part, message.ToolCallPart):
@@ -281,11 +316,16 @@ class TurnExecutor:
                                     arguments_json=part.arguments_json,
                                 )
                             )
-                    yield events.AssistantMessageEvent(
-                        content=message.join_text_parts(msg.parts),
-                        response_id=msg.response_id,
-                        session_id=session_ctx.session_id,
-                    )
+                    if msg.stop_reason != "aborted":
+                        thinking_text = "".join(
+                            part.text for part in msg.parts if isinstance(part, message.ThinkingTextPart)
+                        )
+                        yield events.ResponseCompleteEvent(
+                            content=message.join_text_parts(msg.parts),
+                            response_id=msg.response_id,
+                            session_id=session_ctx.session_id,
+                            thinking_text=thinking_text or None,
+                        )
                     if msg.stop_reason == "aborted":
                         yield events.InterruptEvent(session_id=session_ctx.session_id)
                     if msg.usage:
@@ -296,9 +336,9 @@ class TurnExecutor:
                             metadata.model_name = ctx.llm_client.model_name
                         if metadata.provider is None:
                             metadata.provider = ctx.llm_client.get_llm_config().provider_name or None
-                        yield events.ResponseMetadataEvent(
+                        yield events.UsageEvent(
                             session_id=session_ctx.session_id,
-                            metadata=metadata,
+                            usage=metadata,
                         )
                 case message.StreamErrorItem() as msg:
                     turn_result.stream_error = msg
@@ -309,12 +349,23 @@ class TurnExecutor:
                         debug_type=DebugType.RESPONSE,
                     )
                 case message.ToolCallStartItem() as msg:
-                    yield events.TurnToolCallStartEvent(
+                    if thinking_active:
+                        thinking_active = False
+                        yield events.ThinkingEndEvent(
+                            response_id=msg.response_id,
+                            session_id=session_ctx.session_id,
+                        )
+                    if assistant_text_active:
+                        assistant_text_active = False
+                        yield events.AssistantTextEndEvent(
+                            response_id=msg.response_id,
+                            session_id=session_ctx.session_id,
+                        )
+                    yield events.ToolCallStartEvent(
                         session_id=session_ctx.session_id,
                         response_id=msg.response_id,
                         tool_call_id=msg.call_id,
                         tool_name=msg.name,
-                        arguments="",
                     )
                 case _:
                     continue
@@ -332,20 +383,26 @@ class TurnExecutor:
 
         ctx = self._context
         session_ctx = ctx.session_ctx
-        with tool_context(session_ctx.file_tracker, session_ctx.todo_context):
-            resume_claims_token = current_sub_agent_resume_claims.set(set())
-            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
-                current_sub_agent_resume_claims.reset(resume_claims_token)
+        tool_context = ToolContext(
+            file_tracker=session_ctx.file_tracker,
+            todo_context=session_ctx.todo_context,
+            session_id=session_ctx.session_id,
+            run_subtask=session_ctx.run_subtask,
+            sub_agent_resume_claims=SubAgentResumeClaims(),
+        )
+
+        executor = ToolExecutor(
+            context=tool_context,
+            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
 
     def _persist_partial_assistant_on_cancel(self) -> None:
         """Persist streamed assistant text when a turn is interrupted.

klaude_code/llm/anthropic/client.py

@@ -34,8 +34,8 @@ from klaude_code.llm.client import LLMClientABC
 from klaude_code.llm.input_common import apply_config_defaults
 from klaude_code.llm.registry import register
 from klaude_code.llm.usage import MetadataTracker, error_stream_items
+from klaude_code.log import DebugType, log_debug
 from klaude_code.protocol import llm_param, message, model
-from klaude_code.trace import DebugType, log_debug
 
 
 def _map_anthropic_stop_reason(reason: str) -> model.StopReason | None:

klaude_code/llm/bedrock/client.py

@@ -14,8 +14,8 @@ from klaude_code.llm.client import LLMClientABC
 from klaude_code.llm.input_common import apply_config_defaults
 from klaude_code.llm.registry import register
 from klaude_code.llm.usage import MetadataTracker, error_stream_items
+from klaude_code.log import DebugType, log_debug
 from klaude_code.protocol import llm_param, message
-from klaude_code.trace import DebugType, log_debug
 
 
 @register(llm_param.LLMClientProtocol.BEDROCK)

klaude_code/llm/claude/client.py

@@ -22,8 +22,8 @@ from klaude_code.llm.client import LLMClientABC
 from klaude_code.llm.input_common import apply_config_defaults
 from klaude_code.llm.registry import register
 from klaude_code.llm.usage import MetadataTracker, error_stream_items
+from klaude_code.log import DebugType, log_debug
 from klaude_code.protocol import llm_param, message
-from klaude_code.trace import DebugType, log_debug
 
 _CLAUDE_OAUTH_REQUIRED_BETAS: tuple[str, ...] = (
     ANTHROPIC_BETA_OAUTH,

klaude_code/llm/codex/client.py

@@ -25,8 +25,8 @@ 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, error_stream_items
+from klaude_code.log import DebugType, log_debug
 from klaude_code.protocol import llm_param, message
-from klaude_code.trace import DebugType, log_debug
 
 
 def build_payload(param: llm_param.LLMCallParameter) -> ResponseCreateParamsStreaming:

klaude_code/llm/google/client.py

@@ -26,8 +26,8 @@ from klaude_code.llm.google.input import convert_history_to_contents, convert_to
 from klaude_code.llm.input_common import apply_config_defaults
 from klaude_code.llm.registry import register
 from klaude_code.llm.usage import MetadataTracker
+from klaude_code.log import DebugType, log_debug
 from klaude_code.protocol import llm_param, message, model
-from klaude_code.trace import DebugType, log_debug
 
 
 def _build_config(param: llm_param.LLMCallParameter) -> GenerateContentConfig:

klaude_code/llm/openai_compatible/client.py

@@ -13,8 +13,8 @@ from klaude_code.llm.openai_compatible.input import convert_history_to_input, co
 from klaude_code.llm.openai_compatible.stream import DefaultReasoningHandler, parse_chat_completions_stream
 from klaude_code.llm.registry import register
 from klaude_code.llm.usage import MetadataTracker
+from klaude_code.log import DebugType, log_debug
 from klaude_code.protocol import llm_param, message
-from klaude_code.trace import DebugType, log_debug
 
 
 def build_payload(param: llm_param.LLMCallParameter) -> tuple[CompletionCreateParamsStreaming, dict[str, object]]:

klaude_code/llm/openai_compatible/tool_call_accumulator.py

@@ -4,8 +4,8 @@ from abc import ABC, abstractmethod
 from openai.types.chat.chat_completion_chunk import ChoiceDeltaToolCall
 from pydantic import BaseModel, Field
 
+from klaude_code.log import log_debug
 from klaude_code.protocol import message
-from klaude_code.trace.log import log_debug
 
 
 def normalize_tool_name(name: str) -> str:

klaude_code/llm/openrouter/client.py

@@ -22,8 +22,8 @@ from klaude_code.llm.openrouter.input import convert_history_to_input, is_claude
 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.log import DebugType, is_debug_enabled, log_debug
 from klaude_code.protocol import llm_param, message
-from klaude_code.trace import DebugType, is_debug_enabled, log_debug
 
 
 def build_payload(

klaude_code/llm/openrouter/reasoning.py

@@ -1,8 +1,8 @@
 from pydantic import BaseModel
 
 from klaude_code.llm.openai_compatible.stream import ReasoningDeltaResult, ReasoningHandlerABC
+from klaude_code.log import log
 from klaude_code.protocol import message
-from klaude_code.trace import log
 
 
 class ReasoningDetail(BaseModel):

klaude_code/llm/responses/client.py

@@ -14,8 +14,8 @@ from klaude_code.llm.input_common import apply_config_defaults
 from klaude_code.llm.registry import register
 from klaude_code.llm.responses.input import convert_history_to_input, convert_tool_schema
 from klaude_code.llm.usage import MetadataTracker, error_stream_items
+from klaude_code.log import DebugType, log_debug
 from klaude_code.protocol import llm_param, message, model
-from klaude_code.trace import DebugType, log_debug
 
 if TYPE_CHECKING:
     from openai import AsyncStream

klaude_code/protocol/commands.py

@@ -28,6 +28,7 @@ class CommandName(str, Enum):
     THINKING = "thinking"
     FORK_SESSION = "fork-session"
     RESUME = "resume"
+    COPY = "copy"
     # PLAN and DOC are dynamically registered now, but kept here if needed for reference
     # or we can remove them if no code explicitly imports them.
     # PLAN = "plan"

klaude_code/protocol/events/__init__.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+from klaude_code.protocol.events.base import Event, ResponseEvent
+from klaude_code.protocol.events.chat import DeveloperMessageEvent, TodoChangeEvent, UserMessageEvent
+from klaude_code.protocol.events.lifecycle import TaskFinishEvent, TaskStartEvent, TurnEndEvent, TurnStartEvent
+from klaude_code.protocol.events.metadata import TaskMetadataEvent, UsageEvent
+from klaude_code.protocol.events.streaming import (
+    AssistantImageDeltaEvent,
+    AssistantTextDeltaEvent,
+    AssistantTextEndEvent,
+    AssistantTextStartEvent,
+    ResponseCompleteEvent,
+    ThinkingDeltaEvent,
+    ThinkingEndEvent,
+    ThinkingStartEvent,
+    ToolCallStartEvent,
+)
+from klaude_code.protocol.events.system import (
+    EndEvent,
+    ErrorEvent,
+    InterruptEvent,
+    ReplayEventUnion,
+    ReplayHistoryEvent,
+    WelcomeEvent,
+)
+from klaude_code.protocol.events.tools import ToolCallEvent, ToolResultEvent
+
+__all__ = [
+    "AssistantImageDeltaEvent",
+    "AssistantTextDeltaEvent",
+    "AssistantTextEndEvent",
+    "AssistantTextStartEvent",
+    "DeveloperMessageEvent",
+    "EndEvent",
+    "ErrorEvent",
+    "Event",
+    "InterruptEvent",
+    "ReplayEventUnion",
+    "ReplayHistoryEvent",
+    "ResponseCompleteEvent",
+    "ResponseEvent",
+    "TaskFinishEvent",
+    "TaskMetadataEvent",
+    "TaskStartEvent",
+    "ThinkingDeltaEvent",
+    "ThinkingEndEvent",
+    "ThinkingStartEvent",
+    "TodoChangeEvent",
+    "ToolCallEvent",
+    "ToolCallStartEvent",
+    "ToolResultEvent",
+    "TurnEndEvent",
+    "TurnStartEvent",
+    "UsageEvent",
+    "UserMessageEvent",
+    "WelcomeEvent",
+]

klaude_code/protocol/events/base.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+import time
+
+from pydantic import BaseModel, Field
+
+
+class Event(BaseModel):
+    """Base event."""
+
+    session_id: str
+    timestamp: float = Field(default_factory=time.time)
+
+
+class ResponseEvent(Event):
+    """Event associated with a single model response."""
+
+    response_id: str | None = None

klaude_code/protocol/events/chat.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+from klaude_code.protocol import message, model
+
+from .base import Event
+
+
+class UserMessageEvent(Event):
+    content: str
+    images: list[message.ImageURLPart] | None = None
+
+
+class DeveloperMessageEvent(Event):
+    """DeveloperMessages are reminders in user messages or tool results."""
+
+    item: message.DeveloperMessage
+
+
+class TodoChangeEvent(Event):
+    todos: list[model.TodoItem]

klaude_code/protocol/events/lifecycle.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+from klaude_code.protocol import model
+
+from .base import Event
+
+
+class TaskStartEvent(Event):
+    sub_agent_state: model.SubAgentState | None = None
+
+
+class TaskFinishEvent(Event):
+    task_result: str
+    has_structured_output: bool = False
+
+
+class TurnStartEvent(Event):
+    pass
+
+
+class TurnEndEvent(Event):
+    pass

klaude_code/protocol/events/metadata.py

@@ -0,0 +1,15 @@
+"""Metadata-related protocol events."""
+
+from __future__ import annotations
+
+from klaude_code.protocol import model
+
+from .base import Event, ResponseEvent
+
+
+class UsageEvent(ResponseEvent):
+    usage: model.Usage
+
+
+class TaskMetadataEvent(Event):
+    metadata: model.TaskMetadataItem

klaude_code/protocol/events/streaming.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from .base import ResponseEvent
+
+
+class ThinkingStartEvent(ResponseEvent):
+    pass
+
+
+class ThinkingDeltaEvent(ResponseEvent):
+    content: str
+
+
+class ThinkingEndEvent(ResponseEvent):
+    pass
+
+
+class AssistantTextStartEvent(ResponseEvent):
+    pass
+
+
+class AssistantTextDeltaEvent(ResponseEvent):
+    content: str
+
+
+class AssistantTextEndEvent(ResponseEvent):
+    pass
+
+
+class AssistantImageDeltaEvent(ResponseEvent):
+    file_path: str
+
+
+class ToolCallStartEvent(ResponseEvent):
+    tool_call_id: str
+    tool_name: str
+
+
+class ResponseCompleteEvent(ResponseEvent):
+    """Final snapshot of the model response."""
+
+    content: str
+    thinking_text: str | None = None

klaude_code/protocol/events/system.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from klaude_code.protocol import llm_param
+from klaude_code.protocol.events.chat import DeveloperMessageEvent, UserMessageEvent
+from klaude_code.protocol.events.lifecycle import TaskFinishEvent, TaskStartEvent, TurnStartEvent
+from klaude_code.protocol.events.metadata import TaskMetadataEvent
+from klaude_code.protocol.events.streaming import AssistantImageDeltaEvent, ResponseCompleteEvent
+from klaude_code.protocol.events.tools import ToolCallEvent, ToolResultEvent
+
+from .base import Event
+
+
+class WelcomeEvent(Event):
+    work_dir: str
+    llm_config: llm_param.LLMConfigParameter
+    show_klaude_code_info: bool = True
+
+
+class ErrorEvent(Event):
+    error_message: str
+    can_retry: bool = False
+
+
+class InterruptEvent(Event):
+    pass
+
+
+class EndEvent(Event):
+    """Global display shutdown."""
+
+    session_id: str = "__app__"
+
+
+type ReplayEventUnion = (
+    TaskStartEvent
+    | TaskFinishEvent
+    | TurnStartEvent
+    | AssistantImageDeltaEvent
+    | ResponseCompleteEvent
+    | ToolCallEvent
+    | ToolResultEvent
+    | UserMessageEvent
+    | TaskMetadataEvent
+    | InterruptEvent
+    | DeveloperMessageEvent
+    | ErrorEvent
+)
+
+
+class ReplayHistoryEvent(Event):
+    events: list[ReplayEventUnion]
+    updated_at: float
+    is_load: bool = True

klaude_code/protocol/events/tools.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from klaude_code.protocol import model
+
+from .base import ResponseEvent
+
+
+class ToolCallEvent(ResponseEvent):
+    tool_call_id: str
+    tool_name: str
+    arguments: str
+
+
+class ToolResultEvent(ResponseEvent):
+    tool_call_id: str
+    tool_name: str
+    result: str
+    ui_extra: model.ToolResultUIExtra | None = None
+    status: Literal["success", "error", "aborted"]
+    task_metadata: model.TaskMetadata | None = None
+    is_last_in_turn: bool = True
+
+    @property
+    def is_error(self) -> bool:
+        return self.status in ("error", "aborted")

klaude_code/protocol/op.py

@@ -51,6 +51,11 @@ class RunAgentOperation(Operation):
     type: OperationType = OperationType.RUN_AGENT
     session_id: str
     input: UserInputPayload
+    # Frontends may choose to render the user message themselves (e.g. TUI) to support
+    # event-only commands; in that case the core should skip emitting the UserMessageEvent.
+    emit_user_message_event: bool = True
+    # Frontends may choose to run without persisting input (e.g. some interactive commands).
+    persist_user_input: bool = True
 
     async def execute(self, handler: OperationHandler) -> None:
         await handler.handle_run_agent(self)

klaude_code/protocol/tools.py

@@ -6,7 +6,6 @@ READ = "Read"
 WRITE = "Write"
 TODO_WRITE = "TodoWrite"
 UPDATE_PLAN = "update_plan"
-SKILL = "Skill"
 MERMAID = "Mermaid"
 WEB_FETCH = "WebFetch"
 WEB_SEARCH = "WebSearch"