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

klaude_code/core/tool/memory/memory_tool.py

@@ -13,9 +13,7 @@ from pydantic import BaseModel, Field
 
 from klaude_code.core.tool.tool_abc import ToolABC, load_desc
 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 MEMORY
+from klaude_code.protocol import llm_param, model, tools
 
 MEMORY_VIRTUAL_ROOT = "/memories"
 MEMORY_DIR_NAME = ".claude/memories"
@@ -102,7 +100,7 @@ def _format_numbered_line(line_no: int, content: str) -> str:
     return f"{line_no:>6}|{content}"
 
 
-def _make_diff_ui_extra(before: str, after: str, path: str) -> ToolResultUIExtra:
+def _make_diff_ui_extra(before: str, after: str, path: str) -> model.ToolResultUIExtra:
     diff_lines = list(
         difflib.unified_diff(
             before.splitlines(),
@@ -113,10 +111,10 @@ def _make_diff_ui_extra(before: str, after: str, path: str) -> ToolResultUIExtra
         )
     )
     diff_text = "\n".join(diff_lines)
-    return ToolResultUIExtra(type=ToolResultUIExtraType.DIFF_TEXT, diff_text=diff_text)
+    return model.ToolResultUIExtra(type=model.ToolResultUIExtraType.DIFF_TEXT, diff_text=diff_text)
 
 
-@register(MEMORY)
+@register(tools.MEMORY)
 class MemoryTool(ToolABC):
     class MemoryArguments(BaseModel):
         command: Literal["view", "create", "str_replace", "insert", "delete", "rename"]
@@ -136,9 +134,9 @@ class MemoryTool(ToolABC):
         new_path: str | None = Field(default=None)
 
     @classmethod
-    def schema(cls) -> ToolSchema:
-        return ToolSchema(
-            name=MEMORY,
+    def schema(cls) -> llm_param.ToolSchema:
+        return llm_param.ToolSchema(
+            name=tools.MEMORY,
             type="function",
             description=load_desc(Path(__file__).parent / "memory_tool.md"),
             parameters={
@@ -146,7 +144,14 @@ class MemoryTool(ToolABC):
                 "properties": {
                     "command": {
                         "type": "string",
-                        "enum": ["view", "create", "str_replace", "insert", "delete", "rename"],
+                        "enum": [
+                            "view",
+                            "create",
+                            "str_replace",
+                            "insert",
+                            "delete",
+                            "rename",
+                        ],
                         "description": "The memory operation to perform",
                     },
                     "path": {
@@ -193,11 +198,11 @@ class MemoryTool(ToolABC):
         )
 
     @classmethod
-    async def call(cls, arguments: str) -> ToolResultItem:
+    async def call(cls, arguments: str) -> model.ToolResultItem:
         try:
             args = cls.MemoryArguments.model_validate_json(arguments)
         except Exception as e:
-            return ToolResultItem(status="error", output=f"Invalid arguments: {e}")
+            return model.ToolResultItem(status="error", output=f"Invalid arguments: {e}")
 
         command = args.command
         if command == "view":
@@ -213,37 +218,40 @@ class MemoryTool(ToolABC):
         elif command == "rename":
             return await cls._rename(args)
         else:
-            return ToolResultItem(status="error", output=f"Unknown command: {command}")
+            return model.ToolResultItem(status="error", output=f"Unknown command: {command}")
 
     @classmethod
-    async def _view(cls, args: MemoryArguments) -> ToolResultItem:
+    async def _view(cls, args: MemoryArguments) -> model.ToolResultItem:
         if args.path is None:
-            return ToolResultItem(status="error", output="path is required for view command")
+            return model.ToolResultItem(status="error", output="path is required for view command")
 
         actual_path, error = _validate_path(args.path)
         if error:
-            return ToolResultItem(status="error", output=error)
+            return model.ToolResultItem(status="error", output=error)
         assert actual_path is not None
 
         # Ensure memories directory exists
         _ensure_memories_dir()
 
         if not actual_path.exists():
-            return ToolResultItem(status="error", output=f"Path does not exist: {args.path}")
+            return model.ToolResultItem(status="error", output=f"Path does not exist: {args.path}")
 
         if actual_path.is_dir():
             # List directory contents
             try:
-                entries = sorted(actual_path.iterdir(), key=lambda p: (not p.is_dir(), p.name.lower()))
+                entries = sorted(
+                    actual_path.iterdir(),
+                    key=lambda p: (not p.is_dir(), p.name.lower()),
+                )
                 lines = [f"Directory: {args.path}"]
                 for entry in entries:
                     prefix = "/" if entry.is_dir() else ""
                     lines.append(f"- {entry.name}{prefix}")
                 if len(entries) == 0:
                     lines.append("(empty directory)")
-                return ToolResultItem(status="success", output="\n".join(lines))
+                return model.ToolResultItem(status="success", output="\n".join(lines))
             except Exception as e:
-                return ToolResultItem(status="error", output=f"Failed to list directory: {e}")
+                return model.ToolResultItem(status="error", output=f"Failed to list directory: {e}")
         else:
             # Read file contents
             try:
@@ -259,7 +267,7 @@ class MemoryTool(ToolABC):
                     end = min(total_lines, args.view_range[1])
 
                 if start > total_lines:
-                    return ToolResultItem(
+                    return model.ToolResultItem(
                         status="success",
                         output=f"File has {total_lines} lines, requested start line {start} is beyond end of file",
                     )
@@ -269,25 +277,28 @@ class MemoryTool(ToolABC):
                 output = "\n".join(numbered)
                 if not output:
                     output = "(empty file)"
-                return ToolResultItem(status="success", output=output)
+                return model.ToolResultItem(status="success", output=output)
             except Exception as e:
-                return ToolResultItem(status="error", output=f"Failed to read file: {e}")
+                return model.ToolResultItem(status="error", output=f"Failed to read file: {e}")
 
     @classmethod
-    async def _create(cls, args: MemoryArguments) -> ToolResultItem:
+    async def _create(cls, args: MemoryArguments) -> model.ToolResultItem:
         if args.path is None:
-            return ToolResultItem(status="error", output="path is required for create command")
+            return model.ToolResultItem(status="error", output="path is required for create command")
         if args.file_text is None:
-            return ToolResultItem(status="error", output="file_text is required for create command")
+            return model.ToolResultItem(status="error", output="file_text is required for create command")
 
         actual_path, error = _validate_path(args.path)
         if error:
-            return ToolResultItem(status="error", output=error)
+            return model.ToolResultItem(status="error", output=error)
         assert actual_path is not None
 
         # Cannot create the root directory itself
         if args.path == MEMORY_VIRTUAL_ROOT or args.path == MEMORY_VIRTUAL_ROOT + "/":
-            return ToolResultItem(status="error", output="Cannot create the memories root directory as a file")
+            return model.ToolResultItem(
+                status="error",
+                output="Cannot create the memories root directory as a file",
+            )
 
         try:
             # Read existing content for diff (if file exists)
@@ -300,60 +311,64 @@ class MemoryTool(ToolABC):
             await asyncio.to_thread(actual_path.write_text, args.file_text, encoding="utf-8")
 
             ui_extra = _make_diff_ui_extra(before, args.file_text, args.path)
-            return ToolResultItem(status="success", output=f"File created: {args.path}", ui_extra=ui_extra)
+            return model.ToolResultItem(status="success", output=f"File created: {args.path}", ui_extra=ui_extra)
         except Exception as e:
-            return ToolResultItem(status="error", output=f"Failed to create file: {e}")
+            return model.ToolResultItem(status="error", output=f"Failed to create file: {e}")
 
     @classmethod
-    async def _str_replace(cls, args: MemoryArguments) -> ToolResultItem:
+    async def _str_replace(cls, args: MemoryArguments) -> model.ToolResultItem:
         if args.path is None:
-            return ToolResultItem(status="error", output="path is required for str_replace command")
+            return model.ToolResultItem(status="error", output="path is required for str_replace command")
         if args.old_str is None:
-            return ToolResultItem(status="error", output="old_str is required for str_replace command")
+            return model.ToolResultItem(status="error", output="old_str is required for str_replace command")
         if args.new_str is None:
-            return ToolResultItem(status="error", output="new_str is required for str_replace command")
+            return model.ToolResultItem(status="error", output="new_str is required for str_replace command")
 
         actual_path, error = _validate_path(args.path)
         if error:
-            return ToolResultItem(status="error", output=error)
+            return model.ToolResultItem(status="error", output=error)
         assert actual_path is not None
 
         if not actual_path.exists():
-            return ToolResultItem(status="error", output=f"File does not exist: {args.path}")
+            return model.ToolResultItem(status="error", output=f"File does not exist: {args.path}")
         if actual_path.is_dir():
-            return ToolResultItem(status="error", output="Cannot perform str_replace on a directory")
+            return model.ToolResultItem(status="error", output="Cannot perform str_replace on a directory")
 
         try:
             before = await asyncio.to_thread(actual_path.read_text, encoding="utf-8")
             if args.old_str not in before:
-                return ToolResultItem(status="error", output=f"String not found in file: {args.old_str}")
+                return model.ToolResultItem(status="error", output=f"String not found in file: {args.old_str}")
 
             after = before.replace(args.old_str, args.new_str, 1)
             await asyncio.to_thread(actual_path.write_text, after, encoding="utf-8")
 
             ui_extra = _make_diff_ui_extra(before, after, args.path)
-            return ToolResultItem(status="success", output=f"Replaced text in {args.path}", ui_extra=ui_extra)
+            return model.ToolResultItem(
+                status="success",
+                output=f"Replaced text in {args.path}",
+                ui_extra=ui_extra,
+            )
         except Exception as e:
-            return ToolResultItem(status="error", output=f"Failed to replace text: {e}")
+            return model.ToolResultItem(status="error", output=f"Failed to replace text: {e}")
 
     @classmethod
-    async def _insert(cls, args: MemoryArguments) -> ToolResultItem:
+    async def _insert(cls, args: MemoryArguments) -> model.ToolResultItem:
         if args.path is None:
-            return ToolResultItem(status="error", output="path is required for insert command")
+            return model.ToolResultItem(status="error", output="path is required for insert command")
         if args.insert_line is None:
-            return ToolResultItem(status="error", output="insert_line is required for insert command")
+            return model.ToolResultItem(status="error", output="insert_line is required for insert command")
         if args.insert_text is None:
-            return ToolResultItem(status="error", output="insert_text is required for insert command")
+            return model.ToolResultItem(status="error", output="insert_text is required for insert command")
 
         actual_path, error = _validate_path(args.path)
         if error:
-            return ToolResultItem(status="error", output=error)
+            return model.ToolResultItem(status="error", output=error)
         assert actual_path is not None
 
         if not actual_path.exists():
-            return ToolResultItem(status="error", output=f"File does not exist: {args.path}")
+            return model.ToolResultItem(status="error", output=f"File does not exist: {args.path}")
         if actual_path.is_dir():
-            return ToolResultItem(status="error", output="Cannot insert into a directory")
+            return model.ToolResultItem(status="error", output="Cannot insert into a directory")
 
         try:
             before = await asyncio.to_thread(actual_path.read_text, encoding="utf-8")
@@ -377,69 +392,71 @@ class MemoryTool(ToolABC):
             await asyncio.to_thread(actual_path.write_text, after, encoding="utf-8")
 
             ui_extra = _make_diff_ui_extra(before, after, args.path)
-            return ToolResultItem(
-                status="success", output=f"Inserted text at line {args.insert_line} in {args.path}", ui_extra=ui_extra
+            return model.ToolResultItem(
+                status="success",
+                output=f"Inserted text at line {args.insert_line} in {args.path}",
+                ui_extra=ui_extra,
             )
         except Exception as e:
-            return ToolResultItem(status="error", output=f"Failed to insert text: {e}")
+            return model.ToolResultItem(status="error", output=f"Failed to insert text: {e}")
 
     @classmethod
-    async def _delete(cls, args: MemoryArguments) -> ToolResultItem:
+    async def _delete(cls, args: MemoryArguments) -> model.ToolResultItem:
         if args.path is None:
-            return ToolResultItem(status="error", output="path is required for delete command")
+            return model.ToolResultItem(status="error", output="path is required for delete command")
 
         # Prevent deleting the root memories directory
         if args.path == MEMORY_VIRTUAL_ROOT or args.path == MEMORY_VIRTUAL_ROOT + "/":
-            return ToolResultItem(status="error", output="Cannot delete the memories root directory")
+            return model.ToolResultItem(status="error", output="Cannot delete the memories root directory")
 
         actual_path, error = _validate_path(args.path)
         if error:
-            return ToolResultItem(status="error", output=error)
+            return model.ToolResultItem(status="error", output=error)
         assert actual_path is not None
 
         if not actual_path.exists():
-            return ToolResultItem(status="error", output=f"Path does not exist: {args.path}")
+            return model.ToolResultItem(status="error", output=f"Path does not exist: {args.path}")
 
         try:
             if actual_path.is_dir():
                 await asyncio.to_thread(shutil.rmtree, actual_path)
-                return ToolResultItem(status="success", output=f"Directory deleted: {args.path}")
+                return model.ToolResultItem(status="success", output=f"Directory deleted: {args.path}")
             else:
                 await asyncio.to_thread(os.remove, actual_path)
-                return ToolResultItem(status="success", output=f"File deleted: {args.path}")
+                return model.ToolResultItem(status="success", output=f"File deleted: {args.path}")
         except Exception as e:
-            return ToolResultItem(status="error", output=f"Failed to delete: {e}")
+            return model.ToolResultItem(status="error", output=f"Failed to delete: {e}")
 
     @classmethod
-    async def _rename(cls, args: MemoryArguments) -> ToolResultItem:
+    async def _rename(cls, args: MemoryArguments) -> model.ToolResultItem:
         if args.old_path is None:
-            return ToolResultItem(status="error", output="old_path is required for rename command")
+            return model.ToolResultItem(status="error", output="old_path is required for rename command")
         if args.new_path is None:
-            return ToolResultItem(status="error", output="new_path is required for rename command")
+            return model.ToolResultItem(status="error", output="new_path is required for rename command")
 
         # Prevent renaming the root memories directory
         if args.old_path == MEMORY_VIRTUAL_ROOT or args.old_path == MEMORY_VIRTUAL_ROOT + "/":
-            return ToolResultItem(status="error", output="Cannot rename the memories root directory")
+            return model.ToolResultItem(status="error", output="Cannot rename the memories root directory")
 
         old_actual, error = _validate_path(args.old_path)
         if error:
-            return ToolResultItem(status="error", output=f"Invalid old_path: {error}")
+            return model.ToolResultItem(status="error", output=f"Invalid old_path: {error}")
         assert old_actual is not None
 
         new_actual, error = _validate_path(args.new_path)
         if error:
-            return ToolResultItem(status="error", output=f"Invalid new_path: {error}")
+            return model.ToolResultItem(status="error", output=f"Invalid new_path: {error}")
         assert new_actual is not None
 
         if not old_actual.exists():
-            return ToolResultItem(status="error", output=f"Source path does not exist: {args.old_path}")
+            return model.ToolResultItem(status="error", output=f"Source path does not exist: {args.old_path}")
         if new_actual.exists():
-            return ToolResultItem(status="error", output=f"Destination already exists: {args.new_path}")
+            return model.ToolResultItem(status="error", output=f"Destination already exists: {args.new_path}")
 
         try:
             # Ensure parent directory of destination exists
             new_actual.parent.mkdir(parents=True, exist_ok=True)
             await asyncio.to_thread(shutil.move, str(old_actual), str(new_actual))
-            return ToolResultItem(status="success", output=f"Renamed {args.old_path} to {args.new_path}")
+            return model.ToolResultItem(status="success", output=f"Renamed {args.old_path} to {args.new_path}")
         except Exception as e:
-            return ToolResultItem(status="error", output=f"Failed to rename: {e}")
+            return model.ToolResultItem(status="error", output=f"Failed to rename: {e}")

klaude_code/core/tool/memory/skill_tool.py

@@ -5,12 +5,10 @@ from pydantic import BaseModel
 from klaude_code.core.tool.memory.skill_loader import SkillLoader
 from klaude_code.core.tool.tool_abc import ToolABC, load_desc
 from klaude_code.core.tool.tool_registry import register
-from klaude_code.protocol.llm_parameter import ToolSchema
-from klaude_code.protocol.model import ToolResultItem
-from klaude_code.protocol.tools import SKILL
+from klaude_code.protocol import llm_param, model, tools
 
 
-@register(SKILL)
+@register(tools.SKILL)
 class SkillTool(ToolABC):
     """Tool to execute/load a skill within the main conversation"""
 
@@ -22,12 +20,12 @@ class SkillTool(ToolABC):
         cls._skill_loader = loader
 
     @classmethod
-    def schema(cls) -> ToolSchema:
+    def schema(cls) -> llm_param.ToolSchema:
         """Generate schema with embedded available skills metadata"""
         skills_xml = cls._generate_skills_xml()
 
-        return ToolSchema(
-            name=SKILL,
+        return llm_param.ToolSchema(
+            name=tools.SKILL,
             type="function",
             description=load_desc(Path(__file__).parent / "skill_tool.md", {"skills_xml": skills_xml}),
             parameters={
@@ -61,18 +59,18 @@ class SkillTool(ToolABC):
         command: str
 
     @classmethod
-    async def call(cls, arguments: str) -> ToolResultItem:
+    async def call(cls, arguments: str) -> model.ToolResultItem:
         """Load and return full skill content"""
         try:
             args = cls.SkillArguments.model_validate_json(arguments)
         except ValueError as e:
-            return ToolResultItem(
+            return model.ToolResultItem(
                 status="error",
                 output=f"Invalid arguments: {e}",
             )
 
         if not cls._skill_loader:
-            return ToolResultItem(
+            return model.ToolResultItem(
                 status="error",
                 output="Skill loader not initialized",
             )
@@ -81,7 +79,7 @@ class SkillTool(ToolABC):
 
         if not skill:
             available = ", ".join(cls._skill_loader.list_skills())
-            return ToolResultItem(
+            return model.ToolResultItem(
                 status="error",
                 output=f"Skill '{args.command}' does not exist. Available skills: {available}",
             )
@@ -96,4 +94,4 @@ class SkillTool(ToolABC):
 Base directory for this skill: {base_dir}
 
 {skill.to_prompt()}"""
-        return ToolResultItem(status="success", output=result)
+        return model.ToolResultItem(status="success", output=result)

klaude_code/core/tool/shell/bash_tool.py

@@ -1,24 +1,26 @@
 import asyncio
+import re
 import subprocess
 from pathlib import Path
 
 from pydantic import BaseModel
 
-from klaude_code.const import BASH_DEFAULT_TIMEOUT_MS
+from klaude_code import const
 from klaude_code.core.tool.shell.command_safety import is_safe_command, strip_bash_lc
 from klaude_code.core.tool.tool_abc import ToolABC, load_desc
 from klaude_code.core.tool.tool_registry import register
-from klaude_code.protocol.llm_parameter import ToolSchema
-from klaude_code.protocol.model import ToolResultItem
-from klaude_code.protocol.tools import BASH
+from klaude_code.protocol import llm_param, model, tools
 
+# Regex to strip ANSI escape sequences from command output
+_ANSI_ESCAPE_RE = re.compile(r"\x1b\[[0-9;]*m")
 
-@register(BASH)
+
+@register(tools.BASH)
 class BashTool(ToolABC):
     @classmethod
-    def schema(cls) -> ToolSchema:
-        return ToolSchema(
-            name=BASH,
+    def schema(cls) -> llm_param.ToolSchema:
+        return llm_param.ToolSchema(
+            name=tools.BASH,
             type="function",
             description=load_desc(Path(__file__).parent / "bash_tool.md"),
             parameters={
@@ -30,8 +32,8 @@ class BashTool(ToolABC):
                     },
                     "timeout_ms": {
                         "type": "integer",
-                        "description": f"The timeout for the command in milliseconds, default is {BASH_DEFAULT_TIMEOUT_MS}",
-                        "default": BASH_DEFAULT_TIMEOUT_MS,
+                        "description": f"The timeout for the command in milliseconds, default is {const.BASH_DEFAULT_TIMEOUT_MS}",
+                        "default": const.BASH_DEFAULT_TIMEOUT_MS,
                     },
                 },
                 "required": ["command"],
@@ -40,27 +42,27 @@ class BashTool(ToolABC):
 
     class BashArguments(BaseModel):
         command: str
-        timeout_ms: int = BASH_DEFAULT_TIMEOUT_MS
+        timeout_ms: int = const.BASH_DEFAULT_TIMEOUT_MS
 
     @classmethod
-    async def call(cls, arguments: str) -> ToolResultItem:
+    async def call(cls, arguments: str) -> model.ToolResultItem:
         try:
             args = BashTool.BashArguments.model_validate_json(arguments)
         except ValueError as e:
-            return ToolResultItem(
+            return model.ToolResultItem(
                 status="error",
                 output=f"Invalid arguments: {e}",
             )
         return await cls.call_with_args(args)
 
     @classmethod
-    async def call_with_args(cls, args: BashArguments) -> ToolResultItem:
+    async def call_with_args(cls, args: BashArguments) -> model.ToolResultItem:
         command_str = strip_bash_lc(args.command)
 
         # Safety check: only execute commands proven as "known safe"
         result = is_safe_command(command_str)
         if not result.is_safe:
-            return ToolResultItem(
+            return model.ToolResultItem(
                 status="error",
                 output=f"Command rejected: {result.error_msg}",
             )
@@ -80,8 +82,8 @@ class BashTool(ToolABC):
                 check=False,
             )
 
-            stdout = completed.stdout or ""
-            stderr = completed.stderr or ""
+            stdout = _ANSI_ESCAPE_RE.sub("", completed.stdout or "")
+            stderr = _ANSI_ESCAPE_RE.sub("", completed.stderr or "")
             rc = completed.returncode
 
             if rc == 0:
@@ -89,7 +91,7 @@ class BashTool(ToolABC):
                 # Include stderr if there is useful diagnostics despite success
                 if stderr.strip():
                     output = (output + ("\n" if output else "")) + f"[stderr]\n{stderr}"
-                return ToolResultItem(
+                return model.ToolResultItem(
                     status="success",
                     output=output.strip(),
                 )
@@ -101,23 +103,23 @@ class BashTool(ToolABC):
                     combined += f"[stderr]\n{stderr}"
                 if not combined:
                     combined = f"Command exited with code {rc}"
-                return ToolResultItem(
+                return model.ToolResultItem(
                     status="error",
                     output=combined.strip(),
                 )
 
         except subprocess.TimeoutExpired:
-            return ToolResultItem(
+            return model.ToolResultItem(
                 status="error",
                 output=f"Timeout after {args.timeout_ms} ms running: {command_str}",
             )
         except FileNotFoundError:
-            return ToolResultItem(
+            return model.ToolResultItem(
                 status="error",
                 output="bash not found on system path",
             )
         except Exception as e:  # safeguard against unexpected failures
-            return ToolResultItem(
+            return model.ToolResultItem(
                 status="error",
                 output=f"Execution error: {e}",
             )

klaude_code/core/tool/shell/command_safety.py

@@ -323,7 +323,18 @@ def _is_safe_argv(argv: list[str]) -> SafetyCheckResult:
         return SafetyCheckResult(True)
 
     # Build tools and linters - allow all subcommands
-    if cmd0 in {"cargo", "uv", "go", "ruff", "pyright", "make", "isort", "npm", "pnpm", "bun"}:
+    if cmd0 in {
+        "cargo",
+        "uv",
+        "go",
+        "ruff",
+        "pyright",
+        "make",
+        "isort",
+        "npm",
+        "pnpm",
+        "bun",
+    }:
         return SafetyCheckResult(True)
 
     if cmd0 == "sed":

klaude_code/core/tool/sub_agent_tool.py

@@ -12,11 +12,10 @@ from typing import TYPE_CHECKING, ClassVar
 
 from klaude_code.core.tool.tool_abc import ToolABC
 from klaude_code.core.tool.tool_context import current_run_subtask_callback
-from klaude_code.protocol.llm_parameter import ToolSchema
-from klaude_code.protocol.model import SubAgentState, ToolResultItem, ToolResultUIExtra, ToolResultUIExtraType
+from klaude_code.protocol import llm_param, model
 
 if TYPE_CHECKING:
-    from klaude_code.core.sub_agent import SubAgentProfile
+    from klaude_code.protocol.sub_agent import SubAgentProfile
 
 
 class SubAgentTool(ToolABC):
@@ -38,9 +37,9 @@ class SubAgentTool(ToolABC):
         )
 
     @classmethod
-    def schema(cls) -> ToolSchema:
+    def schema(cls) -> llm_param.ToolSchema:
         profile = cls._profile
-        return ToolSchema(
+        return llm_param.ToolSchema(
             name=profile.name,
             type="function",
             description=profile.description,
@@ -48,17 +47,17 @@ class SubAgentTool(ToolABC):
         )
 
     @classmethod
-    async def call(cls, arguments: str) -> ToolResultItem:
+    async def call(cls, arguments: str) -> model.ToolResultItem:
         profile = cls._profile
 
         try:
             args = json.loads(arguments)
         except json.JSONDecodeError as e:
-            return ToolResultItem(status="error", output=f"Invalid JSON arguments: {e}")
+            return model.ToolResultItem(status="error", output=f"Invalid JSON arguments: {e}")
 
         runner = current_run_subtask_callback.get()
         if runner is None:
-            return ToolResultItem(status="error", output="No subtask runner available in this context")
+            return model.ToolResultItem(status="error", output="No subtask runner available in this context")
 
         # Build the prompt using the profile's prompt builder
         prompt = profile.prompt_builder(args)
@@ -66,7 +65,7 @@ class SubAgentTool(ToolABC):
 
         try:
             result = await runner(
-                SubAgentState(
+                model.SubAgentState(
                     sub_agent_type=profile.name,
                     sub_agent_desc=description,
                     sub_agent_prompt=prompt,
@@ -75,10 +74,10 @@ class SubAgentTool(ToolABC):
         except asyncio.CancelledError:
             raise
         except Exception as e:
-            return ToolResultItem(status="error", output=f"Failed to run subtask: {e}")
+            return model.ToolResultItem(status="error", output=f"Failed to run subtask: {e}")
 
-        return ToolResultItem(
+        return model.ToolResultItem(
             status="success" if not result.error else "error",
             output=result.task_result or "",
-            ui_extra=ToolResultUIExtra(type=ToolResultUIExtraType.SESSION_ID, session_id=result.session_id),
+            ui_extra=model.ToolResultUIExtra(type=model.ToolResultUIExtraType.SESSION_ID, session_id=result.session_id),
         )