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

klaude_code/core/task.py

@@ -2,13 +2,13 @@ from __future__ import annotations
 
 import asyncio
 import time
-from collections.abc import AsyncGenerator, Callable, MutableMapping, Sequence
+from collections.abc import AsyncGenerator, Callable, Sequence
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
 from klaude_code import const
 from klaude_code.core.reminders import Reminder
-from klaude_code.core.tool import TodoContext, ToolABC
+from klaude_code.core.tool import FileTracker, TodoContext, ToolABC
 from klaude_code.core.turn import TurnError, TurnExecutionContext, TurnExecutor
 from klaude_code.protocol import events, model
 from klaude_code.trace import DebugType, log_debug
@@ -25,36 +25,35 @@ class MetadataAccumulator:
     """
 
     def __init__(self, model_name: str) -> None:
-        self._accumulated = model.ResponseMetadataItem(model_name=model_name)
+        self._main_agent = model.TaskMetadata(model_name=model_name)  # Main agent metadata
+        self._sub_agent_metadata: list[model.TaskMetadata] = []
         self._throughput_weighted_sum: float = 0.0
         self._throughput_tracked_tokens: int = 0
+        self._first_token_latency_sum: float = 0.0
+        self._first_token_latency_count: int = 0
+        self._turn_count: int = 0
 
     def add(self, turn_metadata: model.ResponseMetadataItem) -> None:
         """Merge a turn's metadata into the accumulated state."""
-        accumulated = self._accumulated
+        self._turn_count += 1
         usage = turn_metadata.usage
 
         if usage is not None:
-            if accumulated.usage is None:
-                accumulated.usage = model.Usage()
-            acc_usage = accumulated.usage
-            acc_usage.input_tokens += usage.input_tokens
-            acc_usage.cached_tokens += usage.cached_tokens
-            acc_usage.reasoning_tokens += usage.reasoning_tokens
-            acc_usage.output_tokens += usage.output_tokens
-            acc_usage.total_tokens += usage.total_tokens
-
-            if usage.context_usage_percent is not None:
-                acc_usage.context_usage_percent = usage.context_usage_percent
+            if self._main_agent.usage is None:
+                self._main_agent.usage = model.Usage()
+            acc_usage = self._main_agent.usage
+
+            model.TaskMetadata.merge_usage(acc_usage, usage)
+            acc_usage.currency = usage.currency
+
+            if usage.context_size is not None:
+                acc_usage.context_size = usage.context_size
+            if usage.context_limit is not None:
+                acc_usage.context_limit = usage.context_limit
 
             if usage.first_token_latency_ms is not None:
-                if acc_usage.first_token_latency_ms is None:
-                    acc_usage.first_token_latency_ms = usage.first_token_latency_ms
-                else:
-                    acc_usage.first_token_latency_ms = min(
-                        acc_usage.first_token_latency_ms,
-                        usage.first_token_latency_ms,
-                    )
+                self._first_token_latency_sum += usage.first_token_latency_ms
+                self._first_token_latency_count += 1
 
             if usage.throughput_tps is not None:
                 current_output = usage.output_tokens
@@ -62,53 +61,58 @@ class MetadataAccumulator:
                     self._throughput_weighted_sum += usage.throughput_tps * current_output
                     self._throughput_tracked_tokens += current_output
 
-            # Accumulate costs
-            if usage.input_cost is not None:
-                acc_usage.input_cost = (acc_usage.input_cost or 0.0) + usage.input_cost
-            if usage.output_cost is not None:
-                acc_usage.output_cost = (acc_usage.output_cost or 0.0) + usage.output_cost
-            if usage.cache_read_cost is not None:
-                acc_usage.cache_read_cost = (acc_usage.cache_read_cost or 0.0) + usage.cache_read_cost
-            if usage.total_cost is not None:
-                acc_usage.total_cost = (acc_usage.total_cost or 0.0) + usage.total_cost
-
         if turn_metadata.provider is not None:
-            accumulated.provider = turn_metadata.provider
+            self._main_agent.provider = turn_metadata.provider
         if turn_metadata.model_name:
-            accumulated.model_name = turn_metadata.model_name
-        if turn_metadata.response_id:
-            accumulated.response_id = turn_metadata.response_id
-        if turn_metadata.status is not None:
-            accumulated.status = turn_metadata.status
-        if turn_metadata.error_reason is not None:
-            accumulated.error_reason = turn_metadata.error_reason
-
-    def finalize(self, task_duration_s: float) -> model.ResponseMetadataItem:
+            self._main_agent.model_name = turn_metadata.model_name
+
+    def add_sub_agent_metadata(self, sub_agent_metadata: model.TaskMetadata) -> None:
+        """Add sub-agent task metadata to the accumulated state."""
+        self._sub_agent_metadata.append(sub_agent_metadata)
+
+    def finalize(self, task_duration_s: float) -> model.TaskMetadataItem:
         """Return the final accumulated metadata with computed throughput and duration."""
-        accumulated = self._accumulated
-        if accumulated.usage is not None:
+        if self._main_agent.usage is not None:
             if self._throughput_tracked_tokens > 0:
-                accumulated.usage.throughput_tps = self._throughput_weighted_sum / self._throughput_tracked_tokens
+                self._main_agent.usage.throughput_tps = self._throughput_weighted_sum / self._throughput_tracked_tokens
+            else:
+                self._main_agent.usage.throughput_tps = None
+
+            if self._first_token_latency_count > 0:
+                self._main_agent.usage.first_token_latency_ms = (
+                    self._first_token_latency_sum / self._first_token_latency_count
+                )
             else:
-                accumulated.usage.throughput_tps = None
+                self._main_agent.usage.first_token_latency_ms = None
 
-        accumulated.task_duration_s = task_duration_s
-        return accumulated
+        self._main_agent.task_duration_s = task_duration_s
+        self._main_agent.turn_count = self._turn_count
+        return model.TaskMetadataItem(main_agent=self._main_agent, sub_agent_task_metadata=self._sub_agent_metadata)
 
 
 @dataclass
-class TaskExecutionContext:
-    """Execution context required to run a task."""
+class SessionContext:
+    """Shared session-level context for task and turn execution.
+
+    Contains common fields that both TaskExecutionContext and TurnExecutionContext need.
+    """
 
     session_id: str
-    profile: AgentProfile
     get_conversation_history: Callable[[], list[model.ConversationItem]]
     append_history: Callable[[Sequence[model.ConversationItem]], None]
-    tool_registry: dict[str, type[ToolABC]]
-    file_tracker: MutableMapping[str, float]
+    file_tracker: FileTracker
     todo_context: TodoContext
+
+
+@dataclass
+class TaskExecutionContext:
+    """Execution context required to run a task."""
+
+    session_ctx: SessionContext
+    profile: AgentProfile
+    tool_registry: dict[str, type[ToolABC]]
     # For reminder processing - needs access to session
-    process_reminder: Callable[[Reminder], AsyncGenerator[events.DeveloperMessageEvent, None]]
+    process_reminder: Callable[[Reminder], AsyncGenerator[events.DeveloperMessageEvent]]
     sub_agent_state: model.SubAgentState | None
 
 
@@ -122,34 +126,45 @@ class TaskExecutor:
         self._context = context
         self._current_turn: TurnExecutor | None = None
         self._started_at: float = 0.0
+        self._metadata_accumulator: MetadataAccumulator | None = None
 
     @property
     def current_turn(self) -> TurnExecutor | None:
         return self._current_turn
 
     def cancel(self) -> list[events.Event]:
-        """Cancel the current turn and return any resulting events."""
+        """Cancel the current turn and return any resulting events including metadata."""
         ui_events: list[events.Event] = []
         if self._current_turn is not None:
             ui_events.extend(self._current_turn.cancel())
             self._current_turn = None
+
+        # Emit partial metadata on cancellation
+        if self._metadata_accumulator is not None and self._started_at > 0:
+            task_duration_s = time.perf_counter() - self._started_at
+            accumulated = self._metadata_accumulator.finalize(task_duration_s)
+            if accumulated.main_agent.usage is not None:
+                session_id = self._context.session_ctx.session_id
+                ui_events.append(events.TaskMetadataEvent(metadata=accumulated, session_id=session_id))
+                self._context.session_ctx.append_history([accumulated])
+
         return ui_events
 
-    async def run(self, user_input: model.UserInputPayload) -> AsyncGenerator[events.Event, None]:
+    async def run(self, user_input: model.UserInputPayload) -> AsyncGenerator[events.Event]:
         """Execute the task, yielding events as they occur."""
         ctx = self._context
+        session_ctx = ctx.session_ctx
         self._started_at = time.perf_counter()
 
         yield events.TaskStartEvent(
-            session_id=ctx.session_id,
+            session_id=session_ctx.session_id,
             sub_agent_state=ctx.sub_agent_state,
         )
-
-        ctx.append_history([model.UserMessageItem(content=user_input.text, images=user_input.images)])
+        del user_input  # Persisted by the operation handler before launching the task.
 
         profile = ctx.profile
-        metadata_accumulator = MetadataAccumulator(model_name=profile.llm_client.model_name)
-        last_assistant_message: events.AssistantMessageEvent | None = None
+        self._metadata_accumulator = MetadataAccumulator(model_name=profile.llm_client.model_name)
+        metadata_accumulator = self._metadata_accumulator
 
         while True:
             # Process reminders at the start of each turn
@@ -158,15 +173,11 @@ class TaskExecutor:
                     yield event
 
             turn_context = TurnExecutionContext(
-                session_id=ctx.session_id,
-                get_conversation_history=ctx.get_conversation_history,
-                append_history=ctx.append_history,
+                session_ctx=session_ctx,
                 llm_client=profile.llm_client,
                 system_prompt=profile.system_prompt,
                 tools=profile.tools,
                 tool_registry=ctx.tool_registry,
-                file_tracker=ctx.file_tracker,
-                todo_context=ctx.todo_context,
             )
 
             turn: TurnExecutor | None = None
@@ -181,11 +192,22 @@ class TaskExecutor:
                     async for turn_event in turn.run():
                         match turn_event:
                             case events.AssistantMessageEvent() as am:
-                                if am.content.strip() != "":
-                                    last_assistant_message = am
                                 yield am
                             case events.ResponseMetadataEvent() as e:
                                 metadata_accumulator.add(e.metadata)
+                                # Emit context usage event if available
+                                if e.metadata.usage is not None:
+                                    context_percent = e.metadata.usage.context_usage_percent
+                                    if context_percent is not None:
+                                        yield events.ContextUsageEvent(
+                                            session_id=session_ctx.session_id,
+                                            context_percent=context_percent,
+                                        )
+                            case events.ToolResultEvent() as e:
+                                # Collect sub-agent task metadata from tool results
+                                if e.task_metadata is not None:
+                                    metadata_accumulator.add_sub_agent_metadata(e.task_metadata)
+                                yield turn_event
                             case _:
                                 yield turn_event
 
@@ -198,7 +220,9 @@ class TaskExecutor:
                         error_msg = f"Retrying {attempt + 1}/{const.MAX_FAILED_TURN_RETRIES} in {delay:.1f}s"
                         if last_error_message:
                             error_msg = f"{error_msg} - {last_error_message}"
-                        yield events.ErrorEvent(error_message=error_msg, can_retry=True)
+                        yield events.ErrorEvent(
+                            error_message=error_msg, can_retry=True, session_id=session_ctx.session_id
+                        )
                         await asyncio.sleep(delay)
                 finally:
                     self._current_turn = None
@@ -212,21 +236,35 @@ class TaskExecutor:
                 final_error = f"Turn failed after {const.MAX_FAILED_TURN_RETRIES} retries."
                 if last_error_message:
                     final_error = f"{last_error_message}\n{final_error}"
-                yield events.ErrorEvent(error_message=final_error, can_retry=False)
+                yield events.ErrorEvent(error_message=final_error, can_retry=False, session_id=session_ctx.session_id)
                 return
 
-            if turn is None or not turn.has_tool_call:
+            if turn is None or turn.task_finished:
+                # Empty result should retry instead of finishing
+                if turn is not None and not turn.task_result.strip():
+                    if ctx.sub_agent_state is not None:
+                        error_msg = "Sub-agent returned empty result, retrying..."
+                    else:
+                        error_msg = "Agent returned empty result, retrying..."
+                    yield events.ErrorEvent(error_message=error_msg, can_retry=True, session_id=session_ctx.session_id)
+                    continue
                 break
 
         # Finalize metadata
         task_duration_s = time.perf_counter() - self._started_at
         accumulated = metadata_accumulator.finalize(task_duration_s)
 
-        yield events.ResponseMetadataEvent(metadata=accumulated, session_id=ctx.session_id)
-        ctx.append_history([accumulated])
+        yield events.TaskMetadataEvent(metadata=accumulated, session_id=session_ctx.session_id)
+        session_ctx.append_history([accumulated])
+
+        # Get task result from turn
+        task_result = turn.task_result if turn is not None else ""
+        has_structured_output = turn.has_structured_output if turn is not None else False
+
         yield events.TaskFinishEvent(
-            session_id=ctx.session_id,
-            task_result=last_assistant_message.content if last_assistant_message else "",
+            session_id=session_ctx.session_id,
+            task_result=task_result,
+            has_structured_output=has_structured_output,
         )
 
 

klaude_code/core/tool/__init__.py

@@ -1,21 +1,22 @@
 from .file.apply_patch import DiffError, process_patch
 from .file.apply_patch_tool import ApplyPatchTool
 from .file.edit_tool import EditTool
-from .file.multi_edit_tool import MultiEditTool
+from .file.move_tool import MoveTool
 from .file.read_tool import ReadTool
 from .file.write_tool import WriteTool
-from .memory.memory_tool import MEMORY_DIR_NAME, MemoryTool
-from .memory.skill_loader import Skill, SkillLoader
-from .memory.skill_tool import SkillTool
+from .report_back_tool import ReportBackTool
 from .shell.bash_tool import BashTool
 from .shell.command_safety import SafetyCheckResult, is_safe_command
+from .skill.skill_tool import SkillTool
 from .sub_agent_tool import SubAgentTool
 from .todo.todo_write_tool import TodoWriteTool
 from .todo.update_plan_tool import UpdatePlanTool
 from .tool_abc import ToolABC
 from .tool_context import (
+    FileTracker,
     TodoContext,
     ToolContextToken,
+    build_todo_context,
     current_run_subtask_callback,
     reset_tool_context,
     set_tool_context_from_session,
@@ -26,50 +27,42 @@ from .tool_runner import run_tool
 from .truncation import SimpleTruncationStrategy, TruncationStrategy, get_truncation_strategy, set_truncation_strategy
 from .web.mermaid_tool import MermaidTool
 from .web.web_fetch_tool import WebFetchTool
+from .web.web_search_tool import WebSearchTool
 
 __all__ = [
-    # Tools
     "ApplyPatchTool",
     "BashTool",
+    "DiffError",
     "EditTool",
-    "MemoryTool",
+    "FileTracker",
     "MermaidTool",
-    "MultiEditTool",
+    "MoveTool",
     "ReadTool",
+    "ReportBackTool",
+    "SafetyCheckResult",
+    "SimpleTruncationStrategy",
     "SkillTool",
     "SubAgentTool",
+    "TodoContext",
     "TodoWriteTool",
+    "ToolABC",
+    "ToolContextToken",
+    "TruncationStrategy",
     "UpdatePlanTool",
     "WebFetchTool",
+    "WebSearchTool",
     "WriteTool",
-    # Tool ABC
-    "ToolABC",
-    # Tool context
-    "TodoContext",
-    "ToolContextToken",
+    "build_todo_context",
     "current_run_subtask_callback",
-    "reset_tool_context",
-    "set_tool_context_from_session",
-    "tool_context",
-    # Tool registry
-    "load_agent_tools",
     "get_registry",
     "get_tool_schemas",
-    "run_tool",
-    # Truncation
-    "SimpleTruncationStrategy",
-    "TruncationStrategy",
     "get_truncation_strategy",
-    "set_truncation_strategy",
-    # Command safety
-    "SafetyCheckResult",
     "is_safe_command",
-    # Skill
-    "Skill",
-    "SkillLoader",
-    # Memory
-    "MEMORY_DIR_NAME",
-    # Apply patch
-    "DiffError",
+    "load_agent_tools",
     "process_patch",
+    "reset_tool_context",
+    "run_tool",
+    "set_tool_context_from_session",
+    "set_truncation_strategy",
+    "tool_context",
 ]

klaude_code/core/tool/file/_utils.py

@@ -0,0 +1,36 @@
+"""Shared utility functions for file tools."""
+
+from __future__ import annotations
+
+import hashlib
+import os
+from pathlib import Path
+
+
+def is_directory(path: str) -> bool:
+    """Check if path is a directory."""
+    return os.path.isdir(path)
+
+
+def file_exists(path: str) -> bool:
+    """Check if path exists."""
+    return os.path.exists(path)
+
+
+def read_text(path: str) -> str:
+    """Read text from file with UTF-8 encoding."""
+    with open(path, encoding="utf-8", errors="replace") as f:
+        return f.read()
+
+
+def write_text(path: str, content: str) -> None:
+    """Write text to file, creating parent directories if needed."""
+    parent = Path(path).parent
+    parent.mkdir(parents=True, exist_ok=True)
+    with open(path, "w", encoding="utf-8") as f:
+        f.write(content)
+
+
+def hash_text_sha256(content: str) -> str:
+    """Return SHA-256 for the given text content encoded as UTF-8."""
+    return hashlib.sha256(content.encode("utf-8")).hexdigest()

klaude_code/core/tool/file/apply_patch.py

@@ -3,8 +3,8 @@ https://github.com/openai/openai-cookbook/blob/main/examples/gpt-5/apply_patch.p
 """
 
 import os
+from collections.abc import Callable
 from enum import Enum
-from typing import Callable, Optional
 
 from pydantic import BaseModel, Field
 
@@ -17,16 +17,16 @@ class ActionType(str, Enum):
 
 class FileChange(BaseModel):
     type: ActionType
-    old_content: Optional[str] = None
-    new_content: Optional[str] = None
-    move_path: Optional[str] = None
+    old_content: str | None = None
+    new_content: str | None = None
+    move_path: str | None = None
 
 
 class Commit(BaseModel):
     changes: dict[str, FileChange] = Field(default_factory=dict)
 
 
-def assemble_changes(orig: dict[str, Optional[str]], dest: dict[str, Optional[str]]) -> Commit:
+def assemble_changes(orig: dict[str, str | None], dest: dict[str, str | None]) -> Commit:
     commit = Commit()
     for path in sorted(set(orig.keys()).union(dest.keys())):
         old_content = orig.get(path)
@@ -49,7 +49,7 @@ def assemble_changes(orig: dict[str, Optional[str]], dest: dict[str, Optional[st
                     old_content=old_content,
                 )
             else:
-                assert False
+                raise AssertionError()
     return commit
 
 
@@ -71,9 +71,9 @@ def _new_chunk_list() -> list["Chunk"]:
 
 class PatchAction(BaseModel):
     type: ActionType
-    new_file: Optional[str] = None
+    new_file: str | None = None
     chunks: list[Chunk] = Field(default_factory=_new_chunk_list)
-    move_path: Optional[str] = None
+    move_path: str | None = None
 
 
 class Patch(BaseModel):
@@ -87,26 +87,19 @@ class Parser(BaseModel):
     patch: Patch = Field(default_factory=Patch)
     fuzz: int = 0
 
-    def is_done(self, prefixes: Optional[tuple[str, ...]] = None) -> bool:
+    def is_done(self, prefixes: tuple[str, ...] | None = None) -> bool:
         if self.index >= len(self.lines):
             return True
-        if prefixes and self.lines[self.index].startswith(prefixes):
-            return True
-        return False
+        return bool(prefixes and self.lines[self.index].startswith(prefixes))
 
     def startswith(self, prefix: str | tuple[str, ...]) -> bool:
         assert self.index < len(self.lines), f"Index: {self.index} >= {len(self.lines)}"
-        if self.lines[self.index].startswith(prefix):
-            return True
-        return False
+        return self.lines[self.index].startswith(prefix)
 
     def read_str(self, prefix: str = "", return_everything: bool = False) -> str:
         assert self.index < len(self.lines), f"Index: {self.index} >= {len(self.lines)}"
         if self.lines[self.index].startswith(prefix):
-            if return_everything:
-                text = self.lines[self.index]
-            else:
-                text = self.lines[self.index][len(prefix) :]
+            text = self.lines[self.index] if return_everything else self.lines[self.index][len(prefix) :]
             self.index += 1
             return text
         return ""
@@ -167,10 +160,9 @@ class Parser(BaseModel):
         ):
             def_str = self.read_str("@@ ")
             section_str = ""
-            if not def_str:
-                if self.lines[self.index] == "@@":
-                    section_str = self.lines[self.index]
-                    self.index += 1
+            if not def_str and self.lines[self.index] == "@@":
+                section_str = self.lines[self.index]
+                self.index += 1
             if not (def_str or section_str or index == 0):
                 raise DiffError(f"Invalid Line:\n{self.lines[self.index]}")
             if def_str.strip():
@@ -457,7 +449,7 @@ def process_patch(
 
 
 def open_file(path: str) -> str:
-    with open(path, "rt") as f:
+    with open(path) as f:
         return f.read()
 
 
@@ -465,7 +457,7 @@ def write_file(path: str, content: str) -> None:
     if "/" in path:
         parent = "/".join(path.split("/")[:-1])
         os.makedirs(parent, exist_ok=True)
-    with open(path, "wt") as f:
+    with open(path, "w") as f:
         f.write(content)