klaude-code 1.2.1__py3-none-any.whl → 1.2.3__py3-none-any.whl

klaude_code/core/reminders.py

@@ -1,15 +1,11 @@
 import json
-import re
 from pathlib import Path
 from typing import Awaitable, Callable
 
 from pydantic import BaseModel
 
-from klaude_code.const import TODO_REMINDER_TOOL_CALL_THRESHOLD
-from klaude_code.core.clipboard_manifest import load_latest_clipboard_manifest, next_session_token
-from klaude_code.core.tool.file.read_tool import ReadTool
-from klaude_code.core.tool.shell.bash_tool import BashTool
-from klaude_code.core.tool.tool_context import reset_tool_context, set_tool_context_from_session
+from klaude_code import const
+from klaude_code.core.tool import BashTool, ReadTool, reset_tool_context, set_tool_context_from_session
 from klaude_code.protocol import model, tools
 from klaude_code.session import Session
 
@@ -30,7 +26,9 @@ def get_last_new_user_input(session: Session) -> str | None:
     return "\n\n".join(result)
 
 
-async def at_file_reader_reminder(session: Session) -> model.DeveloperMessageItem | None:
+async def at_file_reader_reminder(
+    session: Session,
+) -> model.DeveloperMessageItem | None:
     """Parse @foo/bar to read"""
     last_user_input = get_last_new_user_input(session)
     if not last_user_input or "@" not in last_user_input.strip():
@@ -124,7 +122,9 @@ async def empty_todo_reminder(session: Session) -> model.DeveloperMessageItem |
     return None
 
 
-async def todo_not_used_recently_reminder(session: Session) -> model.DeveloperMessageItem | None:
+async def todo_not_used_recently_reminder(
+    session: Session,
+) -> model.DeveloperMessageItem | None:
     """Remind agent to use TodoWrite tool if it hasn't been used recently (>=10 other tool calls), with cooldown.
 
     Cooldown behavior:
@@ -147,10 +147,10 @@ async def todo_not_used_recently_reminder(session: Session) -> model.DeveloperMe
             if item.name in (tools.TODO_WRITE, tools.UPDATE_PLAN):
                 break
             other_tool_call_count_befor_last_todo += 1
-            if other_tool_call_count_befor_last_todo >= TODO_REMINDER_TOOL_CALL_THRESHOLD:
+            if other_tool_call_count_befor_last_todo >= const.TODO_REMINDER_TOOL_CALL_THRESHOLD:
                 break
 
-    not_used_recently = other_tool_call_count_befor_last_todo >= TODO_REMINDER_TOOL_CALL_THRESHOLD
+    not_used_recently = other_tool_call_count_befor_last_todo >= const.TODO_REMINDER_TOOL_CALL_THRESHOLD
 
     if not not_used_recently:
         return None
@@ -173,7 +173,9 @@ Here are the existing contents of your todo list:
     return None
 
 
-async def file_changed_externally_reminder(session: Session) -> model.DeveloperMessageItem | None:
+async def file_changed_externally_reminder(
+    session: Session,
+) -> model.DeveloperMessageItem | None:
     """Remind agent about user/linter' changes to the files in FileTracker, provding the newest content of the file."""
     changed_files: list[tuple[str, str, list[model.ImageURLPart] | None]] = []
     collected_images: list[model.ImageURLPart] = []
@@ -192,7 +194,13 @@ async def file_changed_externally_reminder(session: Session) -> model.DeveloperM
                                 collected_images.extend(tool_result.images)
                     finally:
                         reset_tool_context(context_token)
-            except (FileNotFoundError, IsADirectoryError, OSError, PermissionError, UnicodeDecodeError):
+            except (
+                FileNotFoundError,
+                IsADirectoryError,
+                OSError,
+                PermissionError,
+                UnicodeDecodeError,
+            ):
                 continue
     if len(changed_files) > 0:
         changed_files_str = "\n\n".join(
@@ -213,8 +221,14 @@ async def file_changed_externally_reminder(session: Session) -> model.DeveloperM
 
 def get_memory_paths() -> list[tuple[Path, str]]:
     return [
-        (Path.home() / ".claude" / "CLAUDE.md", "user's private global instructions for all projects"),
-        (Path.home() / ".codex" / "AGENTS.md", "user's private global instructions for all projects"),
+        (
+            Path.home() / ".claude" / "CLAUDE.md",
+            "user's private global instructions for all projects",
+        ),
+        (
+            Path.home() / ".codex" / "AGENTS.md",
+            "user's private global instructions for all projects",
+        ),
         (Path.cwd() / "AGENTS.md", "project instructions, checked into the codebase"),
         (Path.cwd() / "AGENT.md", "project instructions, checked into the codebase"),
         (Path.cwd() / "CLAUDE.md", "project instructions, checked into the codebase"),
@@ -227,6 +241,28 @@ class Memory(BaseModel):
     content: str
 
 
+def get_last_user_message_image_count(session: Session) -> int:
+    """Get image count from the last user message in conversation history."""
+    for item in reversed(session.conversation_history):
+        if isinstance(item, model.ToolResultItem):
+            return 0
+        if isinstance(item, model.UserMessageItem):
+            return len(item.images) if item.images else 0
+    return 0
+
+
+async def image_reminder(session: Session) -> model.DeveloperMessageItem | None:
+    """Remind agent about images attached by user in the last message."""
+    image_count = get_last_user_message_image_count(session)
+    if image_count == 0:
+        return None
+
+    return model.DeveloperMessageItem(
+        content=f"<system-reminder>User attached {image_count} image{'s' if image_count > 1 else ''} in their message. Make sure to analyze and reference these images as needed.</system-reminder>",
+        user_image_count=image_count,
+    )
+
+
 async def memory_reminder(session: Session) -> model.DeveloperMessageItem | None:
     """CLAUDE.md AGENTS.md"""
     memory_paths = get_memory_paths()
@@ -268,7 +304,14 @@ def get_last_turn_tool_call(session: Session) -> list[model.ToolCallItem]:
     for item in reversed(session.conversation_history):
         if isinstance(item, model.ToolCallItem):
             tool_calls.append(item)
-        if isinstance(item, (model.ReasoningEncryptedItem, model.ReasoningTextItem, model.AssistantMessageItem)):
+        if isinstance(
+            item,
+            (
+                model.ReasoningEncryptedItem,
+                model.ReasoningTextItem,
+                model.AssistantMessageItem,
+            ),
+        ):
             break
     return tool_calls
 
@@ -276,7 +319,9 @@ def get_last_turn_tool_call(session: Session) -> list[model.ToolCallItem]:
 MEMORY_FILE_NAMES = ["CLAUDE.md", "AGENTS.md", "AGENT.md"]
 
 
-async def last_path_memory_reminder(session: Session) -> model.DeveloperMessageItem | None:
+async def last_path_memory_reminder(
+    session: Session,
+) -> model.DeveloperMessageItem | None:
     """When last turn tool call entered a directory (or parent directory) with CLAUDE.md AGENTS.md"""
     tool_calls = get_last_turn_tool_call(session)
     if len(tool_calls) == 0:
@@ -356,63 +401,6 @@ async def last_path_memory_reminder(session: Session) -> model.DeveloperMessageI
         )
 
 
-async def clipboard_image_reminder(session: Session) -> model.DeveloperMessageItem | None:
-    """Parse [Image #N] and attach images from clipboard history."""
-    last_user_input = get_last_new_user_input(session)
-    if not last_user_input or "[Image #" not in last_user_input:
-        return None
-
-    manifest = load_latest_clipboard_manifest()
-    if manifest is None:
-        return None
-    manifest_source = manifest.source_id
-    current_source = next_session_token()
-    if manifest_source and manifest_source != current_source:
-        return None
-    image_map = manifest.tag_map()
-    if not image_map:
-        return None
-
-    collected_images: list[model.ImageURLPart] = []
-
-    # Find all tokens
-    # Regex for [Image #(\d+)]
-    matches = re.findall(r"\[Image #(\d+)\]", last_user_input)
-    requested_tags = [f"[Image #{num}]" for num in matches]
-
-    processed_paths: set[str] = set()
-
-    attached_tags: list[str] = []
-
-    for tag in requested_tags:
-        if tag in image_map:
-            path = image_map[tag]
-            if path in processed_paths:
-                continue
-
-            context_token = set_tool_context_from_session(session)
-            try:
-                # We use ReadTool to get the image object in the correct format
-                # This assumes ReadTool handles image files correctly
-                args = ReadTool.ReadArguments(file_path=path)
-                tool_result = await ReadTool.call_with_args(args)
-                if tool_result.images:
-                    collected_images.extend(tool_result.images)
-                    processed_paths.add(path)
-                    attached_tags.append(tag)
-            finally:
-                reset_tool_context(context_token)
-
-    if not collected_images:
-        return None
-
-    return model.DeveloperMessageItem(
-        content="",
-        images=collected_images,
-        clipboard_images=attached_tags,
-    )
-
-
 ALL_REMINDERS = [
     empty_todo_reminder,
     todo_not_used_recently_reminder,
@@ -420,19 +408,27 @@ ALL_REMINDERS = [
     memory_reminder,
     last_path_memory_reminder,
     at_file_reader_reminder,
-    clipboard_image_reminder,
+    image_reminder,
 ]
 
 
-def get_vanilla_reminders() -> list[Reminder]:
-    return [at_file_reader_reminder, clipboard_image_reminder]
+def load_agent_reminders(
+    model_name: str, sub_agent_type: str | None = None, *, vanilla: bool = False
+) -> list[Reminder]:
+    """Get reminders for an agent based on model and agent type.
 
+    Args:
+        model_name: The model name.
+        sub_agent_type: If None, returns main agent reminders. Otherwise returns sub-agent reminders.
+        vanilla: If True, returns minimal vanilla reminders (ignores sub_agent_type).
+    """
+    if vanilla:
+        return [at_file_reader_reminder]
 
-def get_main_agent_reminders(model_name: str) -> list[Reminder]:
     reminders: list[Reminder] = []
 
-    # For GPT-5, we do not show empty todo and todo not used recently reminders
-    if "gpt-5" not in model_name:
+    # Only main agent (not sub-agent) gets todo reminders, and not for GPT-5
+    if sub_agent_type is None and "gpt-5" not in model_name:
         reminders.append(empty_todo_reminder)
         reminders.append(todo_not_used_recently_reminder)
 
@@ -441,23 +437,8 @@ def get_main_agent_reminders(model_name: str) -> list[Reminder]:
             memory_reminder,
             last_path_memory_reminder,
             at_file_reader_reminder,
-            clipboard_image_reminder,
-            file_changed_externally_reminder,
-        ]
-    )
-
-    return reminders
-
-
-def get_sub_agent_reminders(model_name: str) -> list[Reminder]:
-    reminders: list[Reminder] = []
-    reminders.extend(
-        [
-            memory_reminder,
-            last_path_memory_reminder,
-            at_file_reader_reminder,
-            clipboard_image_reminder,
             file_changed_externally_reminder,
+            image_reminder,
         ]
     )
 

klaude_code/core/task.py

@@ -6,10 +6,9 @@ from collections.abc import AsyncGenerator, Callable, MutableMapping, Sequence
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
-from klaude_code.const import INITIAL_RETRY_DELAY_S, MAX_FAILED_TURN_RETRIES, MAX_RETRY_DELAY_S
+from klaude_code import const
 from klaude_code.core.reminders import Reminder
-from klaude_code.core.tool.tool_abc import ToolABC
-from klaude_code.core.tool.tool_context import TodoContext
+from klaude_code.core.tool import 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
@@ -63,6 +62,16 @@ 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
         if turn_metadata.model_name:
@@ -79,9 +88,7 @@ class MetadataAccumulator:
         accumulated = self._accumulated
         if accumulated.usage is not None:
             if self._throughput_tracked_tokens > 0:
-                accumulated.usage.throughput_tps = (
-                    self._throughput_weighted_sum / self._throughput_tracked_tokens
-                )
+                accumulated.usage.throughput_tps = self._throughput_weighted_sum / self._throughput_tracked_tokens
             else:
                 accumulated.usage.throughput_tps = None
 
@@ -128,7 +135,7 @@ class TaskExecutor:
             self._current_turn = None
         return ui_events
 
-    async def run(self, user_input: str) -> AsyncGenerator[events.Event, None]:
+    async def run(self, user_input: model.UserInputPayload) -> AsyncGenerator[events.Event, None]:
         """Execute the task, yielding events as they occur."""
         ctx = self._context
         self._started_at = time.perf_counter()
@@ -138,7 +145,7 @@ class TaskExecutor:
             sub_agent_state=ctx.sub_agent_state,
         )
 
-        ctx.append_history([model.UserMessageItem(content=user_input)])
+        ctx.append_history([model.UserMessageItem(content=user_input.text, images=user_input.images)])
 
         profile = ctx.profile
         metadata_accumulator = MetadataAccumulator(model_name=profile.llm_client.model_name)
@@ -166,7 +173,7 @@ class TaskExecutor:
             turn_succeeded = False
             last_error_message: str | None = None
 
-            for attempt in range(MAX_FAILED_TURN_RETRIES + 1):
+            for attempt in range(const.MAX_FAILED_TURN_RETRIES + 1):
                 turn = TurnExecutor(turn_context)
                 self._current_turn = turn
 
@@ -186,9 +193,9 @@ class TaskExecutor:
                     break
                 except TurnError as e:
                     last_error_message = str(e)
-                    if attempt < MAX_FAILED_TURN_RETRIES:
+                    if attempt < const.MAX_FAILED_TURN_RETRIES:
                         delay = _retry_delay_seconds(attempt + 1)
-                        error_msg = f"Retrying {attempt + 1}/{MAX_FAILED_TURN_RETRIES} in {delay:.1f}s"
+                        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)
@@ -202,7 +209,7 @@ class TaskExecutor:
                     style="red",
                     debug_type=DebugType.EXECUTION,
                 )
-                final_error = f"Turn failed after {MAX_FAILED_TURN_RETRIES} retries."
+                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)
@@ -226,5 +233,5 @@ class TaskExecutor:
 def _retry_delay_seconds(attempt: int) -> float:
     """Compute exponential backoff delay for the given attempt count."""
     capped_attempt = max(1, attempt)
-    delay = INITIAL_RETRY_DELAY_S * (2 ** (capped_attempt - 1))
-    return min(delay, MAX_RETRY_DELAY_S)
+    delay = const.INITIAL_RETRY_DELAY_S * (2 ** (capped_attempt - 1))
+    return min(delay, const.MAX_RETRY_DELAY_S)

klaude_code/core/tool/__init__.py

@@ -1,41 +1,75 @@
+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.read_tool import ReadTool
 from .file.write_tool import WriteTool
-from .memory.memory_tool import MemoryTool
+from .memory.memory_tool import MEMORY_DIR_NAME, MemoryTool
+from .memory.skill_loader import Skill, SkillLoader
 from .memory.skill_tool import SkillTool
 from .shell.bash_tool import BashTool
+from .shell.command_safety import SafetyCheckResult, is_safe_command
 from .sub_agent_tool import SubAgentTool
 from .todo.todo_write_tool import TodoWriteTool
 from .todo.update_plan_tool import UpdatePlanTool
-from .tool_registry import get_main_agent_tools, get_registry, get_sub_agent_tools, get_tool_schemas
+from .tool_abc import ToolABC
+from .tool_context import (
+    TodoContext,
+    ToolContextToken,
+    current_run_subtask_callback,
+    reset_tool_context,
+    set_tool_context_from_session,
+    tool_context,
+)
+from .tool_registry import get_registry, get_tool_schemas, load_agent_tools
 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
 
 __all__ = [
+    # Tools
+    "ApplyPatchTool",
     "BashTool",
-    "ReadTool",
     "EditTool",
     "MemoryTool",
+    "MermaidTool",
     "MultiEditTool",
+    "ReadTool",
+    "SkillTool",
     "SubAgentTool",
     "TodoWriteTool",
-    "WriteTool",
-    "SkillTool",
     "UpdatePlanTool",
-    "ApplyPatchTool",
-    "MermaidTool",
     "WebFetchTool",
-    "get_tool_schemas",
+    "WriteTool",
+    # Tool ABC
+    "ToolABC",
+    # Tool context
+    "TodoContext",
+    "ToolContextToken",
+    "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",
-    "get_sub_agent_tools",
-    "get_main_agent_tools",
-    "TruncationStrategy",
+    # 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",
+    "process_patch",
 ]

klaude_code/core/tool/file/apply_patch.py

@@ -417,7 +417,11 @@ def load_files(paths: list[str], open_fn: Callable[[str], str]) -> dict[str, str
     return orig
 
 
-def apply_commit(commit: Commit, write_fn: Callable[[str, str], None], remove_fn: Callable[[str], None]) -> None:
+def apply_commit(
+    commit: Commit,
+    write_fn: Callable[[str, str], None],
+    remove_fn: Callable[[str], None],
+) -> None:
     for path, change in commit.changes.items():
         if change.type == ActionType.DELETE:
             remove_fn(path)

klaude_code/core/tool/file/apply_patch_tool.py

@@ -11,24 +11,22 @@ from klaude_code.core.tool.file import apply_patch as apply_patch_module
 from klaude_code.core.tool.tool_abc import ToolABC, load_desc
 from klaude_code.core.tool.tool_context import get_current_file_tracker
 from klaude_code.core.tool.tool_registry import register
-from klaude_code.protocol import tools
-from klaude_code.protocol.llm_parameter import ToolSchema
-from klaude_code.protocol.model import ToolResultItem, ToolResultUIExtra, ToolResultUIExtraType
+from klaude_code.protocol import llm_param, model, tools
 
 
 class ApplyPatchHandler:
     @classmethod
-    async def handle_apply_patch(cls, patch_text: str) -> ToolResultItem:
+    async def handle_apply_patch(cls, patch_text: str) -> model.ToolResultItem:
         try:
             output, diff_text = await asyncio.to_thread(cls._apply_patch_in_thread, patch_text)
         except apply_patch_module.DiffError as error:
-            return ToolResultItem(status="error", output=str(error))
+            return model.ToolResultItem(status="error", output=str(error))
         except Exception as error:  # pragma: no cover  # unexpected errors bubbled to tool result
-            return ToolResultItem(status="error", output=f"Execution error: {error}")
-        return ToolResultItem(
+            return model.ToolResultItem(status="error", output=f"Execution error: {error}")
+        return model.ToolResultItem(
             status="success",
             output=output,
-            ui_extra=ToolResultUIExtra(type=ToolResultUIExtraType.DIFF_TEXT, diff_text=diff_text),
+            ui_extra=model.ToolResultUIExtra(type=model.ToolResultUIExtraType.DIFF_TEXT, diff_text=diff_text),
         )
 
     @staticmethod
@@ -176,8 +174,8 @@ class ApplyPatchTool(ToolABC):
         patch: str
 
     @classmethod
-    def schema(cls) -> ToolSchema:
-        return ToolSchema(
+    def schema(cls) -> llm_param.ToolSchema:
+        return llm_param.ToolSchema(
             name=tools.APPLY_PATCH,
             type="function",
             description=load_desc(Path(__file__).parent / "apply_patch_tool.md"),
@@ -194,13 +192,13 @@ class ApplyPatchTool(ToolABC):
         )
 
     @classmethod
-    async def call(cls, arguments: str) -> ToolResultItem:
+    async def call(cls, arguments: str) -> model.ToolResultItem:
         try:
             args = cls.ApplyPatchArguments.model_validate_json(arguments)
         except ValueError as exc:
-            return ToolResultItem(status="error", output=f"Invalid arguments: {exc}")
+            return model.ToolResultItem(status="error", output=f"Invalid arguments: {exc}")
         return await cls.call_with_args(args)
 
     @classmethod
-    async def call_with_args(cls, args: ApplyPatchArguments) -> ToolResultItem:
+    async def call_with_args(cls, args: ApplyPatchArguments) -> model.ToolResultItem:
         return await ApplyPatchHandler.handle_apply_patch(args.patch)

klaude_code/core/tool/file/edit_tool.py

@@ -10,9 +10,7 @@ from pydantic import BaseModel, Field
 from klaude_code.core.tool.tool_abc import ToolABC, load_desc
 from klaude_code.core.tool.tool_context import get_current_file_tracker
 from klaude_code.core.tool.tool_registry import register
-from klaude_code.protocol.llm_parameter import ToolSchema
-from klaude_code.protocol.model import ToolResultItem, ToolResultUIExtra, ToolResultUIExtraType
-from klaude_code.protocol.tools import EDIT
+from klaude_code.protocol import llm_param, model, tools
 
 
 def _is_directory(path: str) -> bool:
@@ -41,7 +39,7 @@ def _write_text(path: str, content: str) -> None:
         f.write(content)
 
 
-@register(EDIT)
+@register(tools.EDIT)
 class EditTool(ToolABC):
     class EditArguments(BaseModel):
         file_path: str
@@ -50,9 +48,9 @@ class EditTool(ToolABC):
         replace_all: bool = Field(default=False)
 
     @classmethod
-    def schema(cls) -> ToolSchema:
-        return ToolSchema(
-            name=EDIT,
+    def schema(cls) -> llm_param.ToolSchema:
+        return llm_param.ToolSchema(
+            name=tools.EDIT,
             type="function",
             description=load_desc(Path(__file__).parent / "edit_tool.md"),
             parameters={
@@ -112,23 +110,23 @@ class EditTool(ToolABC):
         return content.replace(old_string, new_string, 1)
 
     @classmethod
-    async def call(cls, arguments: str) -> ToolResultItem:
+    async def call(cls, arguments: str) -> model.ToolResultItem:
         try:
             args = EditTool.EditArguments.model_validate_json(arguments)
         except Exception as e:  # pragma: no cover - defensive
-            return ToolResultItem(status="error", output=f"Invalid arguments: {e}")
+            return model.ToolResultItem(status="error", output=f"Invalid arguments: {e}")
 
         file_path = os.path.abspath(args.file_path)
 
         # Common file errors
         if _is_directory(file_path):
-            return ToolResultItem(
+            return model.ToolResultItem(
                 status="error",
                 output="<tool_use_error>Illegal operation on a directory. edit</tool_use_error>",
             )
 
         if args.old_string == "":
-            return ToolResultItem(
+            return model.ToolResultItem(
                 status="error",
                 output=(
                     "<tool_use_error>old_string must not be empty for Edit. "
@@ -140,14 +138,14 @@ class EditTool(ToolABC):
         file_tracker = get_current_file_tracker()
         if not _file_exists(file_path):
             # We require reading before editing
-            return ToolResultItem(
+            return model.ToolResultItem(
                 status="error",
                 output=("File has not been read yet. Read it first before writing to it."),
             )
         if file_tracker is not None:
             tracked = file_tracker.get(file_path)
             if tracked is None:
-                return ToolResultItem(
+                return model.ToolResultItem(
                     status="error",
                     output=("File has not been read yet. Read it first before writing to it."),
                 )
@@ -156,7 +154,7 @@ class EditTool(ToolABC):
             except Exception:
                 current_mtime = tracked
             if current_mtime != tracked:
-                return ToolResultItem(
+                return model.ToolResultItem(
                     status="error",
                     output=(
                         "File has been modified externally. Either by user or a linter. Read it first before writing to it."
@@ -167,24 +165,30 @@ class EditTool(ToolABC):
         try:
             before = await asyncio.to_thread(_read_text, file_path)
         except FileNotFoundError:
-            return ToolResultItem(
+            return model.ToolResultItem(
                 status="error",
                 output="File has not been read yet. Read it first before writing to it.",
             )
 
         err = cls.valid(
-            content=before, old_string=args.old_string, new_string=args.new_string, replace_all=args.replace_all
+            content=before,
+            old_string=args.old_string,
+            new_string=args.new_string,
+            replace_all=args.replace_all,
         )
         if err is not None:
-            return ToolResultItem(status="error", output=err)
+            return model.ToolResultItem(status="error", output=err)
 
         after = cls.execute(
-            content=before, old_string=args.old_string, new_string=args.new_string, replace_all=args.replace_all
+            content=before,
+            old_string=args.old_string,
+            new_string=args.new_string,
+            replace_all=args.replace_all,
         )
 
         # If nothing changed due to replacement semantics (should not happen after valid), guard anyway
         if before == after:
-            return ToolResultItem(
+            return model.ToolResultItem(
                 status="error",
                 output=(
                     "<tool_use_error>No changes to make: old_string and new_string are exactly the same.</tool_use_error>"
@@ -195,7 +199,7 @@ class EditTool(ToolABC):
         try:
             await asyncio.to_thread(_write_text, file_path, after)
         except Exception as e:  # pragma: no cover
-            return ToolResultItem(status="error", output=f"<tool_use_error>{e}</tool_use_error>")
+            return model.ToolResultItem(status="error", output=f"<tool_use_error>{e}</tool_use_error>")
 
         # Prepare UI extra: unified diff with 3 context lines
         diff_lines = list(
@@ -208,7 +212,7 @@ class EditTool(ToolABC):
             )
         )
         diff_text = "\n".join(diff_lines)
-        ui_extra = ToolResultUIExtra(type=ToolResultUIExtraType.DIFF_TEXT, diff_text=diff_text)
+        ui_extra = model.ToolResultUIExtra(type=model.ToolResultUIExtraType.DIFF_TEXT, diff_text=diff_text)
 
         # Update tracker with new mtime
         if file_tracker is not None:
@@ -220,7 +224,7 @@ class EditTool(ToolABC):
         # Build output message
         if args.replace_all:
             msg = f"The file {file_path} has been updated. All occurrences of '{args.old_string}' were successfully replaced with '{args.new_string}'."
-            return ToolResultItem(status="success", output=msg, ui_extra=ui_extra)
+            return model.ToolResultItem(status="success", output=msg, ui_extra=ui_extra)
 
         # For single replacement, show a snippet consisting of context + added lines only
         # Parse the diff to collect target line numbers in the 'after' file
@@ -267,4 +271,4 @@ class EditTool(ToolABC):
             f"The file {file_path} has been updated. Here's the result of running `cat -n` on a snippet of the edited file:\n"
             f"{snippet}"
         )
-        return ToolResultItem(status="success", output=output, ui_extra=ui_extra)
+        return model.ToolResultItem(status="success", output=output, ui_extra=ui_extra)