klaude-code 1.9.0__py3-none-any.whl → 2.0.0__py3-none-any.whl

klaude_code/core/tool/todo/update_plan_tool.py

@@ -9,7 +9,7 @@ from pydantic import BaseModel, field_validator
 from klaude_code.core.tool.tool_abc import ToolABC, load_desc
 from klaude_code.core.tool.tool_context import get_current_todo_context
 from klaude_code.core.tool.tool_registry import register
-from klaude_code.protocol import llm_param, model, tools
+from klaude_code.protocol import llm_param, message, model, tools
 
 from .todo_write_tool import get_new_completed_todos
 
@@ -79,15 +79,15 @@ class UpdatePlanTool(ToolABC):
         )
 
     @classmethod
-    async def call(cls, arguments: str) -> model.ToolResultItem:
+    async def call(cls, arguments: str) -> message.ToolResultMessage:
         try:
             args = UpdatePlanArguments.model_validate_json(arguments)
         except ValueError as exc:
-            return model.ToolResultItem(status="error", output=f"Invalid arguments: {exc}")
+            return message.ToolResultMessage(status="error", output_text=f"Invalid arguments: {exc}")
 
         todo_context = get_current_todo_context()
         if todo_context is None:
-            return model.ToolResultItem(status="error", output="No active session found")
+            return message.ToolResultMessage(status="error", output_text="No active session found")
 
         new_todos = [model.TodoItem(content=item.step, status=item.status) for item in args.plan]
         old_todos = todo_context.get_todos()
@@ -96,9 +96,9 @@ class UpdatePlanTool(ToolABC):
 
         ui_extra = model.TodoUIExtra(todos=new_todos, new_completed=new_completed)
 
-        return model.ToolResultItem(
+        return message.ToolResultMessage(
             status="success",
-            output="Plan updated",
+            output_text="Plan updated",
             ui_extra=model.TodoListUIExtra(todo_list=ui_extra),
             side_effects=[model.ToolSideEffect.TODO_CHANGE],
         )

klaude_code/core/tool/tool_abc.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
 
-from klaude_code.protocol import llm_param, model
+from klaude_code.protocol import llm_param, message
 
 
 def load_desc(path: Path, substitutions: dict[str, str] | None = None) -> str:
@@ -27,7 +27,7 @@ class ToolABC(ABC):
 
     @classmethod
     @abstractmethod
-    async def call(cls, arguments: str) -> model.ToolResultItem:
+    async def call(cls, arguments: str) -> message.ToolResultMessage:
         raise NotImplementedError
 
 

klaude_code/core/tool/tool_context.py

@@ -119,3 +119,30 @@ def get_current_todo_context() -> TodoContext | None:
 current_run_subtask_callback: ContextVar[Callable[[model.SubAgentState], Awaitable[SubAgentResult]] | None] = (
     ContextVar("current_run_subtask_callback", default=None)
 )
+
+
+# Allows sub-agent execution to record the created/used session id for the currently
+# executing tool call (used by ToolExecutor.cancel() to include session_id in UIExtra).
+current_sub_agent_session_id_recorder: ContextVar[Callable[[str], None] | None] = ContextVar(
+    "current_sub_agent_session_id_recorder",
+    default=None,
+)
+
+
+def record_sub_agent_session_id(session_id: str) -> None:
+    """Record the sub-agent session id for the current tool call, if supported."""
+
+    recorder = current_sub_agent_session_id_recorder.get()
+    if recorder is None:
+        return
+    recorder(session_id)
+
+
+# Tracks sub-agent resume claims for the current turn.
+#
+# This is used to reject multiple sub-agent tool calls in the same LLM response
+# that attempt to resume the same agent ID.
+current_sub_agent_resume_claims: ContextVar[set[str] | None] = ContextVar(
+    "current_sub_agent_resume_claims",
+    default=None,
+)

klaude_code/core/tool/tool_runner.py

@@ -2,14 +2,23 @@ import asyncio
 from collections.abc import AsyncGenerator, Callable, Iterable, Sequence
 from dataclasses import dataclass
 
-from klaude_code import const
+from klaude_code.const import CANCEL_OUTPUT
 from klaude_code.core.tool.report_back_tool import ReportBackTool
 from klaude_code.core.tool.tool_abc import ToolABC, ToolConcurrencyPolicy
+from klaude_code.core.tool.tool_context import current_sub_agent_session_id_recorder
 from klaude_code.core.tool.truncation import truncate_tool_output
-from klaude_code.protocol import model, tools
+from klaude_code.protocol import message, model, tools
 
 
-async def run_tool(tool_call: model.ToolCallItem, registry: dict[str, type[ToolABC]]) -> model.ToolResultItem:
+@dataclass(frozen=True)
+class ToolCallRequest:
+    response_id: str | None
+    call_id: str
+    tool_name: str
+    arguments_json: str
+
+
+async def run_tool(tool_call: ToolCallRequest, registry: dict[str, type[ToolABC]]) -> message.ToolResultMessage:
     """Execute a tool call and return the result.
 
     Args:
@@ -20,26 +29,26 @@ async def run_tool(tool_call: model.ToolCallItem, registry: dict[str, type[ToolA
         The result of the tool execution.
     """
     # Special handling for report_back tool (not registered in global registry)
-    if tool_call.name == tools.REPORT_BACK:
-        tool_result = await ReportBackTool.call(tool_call.arguments)
+    if tool_call.tool_name == tools.REPORT_BACK:
+        tool_result = await ReportBackTool.call(tool_call.arguments_json)
         tool_result.call_id = tool_call.call_id
-        tool_result.tool_name = tool_call.name
+        tool_result.tool_name = tool_call.tool_name
         return tool_result
 
-    if tool_call.name not in registry:
-        return model.ToolResultItem(
+    if tool_call.tool_name not in registry:
+        return message.ToolResultMessage(
             call_id=tool_call.call_id,
-            output=f"Tool {tool_call.name} not exists",
+            output_text=f"Tool {tool_call.tool_name} not exists",
             status="error",
-            tool_name=tool_call.name,
+            tool_name=tool_call.tool_name,
         )
     try:
-        tool_result = await registry[tool_call.name].call(tool_call.arguments)
+        tool_result = await registry[tool_call.tool_name].call(tool_call.arguments_json)
         tool_result.call_id = tool_call.call_id
-        tool_result.tool_name = tool_call.name
-        if tool_result.output:
-            truncation_result = truncate_tool_output(tool_result.output, tool_call)
-            tool_result.output = truncation_result.output
+        tool_result.tool_name = tool_call.tool_name
+        if tool_result.output_text:
+            truncation_result = truncate_tool_output(tool_result.output_text, tool_call)
+            tool_result.output_text = truncation_result.output
             if truncation_result.was_truncated and truncation_result.saved_file_path:
                 tool_result.ui_extra = model.TruncationUIExtra(
                     saved_file_path=truncation_result.saved_file_path,
@@ -51,11 +60,11 @@ async def run_tool(tool_call: model.ToolCallItem, registry: dict[str, type[ToolA
         # Propagate cooperative cancellation so outer layers can handle interrupts correctly.
         raise
     except Exception as e:
-        return model.ToolResultItem(
+        return message.ToolResultMessage(
             call_id=tool_call.call_id,
-            output=f"Tool {tool_call.name} execution error: {e.__class__.__name__} {e}",
+            output_text=f"Tool {tool_call.tool_name} execution error: {e.__class__.__name__} {e}",
             status="error",
-            tool_name=tool_call.name,
+            tool_name=tool_call.tool_name,
         )
 
 
@@ -63,15 +72,18 @@ async def run_tool(tool_call: model.ToolCallItem, registry: dict[str, type[ToolA
 class ToolExecutionCallStarted:
     """Represents the start of a tool call execution."""
 
-    tool_call: model.ToolCallItem
+    tool_call: ToolCallRequest
 
 
 @dataclass
 class ToolExecutionResult:
     """Represents the completion of a tool call with its result."""
 
-    tool_call: model.ToolCallItem
-    tool_result: model.ToolResultItem
+    tool_call: ToolCallRequest
+    tool_result: message.ToolResultMessage
+    # Whether this is the last ToolExecutionResult emitted in the current turn.
+    # Used by UI to decide whether to close the tree prefix.
+    is_last_in_turn: bool = False
 
 
 @dataclass
@@ -98,16 +110,17 @@ class ToolExecutor:
         self,
         *,
         registry: dict[str, type[ToolABC]],
-        append_history: Callable[[Sequence[model.ConversationItem]], None],
+        append_history: Callable[[Sequence[message.HistoryEvent]], None],
     ) -> None:
         self._registry = registry
         self._append_history = append_history
 
-        self._unfinished_calls: dict[str, model.ToolCallItem] = {}
+        self._unfinished_calls: dict[str, ToolCallRequest] = {}
         self._call_event_emitted: set[str] = set()
         self._concurrent_tasks: set[asyncio.Task[list[ToolExecutorEvent]]] = set()
+        self._sub_agent_session_ids: dict[str, str] = {}
 
-    async def run_tools(self, tool_calls: list[model.ToolCallItem]) -> AsyncGenerator[ToolExecutorEvent]:
+    async def run_tools(self, tool_calls: list[ToolCallRequest]) -> AsyncGenerator[ToolExecutorEvent]:
         """Run the given tool calls and yield execution events.
 
         Tool calls are partitioned into regular tools and sub-agent tools. Regular tools
@@ -120,8 +133,15 @@ class ToolExecutor:
 
         sequential_tool_calls, concurrent_tool_calls = self._partition_tool_calls(tool_calls)
 
+        def _mark_last_in_turn(events_to_mark: list[ToolExecutorEvent], *, is_last_in_turn: bool) -> None:
+            if not events_to_mark:
+                return
+            first = events_to_mark[0]
+            if isinstance(first, ToolExecutionResult):
+                first.is_last_in_turn = is_last_in_turn
+
         # Run sequential tools one by one.
-        for tool_call in sequential_tool_calls:
+        for idx, tool_call in enumerate(sequential_tool_calls):
             tool_call_event = self._build_tool_call_started(tool_call)
             self._call_event_emitted.add(tool_call.call_id)
             yield tool_call_event
@@ -132,6 +152,9 @@ class ToolExecutor:
                 # Propagate cooperative cancellation so the agent task can be stopped.
                 raise
 
+            is_last_in_turn = idx == len(sequential_tool_calls) - 1 and not concurrent_tool_calls
+            _mark_last_in_turn(result_events, is_last_in_turn=is_last_in_turn)
+
             for exec_event in result_events:
                 yield exec_event
 
@@ -147,6 +170,7 @@ class ToolExecutor:
                 self._register_concurrent_task(task)
                 execution_tasks.append(task)
 
+            remaining = len(execution_tasks)
             for task in asyncio.as_completed(execution_tasks):
                 # Do not swallow asyncio.CancelledError here:
                 # - If the user interrupts the main agent, the executor cancels the
@@ -158,6 +182,9 @@ class ToolExecutor:
                 #   calling agent can stop cleanly, matching pre-refactor behavior.
                 result_events = await task
 
+                remaining -= 1
+                _mark_last_in_turn(result_events, is_last_in_turn=remaining == 0)
+
                 for exec_event in result_events:
                     yield exec_event
 
@@ -168,7 +195,7 @@ class ToolExecutor:
         - For each unfinished tool call, yields a ToolExecutionCallStarted (if not
           already emitted for this turn) followed by a ToolExecutionResult with
           error status and a standard cancellation output. The corresponding
-          ToolResultItem is appended to history via `append_history`.
+          ToolResultMessage is appended to history via `append_history`.
         """
 
         events_to_yield: list[ToolExecutorEvent] = []
@@ -182,23 +209,32 @@ class ToolExecutor:
         if not self._unfinished_calls:
             return events_to_yield
 
-        for call_id, tool_call in list(self._unfinished_calls.items()):
-            cancel_result = model.ToolResultItem(
+        unfinished = list(self._unfinished_calls.items())
+        for idx, (call_id, tool_call) in enumerate(unfinished):
+            session_id = self._sub_agent_session_ids.get(call_id)
+            cancel_result = message.ToolResultMessage(
                 call_id=tool_call.call_id,
-                output=const.CANCEL_OUTPUT,
-                status="error",
-                tool_name=tool_call.name,
-                ui_extra=None,
+                output_text=CANCEL_OUTPUT,
+                status="aborted",
+                tool_name=tool_call.tool_name,
+                ui_extra=model.SessionIdUIExtra(session_id=session_id) if session_id else None,
             )
 
             if call_id not in self._call_event_emitted:
                 events_to_yield.append(ToolExecutionCallStarted(tool_call=tool_call))
                 self._call_event_emitted.add(call_id)
 
-            events_to_yield.append(ToolExecutionResult(tool_call=tool_call, tool_result=cancel_result))
+            events_to_yield.append(
+                ToolExecutionResult(
+                    tool_call=tool_call,
+                    tool_result=cancel_result,
+                    is_last_in_turn=idx == len(unfinished) - 1,
+                )
+            )
 
             self._append_history([cancel_result])
             self._unfinished_calls.pop(call_id, None)
+            self._sub_agent_session_ids.pop(call_id, None)
 
         return events_to_yield
 
@@ -212,12 +248,12 @@ class ToolExecutor:
 
     def _partition_tool_calls(
         self,
-        tool_calls: list[model.ToolCallItem],
-    ) -> tuple[list[model.ToolCallItem], list[model.ToolCallItem]]:
-        sequential_tool_calls: list[model.ToolCallItem] = []
-        concurrent_tool_calls: list[model.ToolCallItem] = []
+        tool_calls: list[ToolCallRequest],
+    ) -> tuple[list[ToolCallRequest], list[ToolCallRequest]]:
+        sequential_tool_calls: list[ToolCallRequest] = []
+        concurrent_tool_calls: list[ToolCallRequest] = []
         for tool_call in tool_calls:
-            tool_cls = self._registry.get(tool_call.name)
+            tool_cls = self._registry.get(tool_call.tool_name)
             policy = (
                 tool_cls.metadata().concurrency_policy if tool_cls is not None else ToolConcurrencyPolicy.SEQUENTIAL
             )
@@ -227,22 +263,32 @@ class ToolExecutor:
                 sequential_tool_calls.append(tool_call)
         return sequential_tool_calls, concurrent_tool_calls
 
-    def _build_tool_call_started(self, tool_call: model.ToolCallItem) -> ToolExecutionCallStarted:
+    def _build_tool_call_started(self, tool_call: ToolCallRequest) -> ToolExecutionCallStarted:
         return ToolExecutionCallStarted(tool_call=tool_call)
 
-    async def _run_single_tool_call(self, tool_call: model.ToolCallItem) -> list[ToolExecutorEvent]:
-        tool_result: model.ToolResultItem = await run_tool(tool_call, self._registry)
+    async def _run_single_tool_call(self, tool_call: ToolCallRequest) -> list[ToolExecutorEvent]:
+        def _record_sub_agent_session_id(session_id: str) -> None:
+            # Keep the first recorded id if multiple writes happen.
+            if tool_call.call_id not in self._sub_agent_session_ids:
+                self._sub_agent_session_ids[tool_call.call_id] = session_id
+
+        recorder_token = current_sub_agent_session_id_recorder.set(_record_sub_agent_session_id)
+        try:
+            tool_result: message.ToolResultMessage = await run_tool(tool_call, self._registry)
+        finally:
+            current_sub_agent_session_id_recorder.reset(recorder_token)
 
         self._append_history([tool_result])
 
         result_event = ToolExecutionResult(tool_call=tool_call, tool_result=tool_result)
 
         self._unfinished_calls.pop(tool_call.call_id, None)
+        self._sub_agent_session_ids.pop(tool_call.call_id, None)
 
         extra_events = self._build_tool_side_effect_events(tool_result)
         return [result_event, *extra_events]
 
-    def _build_tool_side_effect_events(self, tool_result: model.ToolResultItem) -> list[ToolExecutorEvent]:
+    def _build_tool_side_effect_events(self, tool_result: message.ToolResultMessage) -> list[ToolExecutorEvent]:
         side_effects = tool_result.side_effects
         if not side_effects:
             return []

klaude_code/core/tool/truncation.py

@@ -4,10 +4,28 @@ import time
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from pathlib import Path
+from typing import Protocol
 from urllib.parse import urlparse
 
-from klaude_code import const
-from klaude_code.protocol import model, tools
+from klaude_code.const import (
+    TOOL_OUTPUT_DISPLAY_HEAD,
+    TOOL_OUTPUT_DISPLAY_TAIL,
+    TOOL_OUTPUT_MAX_LENGTH,
+    TOOL_OUTPUT_TRUNCATION_DIR,
+    URL_FILENAME_MAX_LENGTH,
+)
+from klaude_code.protocol import tools
+
+
+class ToolCallLike(Protocol):
+    @property
+    def tool_name(self) -> str: ...
+
+    @property
+    def call_id(self) -> str: ...
+
+    @property
+    def arguments_json(self) -> str: ...
 
 
 @dataclass
@@ -40,14 +58,14 @@ def _extract_url_filename(url: str) -> str:
     # Sanitize: keep only alphanumeric, underscore, hyphen
     name = re.sub(r"[^a-zA-Z0-9_\-]", "_", name)
     # Limit length
-    return name[:80] if len(name) > 80 else name
+    return name[:URL_FILENAME_MAX_LENGTH] if len(name) > URL_FILENAME_MAX_LENGTH else name
 
 
 class TruncationStrategy(ABC):
     """Abstract base class for tool output truncation strategies."""
 
     @abstractmethod
-    def truncate(self, output: str, tool_call: model.ToolCallItem | None = None) -> TruncationResult:
+    def truncate(self, output: str, tool_call: ToolCallLike | None = None) -> TruncationResult:
         """Truncate the output according to the strategy."""
         ...
 
@@ -55,13 +73,13 @@ class TruncationStrategy(ABC):
 class SimpleTruncationStrategy(TruncationStrategy):
     """Simple character-based truncation strategy."""
 
-    def __init__(self, max_length: int = const.TOOL_OUTPUT_MAX_LENGTH):
+    def __init__(self, max_length: int = TOOL_OUTPUT_MAX_LENGTH):
         self.max_length = max_length
 
-    def truncate(self, output: str, tool_call: model.ToolCallItem | None = None) -> TruncationResult:
+    def truncate(self, output: str, tool_call: ToolCallLike | None = None) -> TruncationResult:
         if len(output) > self.max_length:
             truncated_length = len(output) - self.max_length
-            truncated_output = output[: self.max_length] + f"... (truncated {truncated_length} characters)"
+            truncated_output = output[: self.max_length] + f"… (truncated {truncated_length} characters)"
             return TruncationResult(
                 output=truncated_output,
                 was_truncated=True,
@@ -76,21 +94,21 @@ class SmartTruncationStrategy(TruncationStrategy):
 
     def __init__(
         self,
-        max_length: int = const.TOOL_OUTPUT_MAX_LENGTH,
-        head_chars: int = const.TOOL_OUTPUT_DISPLAY_HEAD,
-        tail_chars: int = const.TOOL_OUTPUT_DISPLAY_TAIL,
-        truncation_dir: str = const.TOOL_OUTPUT_TRUNCATION_DIR,
+        max_length: int = TOOL_OUTPUT_MAX_LENGTH,
+        head_chars: int = TOOL_OUTPUT_DISPLAY_HEAD,
+        tail_chars: int = TOOL_OUTPUT_DISPLAY_TAIL,
+        truncation_dir: str = TOOL_OUTPUT_TRUNCATION_DIR,
     ):
         self.max_length = max_length
         self.head_chars = head_chars
         self.tail_chars = tail_chars
         self.truncation_dir = Path(truncation_dir)
 
-    def _get_file_identifier(self, tool_call: model.ToolCallItem | None) -> str:
+    def _get_file_identifier(self, tool_call: ToolCallLike | None) -> str:
         """Get a file identifier based on tool call. For WebFetch, use URL; otherwise use call_id."""
-        if tool_call and tool_call.name == tools.WEB_FETCH:
+        if tool_call and tool_call.tool_name == tools.WEB_FETCH:
             try:
-                args = json.loads(tool_call.arguments)
+                args = json.loads(tool_call.arguments_json)
                 url = args.get("url", "")
                 if url:
                     return _extract_url_filename(url)
@@ -101,12 +119,12 @@ class SmartTruncationStrategy(TruncationStrategy):
             return tool_call.call_id.replace("/", "_")
         return "unknown"
 
-    def _save_to_file(self, output: str, tool_call: model.ToolCallItem | None) -> str | None:
+    def _save_to_file(self, output: str, tool_call: ToolCallLike | None) -> str | None:
         """Save full output to file. Returns file path or None on failure."""
         try:
             self.truncation_dir.mkdir(parents=True, exist_ok=True)
             timestamp = int(time.time())
-            tool_name = (tool_call.name if tool_call else "unknown").replace("/", "_")
+            tool_name = (tool_call.tool_name if tool_call else "unknown").replace("/", "_")
             identifier = self._get_file_identifier(tool_call)
             filename = f"{tool_name}-{identifier}-{timestamp}.txt"
             file_path = self.truncation_dir / filename
@@ -115,8 +133,8 @@ class SmartTruncationStrategy(TruncationStrategy):
         except OSError:
             return None
 
-    def truncate(self, output: str, tool_call: model.ToolCallItem | None = None) -> TruncationResult:
-        if tool_call and tool_call.name == tools.READ:
+    def truncate(self, output: str, tool_call: ToolCallLike | None = None) -> TruncationResult:
+        if tool_call and tool_call.tool_name == tools.READ:
             # Do not truncate Read tool outputs
             return TruncationResult(output=output, was_truncated=False, original_length=len(output))
 
@@ -153,7 +171,7 @@ class SmartTruncationStrategy(TruncationStrategy):
 
         truncated_output = (
             f"{header}{head_content}\n\n"
-            f"<system-reminder>... {truncated_length} characters omitted ...</system-reminder>\n\n"
+            f"<system-reminder>… {truncated_length} characters omitted …</system-reminder>\n\n"
             f"{tail_content}"
         )
 
@@ -180,6 +198,6 @@ def set_truncation_strategy(strategy: TruncationStrategy) -> None:
     _default_strategy = strategy
 
 
-def truncate_tool_output(output: str, tool_call: model.ToolCallItem | None = None) -> TruncationResult:
+def truncate_tool_output(output: str, tool_call: ToolCallLike | None = None) -> TruncationResult:
     """Truncate tool output using the current strategy."""
     return get_truncation_strategy().truncate(output, tool_call)

klaude_code/core/tool/web/mermaid_tool.py

@@ -7,11 +7,10 @@ from pathlib import Path
 
 from pydantic import BaseModel, Field
 
+from klaude_code.const import MERMAID_LIVE_PREFIX
 from klaude_code.core.tool.tool_abc import ToolABC, load_desc
 from klaude_code.core.tool.tool_registry import register
-from klaude_code.protocol import llm_param, model, tools
-
-_MERMAID_LIVE_PREFIX = "https://mermaid.live/view#pako:"
+from klaude_code.protocol import llm_param, message, model, tools
 
 
 @register(tools.MERMAID)
@@ -41,17 +40,17 @@ class MermaidTool(ToolABC):
         )
 
     @classmethod
-    async def call(cls, arguments: str) -> model.ToolResultItem:
+    async def call(cls, arguments: str) -> message.ToolResultMessage:
         try:
             args = cls.MermaidArguments.model_validate_json(arguments)
         except Exception as exc:  # pragma: no cover - defensive
-            return model.ToolResultItem(status="error", output=f"Invalid arguments: {exc}")
+            return message.ToolResultMessage(status="error", output_text=f"Invalid arguments: {exc}")
 
         link = cls._build_link(args.code)
         line_count = cls._count_lines(args.code)
         ui_extra = model.MermaidLinkUIExtra(code=args.code, link=link, line_count=line_count)
         output = f"Mermaid diagram rendered successfully ({line_count} lines)."
-        return model.ToolResultItem(status="success", output=output, ui_extra=ui_extra)
+        return message.ToolResultMessage(status="success", output_text=output, ui_extra=ui_extra)
 
     @staticmethod
     def _build_link(code: str) -> str:
@@ -64,7 +63,7 @@ class MermaidTool(ToolABC):
         json_payload = json.dumps(state, ensure_ascii=False)
         compressed = zlib.compress(json_payload.encode("utf-8"), level=9)
         encoded = base64.urlsafe_b64encode(compressed).decode("ascii").rstrip("=")
-        return f"{_MERMAID_LIVE_PREFIX}{encoded}"
+        return f"{MERMAID_LIVE_PREFIX}{encoded}"
 
     @staticmethod
     def _count_lines(code: str) -> int: