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

klaude_code/core/tool/todo/todo_write_tool.py

@@ -5,20 +5,10 @@ from pydantic import BaseModel
 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.llm_parameter import ToolSchema
-from klaude_code.protocol.model import (
-    TodoItem,
-    TodoUIExtra,
-    ToolResultItem,
-    ToolResultUIExtra,
-    ToolResultUIExtraType,
-    ToolSideEffect,
-    todo_list_str,
-)
-from klaude_code.protocol.tools import TODO_WRITE
-
-
-def get_new_completed_todos(old_todos: list[TodoItem], new_todos: list[TodoItem]) -> list[str]:
+from klaude_code.protocol import llm_param, model, tools
+
+
+def get_new_completed_todos(old_todos: list[model.TodoItem], new_todos: list[model.TodoItem]) -> list[str]:
     """
     Compare old and new todo lists to find newly completed todos.
 
@@ -49,15 +39,15 @@ def get_new_completed_todos(old_todos: list[TodoItem], new_todos: list[TodoItem]
 
 
 class TodoWriteArguments(BaseModel):
-    todos: list[TodoItem]
+    todos: list[model.TodoItem]
 
 
-@register(TODO_WRITE)
+@register(tools.TODO_WRITE)
 class TodoWriteTool(ToolABC):
     @classmethod
-    def schema(cls) -> ToolSchema:
-        return ToolSchema(
-            name=TODO_WRITE,
+    def schema(cls) -> llm_param.ToolSchema:
+        return llm_param.ToolSchema(
+            name=tools.TODO_WRITE,
             type="function",
             description=load_desc(Path(__file__).parent / "todo_write_tool.md"),
             parameters={
@@ -69,7 +59,10 @@ class TodoWriteTool(ToolABC):
                             "type": "object",
                             "properties": {
                                 "content": {"type": "string", "minLength": 1},
-                                "status": {"type": "string", "enum": ["pending", "in_progress", "completed"]},
+                                "status": {
+                                    "type": "string",
+                                    "enum": ["pending", "in_progress", "completed"],
+                                },
                                 "activeForm": {"type": "string", "minLength": 1},
                             },
                             "required": ["content", "status", "activeForm"],
@@ -84,11 +77,11 @@ class TodoWriteTool(ToolABC):
         )
 
     @classmethod
-    async def call(cls, arguments: str) -> ToolResultItem:
+    async def call(cls, arguments: str) -> model.ToolResultItem:
         try:
             args = TodoWriteArguments.model_validate_json(arguments)
         except ValueError as e:
-            return ToolResultItem(
+            return model.ToolResultItem(
                 status="error",
                 output=f"Invalid arguments: {e}",
             )
@@ -96,7 +89,7 @@ class TodoWriteTool(ToolABC):
         # Get current todo context to store todos
         todo_context = get_current_todo_context()
         if todo_context is None:
-            return ToolResultItem(
+            return model.ToolResultItem(
                 status="error",
                 output="No active session found",
             )
@@ -110,19 +103,19 @@ class TodoWriteTool(ToolABC):
         # Store todos via todo context
         todo_context.set_todos(args.todos)
 
-        ui_extra = TodoUIExtra(todos=args.todos, new_completed=new_completed)
+        ui_extra = model.TodoUIExtra(todos=args.todos, new_completed=new_completed)
 
         response = f"""Todos have been modified successfully. Ensure that you continue to use the todo list to track your progress. Please proceed with the current tasks if applicable
 
 <system-reminder>
 Your todo list has changed. DO NOT mention this explicitly to the user. Here are the latest contents of your todo list:
 
-{todo_list_str(args.todos)}. Continue on with the tasks at hand if applicable.
+{model.todo_list_str(args.todos)}. Continue on with the tasks at hand if applicable.
 </system-reminder>"""
 
-        return ToolResultItem(
+        return model.ToolResultItem(
             status="success",
             output=response,
-            ui_extra=ToolResultUIExtra(type=ToolResultUIExtraType.TODO_LIST, todo_list=ui_extra),
-            side_effects=[ToolSideEffect.TODO_CHANGE],
+            ui_extra=model.ToolResultUIExtra(type=model.ToolResultUIExtraType.TODO_LIST, todo_list=ui_extra),
+            side_effects=[model.ToolSideEffect.TODO_CHANGE],
         )

klaude_code/core/tool/todo/update_plan_tool.py

@@ -9,24 +9,14 @@ 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.llm_parameter import ToolSchema
-from klaude_code.protocol.model import (
-    TodoItem,
-    TodoStatusType,
-    TodoUIExtra,
-    ToolResultItem,
-    ToolResultUIExtra,
-    ToolResultUIExtraType,
-    ToolSideEffect,
-)
-from klaude_code.protocol.tools import UPDATE_PLAN
+from klaude_code.protocol import llm_param, model, tools
 
 from .todo_write_tool import get_new_completed_todos
 
 
 class PlanItemArguments(BaseModel):
     step: str
-    status: TodoStatusType
+    status: model.TodoStatusType
 
     @field_validator("step")
     @classmethod
@@ -51,12 +41,12 @@ class UpdatePlanArguments(BaseModel):
         return value
 
 
-@register(UPDATE_PLAN)
+@register(tools.UPDATE_PLAN)
 class UpdatePlanTool(ToolABC):
     @classmethod
-    def schema(cls) -> ToolSchema:
-        return ToolSchema(
-            name=UPDATE_PLAN,
+    def schema(cls) -> llm_param.ToolSchema:
+        return llm_param.ToolSchema(
+            name=tools.UPDATE_PLAN,
             type="function",
             description=load_desc(Path(__file__).parent / "update_plan_tool.md"),
             parameters={
@@ -89,26 +79,26 @@ class UpdatePlanTool(ToolABC):
         )
 
     @classmethod
-    async def call(cls, arguments: str) -> ToolResultItem:
+    async def call(cls, arguments: str) -> model.ToolResultItem:
         try:
             args = UpdatePlanArguments.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}")
 
         todo_context = get_current_todo_context()
         if todo_context is None:
-            return ToolResultItem(status="error", output="No active session found")
+            return model.ToolResultItem(status="error", output="No active session found")
 
-        new_todos = [TodoItem(content=item.step, status=item.status) for item in args.plan]
+        new_todos = [model.TodoItem(content=item.step, status=item.status) for item in args.plan]
         old_todos = todo_context.get_todos()
         new_completed = get_new_completed_todos(old_todos, new_todos)
         todo_context.set_todos(new_todos)
 
-        ui_extra = TodoUIExtra(todos=new_todos, new_completed=new_completed)
+        ui_extra = model.TodoUIExtra(todos=new_todos, new_completed=new_completed)
 
-        return ToolResultItem(
+        return model.ToolResultItem(
             status="success",
             output="Plan updated",
-            ui_extra=ToolResultUIExtra(type=ToolResultUIExtraType.TODO_LIST, todo_list=ui_extra),
-            side_effects=[ToolSideEffect.TODO_CHANGE],
+            ui_extra=model.ToolResultUIExtra(type=model.ToolResultUIExtraType.TODO_LIST, todo_list=ui_extra),
+            side_effects=[model.ToolSideEffect.TODO_CHANGE],
         )

klaude_code/core/tool/tool_abc.py

@@ -2,8 +2,7 @@ import string
 from abc import ABC, abstractmethod
 from pathlib import Path
 
-from klaude_code.protocol.llm_parameter import ToolSchema
-from klaude_code.protocol.model import ToolResultItem
+from klaude_code.protocol import llm_param, model
 
 
 def load_desc(path: Path, substitutions: dict[str, str] | None = None) -> str:
@@ -17,10 +16,10 @@ def load_desc(path: Path, substitutions: dict[str, str] | None = None) -> str:
 class ToolABC(ABC):
     @classmethod
     @abstractmethod
-    def schema(cls) -> ToolSchema:
+    def schema(cls) -> llm_param.ToolSchema:
         raise NotImplementedError
 
     @classmethod
     @abstractmethod
-    async def call(cls, arguments: str) -> ToolResultItem:
+    async def call(cls, arguments: str) -> model.ToolResultItem:
         raise NotImplementedError

klaude_code/core/tool/tool_context.py

@@ -5,8 +5,8 @@ from contextlib import contextmanager
 from contextvars import ContextVar, Token
 from dataclasses import dataclass
 
-from klaude_code.core.sub_agent import SubAgentResult
-from klaude_code.protocol.model import SubAgentState, TodoItem
+from klaude_code.protocol import model
+from klaude_code.protocol.sub_agent import SubAgentResult
 from klaude_code.session.session import Session
 
 
@@ -18,8 +18,8 @@ class TodoContext:
     a new list; they cannot access the full Session object.
     """
 
-    get_todos: Callable[[], list[TodoItem]]
-    set_todos: Callable[[list[TodoItem]], None]
+    get_todos: Callable[[], list[model.TodoItem]]
+    set_todos: Callable[[list[model.TodoItem]], None]
 
 
 @dataclass
@@ -100,7 +100,7 @@ def get_current_todo_context() -> TodoContext | None:
 
 
 # Holds a handle to run a nested subtask (sub-agent) from within a tool call.
-# The callable takes a SubAgentState and returns a SubAgentResult.
-current_run_subtask_callback: ContextVar[Callable[[SubAgentState], Awaitable[SubAgentResult]] | None] = ContextVar(
-    "current_run_subtask_callback", default=None
+# The callable takes a model.SubAgentState and returns a SubAgentResult.
+current_run_subtask_callback: ContextVar[Callable[[model.SubAgentState], Awaitable[SubAgentResult]] | None] = (
+    ContextVar("current_run_subtask_callback", default=None)
 )

klaude_code/core/tool/tool_registry.py

@@ -1,10 +1,9 @@
 from typing import Callable, TypeVar
 
-from klaude_code.core.sub_agent import get_sub_agent_profile, iter_sub_agent_profiles, sub_agent_tool_names
 from klaude_code.core.tool.sub_agent_tool import SubAgentTool
 from klaude_code.core.tool.tool_abc import ToolABC
-from klaude_code.protocol import tools
-from klaude_code.protocol.llm_parameter import ToolSchema
+from klaude_code.protocol import llm_param, tools
+from klaude_code.protocol.sub_agent import get_sub_agent_profile, iter_sub_agent_profiles, sub_agent_tool_names
 
 _REGISTRY: dict[str, type[ToolABC]] = {}
 
@@ -33,8 +32,8 @@ def list_tools() -> list[str]:
     return list(_REGISTRY.keys())
 
 
-def get_tool_schemas(tool_names: list[str]) -> list[ToolSchema]:
-    schemas: list[ToolSchema] = []
+def get_tool_schemas(tool_names: list[str]) -> list[llm_param.ToolSchema]:
+    schemas: list[llm_param.ToolSchema] = []
     for tool_name in tool_names:
         if tool_name not in _REGISTRY:
             raise ValueError(f"Unknown Tool: {tool_name}")
@@ -47,47 +46,31 @@ def get_registry() -> dict[str, type[ToolABC]]:
     return _REGISTRY
 
 
-def get_vanilla_tools() -> list[ToolSchema]:
-    base_tool_names = [
-        tools.BASH,
-        tools.EDIT,
-        tools.WRITE,
-        tools.READ,
-    ]
-    return get_tool_schemas(base_tool_names)
-
-
-def get_main_agent_tools(model_name: str) -> list[ToolSchema]:
-    def _base_main_tools(name: str) -> list[str]:
-        if "gpt-5" in name:
-            return [
-                tools.BASH,
-                tools.READ,
-                tools.APPLY_PATCH,
-                tools.UPDATE_PLAN,
-            ]
-        return [
-            tools.BASH,
-            tools.READ,
-            tools.EDIT,
-            tools.WRITE,
-            tools.TODO_WRITE,
-        ]
-
-    tool_names = _base_main_tools(model_name)
-    tool_names.extend(sub_agent_tool_names(enabled_only=True, model_name=model_name))
-    tool_names.extend(
-        [
-            tools.SKILL,
-            tools.MERMAID,
-            tools.MEMORY,
-        ]
-    )
-    return get_tool_schemas(tool_names)
+def load_agent_tools(
+    model_name: str, sub_agent_type: tools.SubAgentType | None = None, *, vanilla: bool = False
+) -> list[llm_param.ToolSchema]:
+    """Get tools for an agent based on model and agent type.
+
+    Args:
+        model_name: The model name.
+        sub_agent_type: If None, returns main agent tools. Otherwise returns sub-agent tools.
+        vanilla: If True, returns minimal vanilla tools (ignores sub_agent_type).
+    """
+    if vanilla:
+        return get_tool_schemas([tools.BASH, tools.EDIT, tools.WRITE, tools.READ])
 
+    if sub_agent_type is not None:
+        profile = get_sub_agent_profile(sub_agent_type)
+        if not profile.enabled_for_model(model_name):
+            return []
+        return get_tool_schemas(list(profile.tool_set))
 
-def get_sub_agent_tools(model_name: str, sub_agent_type: tools.SubAgentType) -> list[ToolSchema]:
-    profile = get_sub_agent_profile(sub_agent_type)
-    if not profile.enabled_for_model(model_name):
-        return []
-    return get_tool_schemas(list(profile.tool_set))
+    # Main agent tools
+    if "gpt-5" in model_name:
+        tool_names = [tools.BASH, tools.READ, tools.APPLY_PATCH, tools.UPDATE_PLAN]
+    else:
+        tool_names = [tools.BASH, tools.READ, tools.EDIT, tools.WRITE, tools.TODO_WRITE]
+
+    tool_names.extend(sub_agent_tool_names(enabled_only=True, model_name=model_name))
+    tool_names.extend([tools.SKILL, tools.MERMAID, tools.MEMORY])
+    return get_tool_schemas(tool_names)

klaude_code/core/tool/tool_runner.py

@@ -2,22 +2,14 @@ import asyncio
 from collections.abc import AsyncGenerator, Callable, Iterable, Sequence
 from dataclasses import dataclass
 
-from klaude_code.const import CANCEL_OUTPUT
-from klaude_code.core.sub_agent import is_sub_agent_tool
+from klaude_code import const
 from klaude_code.core.tool.tool_abc import ToolABC
 from klaude_code.core.tool.truncation import truncate_tool_output
 from klaude_code.protocol import model
-from klaude_code.protocol.model import (
-    ToolCallItem,
-    ToolResultItem,
-    ToolResultUIExtra,
-    ToolResultUIExtraType,
-    ToolSideEffect,
-    TruncationUIExtra,
-)
+from klaude_code.protocol.sub_agent import is_sub_agent_tool
 
 
-async def run_tool(tool_call: ToolCallItem, registry: dict[str, type[ToolABC]]) -> ToolResultItem:
+async def run_tool(tool_call: model.ToolCallItem, registry: dict[str, type[ToolABC]]) -> model.ToolResultItem:
     """Execute a tool call and return the result.
 
     Args:
@@ -28,7 +20,7 @@ async def run_tool(tool_call: ToolCallItem, registry: dict[str, type[ToolABC]])
         The result of the tool execution.
     """
     if tool_call.name not in registry:
-        return ToolResultItem(
+        return model.ToolResultItem(
             call_id=tool_call.call_id,
             output=f"Tool {tool_call.name} not exists",
             status="error",
@@ -42,9 +34,9 @@ async def run_tool(tool_call: ToolCallItem, registry: dict[str, type[ToolABC]])
             truncation_result = truncate_tool_output(tool_result.output, tool_call)
             tool_result.output = truncation_result.output
             if truncation_result.was_truncated and truncation_result.saved_file_path:
-                tool_result.ui_extra = ToolResultUIExtra(
-                    type=ToolResultUIExtraType.TRUNCATION,
-                    truncation=TruncationUIExtra(
+                tool_result.ui_extra = model.ToolResultUIExtra(
+                    type=model.ToolResultUIExtraType.TRUNCATION,
+                    truncation=model.TruncationUIExtra(
                         saved_file_path=truncation_result.saved_file_path,
                         original_length=truncation_result.original_length,
                         truncated_length=truncation_result.truncated_length,
@@ -55,7 +47,7 @@ async def run_tool(tool_call: ToolCallItem, registry: dict[str, type[ToolABC]])
         # Propagate cooperative cancellation so outer layers can handle interrupts correctly.
         raise
     except Exception as e:
-        return ToolResultItem(
+        return model.ToolResultItem(
             call_id=tool_call.call_id,
             output=f"Tool {tool_call.name} execution error: {e.__class__.__name__} {e}",
             status="error",
@@ -67,15 +59,15 @@ async def run_tool(tool_call: ToolCallItem, registry: dict[str, type[ToolABC]])
 class ToolExecutionCallStarted:
     """Represents the start of a tool call execution."""
 
-    tool_call: ToolCallItem
+    tool_call: model.ToolCallItem
 
 
 @dataclass
 class ToolExecutionResult:
     """Represents the completion of a tool call with its result."""
 
-    tool_call: ToolCallItem
-    tool_result: ToolResultItem
+    tool_call: model.ToolCallItem
+    tool_result: model.ToolResultItem
 
 
 @dataclass
@@ -107,11 +99,11 @@ class ToolExecutor:
         self._registry = registry
         self._append_history = append_history
 
-        self._unfinished_calls: dict[str, ToolCallItem] = {}
+        self._unfinished_calls: dict[str, model.ToolCallItem] = {}
         self._call_event_emitted: set[str] = set()
         self._sub_agent_tasks: set[asyncio.Task[list[ToolExecutorEvent]]] = set()
 
-    async def run_tools(self, tool_calls: list[ToolCallItem]) -> AsyncGenerator[ToolExecutorEvent, None]:
+    async def run_tools(self, tool_calls: list[model.ToolCallItem]) -> AsyncGenerator[ToolExecutorEvent, None]:
         """Run the given tool calls and yield execution events.
 
         Tool calls are partitioned into regular tools and sub-agent tools. Regular tools
@@ -152,11 +144,15 @@ class ToolExecutor:
                 execution_tasks.append(task)
 
             for task in asyncio.as_completed(execution_tasks):
-                try:
-                    result_events = await task
-                except asyncio.CancelledError:
-                    # Task was cancelled by ToolExecutor.cancel(); skip synthetic events here.
-                    continue
+                # Do not swallow asyncio.CancelledError here:
+                # - If the user interrupts the main agent, the executor cancels the
+                #   outer agent task, which should propagate cancellation up through
+                #   tool execution so the task can terminate and emit TaskFinishEvent.
+                # - Sub-agent tool tasks cancelled via ToolExecutor.cancel() are
+                #   handled by synthesizing ToolExecutionResult events; any
+                #   CancelledError raised here should still bubble up so the
+                #   calling agent can stop cleanly, matching pre-refactor behavior.
+                result_events = await task
 
                 for exec_event in result_events:
                     yield exec_event
@@ -183,23 +179,19 @@ class ToolExecutor:
             return events_to_yield
 
         for call_id, tool_call in list(self._unfinished_calls.items()):
-            cancel_result = ToolResultItem(
+            cancel_result = model.ToolResultItem(
                 call_id=tool_call.call_id,
-                output=CANCEL_OUTPUT,
+                output=const.CANCEL_OUTPUT,
                 status="error",
                 tool_name=tool_call.name,
                 ui_extra=None,
             )
 
             if call_id not in self._call_event_emitted:
-                events_to_yield.append(
-                    ToolExecutionCallStarted(tool_call=tool_call)
-                )
+                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))
 
             self._append_history([cancel_result])
             self._unfinished_calls.pop(call_id, None)
@@ -216,10 +208,10 @@ class ToolExecutor:
 
     @staticmethod
     def _partition_tool_calls(
-        tool_calls: list[ToolCallItem],
-    ) -> tuple[list[ToolCallItem], list[ToolCallItem]]:
-        regular_tool_calls: list[ToolCallItem] = []
-        sub_agent_tool_calls: list[ToolCallItem] = []
+        tool_calls: list[model.ToolCallItem],
+    ) -> tuple[list[model.ToolCallItem], list[model.ToolCallItem]]:
+        regular_tool_calls: list[model.ToolCallItem] = []
+        sub_agent_tool_calls: list[model.ToolCallItem] = []
         for tool_call in tool_calls:
             if is_sub_agent_tool(tool_call.name):
                 sub_agent_tool_calls.append(tool_call)
@@ -227,11 +219,11 @@ class ToolExecutor:
                 regular_tool_calls.append(tool_call)
         return regular_tool_calls, sub_agent_tool_calls
 
-    def _build_tool_call_started(self, tool_call: ToolCallItem) -> ToolExecutionCallStarted:
+    def _build_tool_call_started(self, tool_call: model.ToolCallItem) -> ToolExecutionCallStarted:
         return ToolExecutionCallStarted(tool_call=tool_call)
 
-    async def _run_single_tool_call(self, tool_call: ToolCallItem) -> list[ToolExecutorEvent]:
-        tool_result: ToolResultItem = await run_tool(tool_call, self._registry)
+    async def _run_single_tool_call(self, tool_call: model.ToolCallItem) -> list[ToolExecutorEvent]:
+        tool_result: model.ToolResultItem = await run_tool(tool_call, self._registry)
 
         self._append_history([tool_result])
 
@@ -242,7 +234,7 @@ class ToolExecutor:
         extra_events = self._build_tool_side_effect_events(tool_result)
         return [result_event, *extra_events]
 
-    def _build_tool_side_effect_events(self, tool_result: ToolResultItem) -> list[ToolExecutorEvent]:
+    def _build_tool_side_effect_events(self, tool_result: model.ToolResultItem) -> list[ToolExecutorEvent]:
         side_effects = tool_result.side_effects
         if not side_effects:
             return []
@@ -250,7 +242,7 @@ class ToolExecutor:
         side_effect_events: list[ToolExecutorEvent] = []
 
         for side_effect in side_effects:
-            if side_effect == ToolSideEffect.TODO_CHANGE:
+            if side_effect == model.ToolSideEffect.TODO_CHANGE:
                 todos: list[model.TodoItem] | None = None
                 if tool_result.ui_extra is not None and tool_result.ui_extra.todo_list is not None:
                     todos = tool_result.ui_extra.todo_list.todos

klaude_code/core/tool/truncation.py

@@ -6,14 +6,8 @@ from dataclasses import dataclass
 from pathlib import Path
 from urllib.parse import urlparse
 
-from klaude_code.const import (
-    TOOL_OUTPUT_DISPLAY_HEAD,
-    TOOL_OUTPUT_DISPLAY_TAIL,
-    TOOL_OUTPUT_MAX_LENGTH,
-    TOOL_OUTPUT_TRUNCATION_DIR,
-)
-from klaude_code.protocol.model import ToolCallItem
-from klaude_code.protocol.tools import WEB_FETCH
+from klaude_code import const
+from klaude_code.protocol import model, tools
 
 
 @dataclass
@@ -47,7 +41,7 @@ class TruncationStrategy(ABC):
     """Abstract base class for tool output truncation strategies."""
 
     @abstractmethod
-    def truncate(self, output: str, tool_call: ToolCallItem | None = None) -> TruncationResult:
+    def truncate(self, output: str, tool_call: model.ToolCallItem | None = None) -> TruncationResult:
         """Truncate the output according to the strategy."""
         ...
 
@@ -55,10 +49,10 @@ class TruncationStrategy(ABC):
 class SimpleTruncationStrategy(TruncationStrategy):
     """Simple character-based truncation strategy."""
 
-    def __init__(self, max_length: int = TOOL_OUTPUT_MAX_LENGTH):
+    def __init__(self, max_length: int = const.TOOL_OUTPUT_MAX_LENGTH):
         self.max_length = max_length
 
-    def truncate(self, output: str, tool_call: ToolCallItem | None = None) -> TruncationResult:
+    def truncate(self, output: str, tool_call: model.ToolCallItem | 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)"
@@ -76,19 +70,19 @@ class SmartTruncationStrategy(TruncationStrategy):
 
     def __init__(
         self,
-        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,
+        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,
     ):
         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: ToolCallItem | None) -> str:
+    def _get_file_identifier(self, tool_call: model.ToolCallItem | 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 == WEB_FETCH:
+        if tool_call and tool_call.name == tools.WEB_FETCH:
             try:
                 args = json.loads(tool_call.arguments)
                 url = args.get("url", "")
@@ -101,7 +95,7 @@ class SmartTruncationStrategy(TruncationStrategy):
             return tool_call.call_id.replace("/", "_")
         return "unknown"
 
-    def _save_to_file(self, output: str, tool_call: ToolCallItem | None) -> str | None:
+    def _save_to_file(self, output: str, tool_call: model.ToolCallItem | 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)
@@ -115,7 +109,7 @@ class SmartTruncationStrategy(TruncationStrategy):
         except (OSError, IOError):
             return None
 
-    def truncate(self, output: str, tool_call: ToolCallItem | None = None) -> TruncationResult:
+    def truncate(self, output: str, tool_call: model.ToolCallItem | None = None) -> TruncationResult:
         original_length = len(output)
 
         if original_length <= self.max_length:
@@ -171,6 +165,6 @@ def set_truncation_strategy(strategy: TruncationStrategy) -> None:
     _default_strategy = strategy
 
 
-def truncate_tool_output(output: str, tool_call: ToolCallItem | None = None) -> TruncationResult:
+def truncate_tool_output(output: str, tool_call: model.ToolCallItem | None = None) -> TruncationResult:
     """Truncate tool output using the current strategy."""
     return get_truncation_strategy().truncate(output, tool_call)