第五章：工具层设计¶

"""
Claude Code 风格的工具接口
"""

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Generic, TypeVar, Optional, Dict, Any, List
import asyncio

# 类型变量
# （Claude Code 的 Tool<Input,Output,Progress> 三泛型中的 Progress，
#   本简化实现用下方具体的 ToolProgress 类承载，不再单设类型变量）
InputType = TypeVar('InputType')
OutputType = TypeVar('OutputType')

@dataclass
class ToolProgress:
    """工具执行进度"""
    step: int              # 当前步骤号
    total_steps: int       # 总步骤数
    current_status: str    # 当前状态描述
    estimated_time_remaining: Optional[float] = None  # 剩余时间(秒)

    def percentage(self) -> float:
        """进度百分比"""
        return (self.step / self.total_steps * 100) if self.total_steps > 0 else 0

class Tool(ABC, Generic[InputType, OutputType]):
    """通用工具抽象"""

    @abstractmethod
    async def call(self, input_data: InputType) -> OutputType:
        """
        执行工具

        Args:
            input_data: 工具输入,已验证和类型检查过

        Returns:
            工具执行结果

        Raises:
            ToolExecutionError: 工具执行失败
        """
        pass

    @abstractmethod
    def name(self) -> str:
        """工具名称"""
        pass

    @abstractmethod
    def description(self) -> str:
        """工具描述"""
        pass

    @abstractmethod
    def input_schema(self) -> Dict[str, Any]:
        """输入的 JSON Schema"""
        pass

    def check_permissions(self, context: Any) -> bool:
        """
        检查当前上下文是否有权限调用此工具

        Args:
            context: 执行上下文(包含用户、会话信息等)

        Returns:
            是否有权限
        """
        return True  # 默认允许

    async def get_progress(self) -> Optional[ToolProgress]:
        """
        获取工具执行进度

        Returns:
            进度对象,如果工具不支持进度报告则返回 None
        """
        return None

    def supports_streaming(self) -> bool:
        """工具是否支持流式输出"""
        return False

    async def stream_output(self, input_data: InputType):
        """
        流式输出(可选)

        Yields:
            OutputType: 逐个产生的输出
        """
        if not self.supports_streaming():
            output = await self.call(input_data)
            yield output

@dataclass
class BashInput:
    command: str
    timeout: int = 30
    cwd: Optional[str] = None

@dataclass
class BashOutput:
    stdout: str
    stderr: str
    returncode: int
    execution_time: float

class BashTool(Tool[BashInput, BashOutput]):
    """Bash 命令执行工具"""

    async def call(self, input_data: BashInput) -> BashOutput:
        """执行 bash 命令"""
        import shlex
        import subprocess
        import time

        start_time = time.time()
        allowed_commands = {"echo", "pwd", "true", "false", "sleep"}
        argv = shlex.split(input_data.command)
        if not argv or argv[0] not in allowed_commands:
            raise ToolExecutionError("Command is not allowed by policy")

        try:
            result = subprocess.run(
                argv,
                shell=False,
                capture_output=True,
                timeout=input_data.timeout,
                text=True,
                cwd=input_data.cwd
            )

            return BashOutput(
                stdout=result.stdout,
                stderr=result.stderr,
                returncode=result.returncode,
                execution_time=time.time() - start_time
            )

        except subprocess.TimeoutExpired:
            raise ToolExecutionError(
                f"Bash command timeout after {input_data.timeout}s"
            )
        except Exception as e:
            raise ToolExecutionError(f"Bash execution failed: {str(e)}")

    def name(self) -> str:
        return "bash_exec"

    def description(self) -> str:
        return "Execute bash commands on the system"

    def input_schema(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "command": {
                    "type": "string",
                    "description": "The bash command to execute"
                },
                "timeout": {
                    "type": "integer",
                    "description": "Command timeout in seconds",
                    "default": 30
                },
                "cwd": {
                    "type": "string",
                    "description": "Current working directory"
                }
            },
            "required": ["command"]
        }

    def check_permissions(self, context: Any) -> bool:
        """Bash 工具需要特殊权限"""
        if hasattr(context, 'is_admin'):
            return context.is_admin
        return False

@dataclass
class FileCopyInput:
    source: str
    destination: str
    overwrite: bool = False

@dataclass
class FileCopyOutput:
    success: bool
    bytes_copied: int
    destination: str

class FileCopyTool(Tool[FileCopyInput, FileCopyOutput]):
    """文件复制工具,支持进度报告"""

    def __init__(self):
        self._current_progress = None

    async def call(self, input_data: FileCopyInput) -> FileCopyOutput:
        """复制文件"""
        import shutil
        import os
        from pathlib import Path
        from mini_harness.security.path_validator import PathValidator

        validator = PathValidator(base_path="./workspace")
        source = Path(validator.validate(input_data.source))
        destination = Path(validator.validate(input_data.destination))

        # 检查源文件是否存在
        if not source.exists() or source.is_symlink():
            raise ToolExecutionError(f"Source file not found or not allowed: {input_data.source}")

        # 检查目标是否已存在
        if destination.exists() and not input_data.overwrite:
            raise ToolExecutionError(
                f"Destination already exists: {input_data.destination}"
            )
        if destination.exists() and destination.is_symlink():
            raise ToolExecutionError("Symlink destinations are not allowed")

        # 获取文件大小用于进度报告
        file_size = source.stat().st_size
        chunk_size = 1024 * 1024  # 1MB chunks

        bytes_copied = 0

        # 自定义复制逻辑以支持进度报告
        destination.parent.mkdir(parents=True, exist_ok=True)
        with source.open('rb') as src, destination.open('wb') as dst:
            while True:
                chunk = src.read(chunk_size)
                if not chunk:
                    break

                dst.write(chunk)
                bytes_copied += len(chunk)

                # 更新进度
                self._current_progress = ToolProgress(
                    step=bytes_copied,
                    total_steps=file_size,
                    current_status=f"Copying {bytes_copied}/{file_size} bytes"
                )

        return FileCopyOutput(
            success=True,
            bytes_copied=bytes_copied,
            destination=input_data.destination
        )

    async def get_progress(self) -> Optional[ToolProgress]:
        """获取进度"""
        return self._current_progress

    def name(self) -> str:
        return "file_copy"

    def description(self) -> str:
        return "Copy files with progress reporting"

    def input_schema(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "source": {
                    "type": "string",
                    "description": "Source file path"
                },
                "destination": {
                    "type": "string",
                    "description": "Destination file path"
                },
                "overwrite": {
                    "type": "boolean",
                    "description": "Whether to overwrite existing file",
                    "default": False
                }
            },
            "required": ["source", "destination"]
        }

from dataclasses import dataclass
from typing import Dict, Any
import asyncio

@dataclass
class DatabaseQueryInput:
    query: str
    limit: int = 1000

@dataclass
class DatabaseQueryRow:
    columns: Dict[str, Any]

class DatabaseQueryTool(Tool[DatabaseQueryInput, DatabaseQueryRow]):
    """数据库查询工具,支持流式输出"""

    def __init__(self, connection_string: str):
        self.connection_string = connection_string

    async def call(self, input_data: DatabaseQueryInput) -> str:
        """执行查询并返回结果"""
        # 这个实现中,我们只返回所有结果的摘要
        # 实际使用应该使用流式输出
        all_rows = []
        async for row in self.stream_output(input_data):
            all_rows.append(row)

        return f"Query returned {len(all_rows)} rows"

    async def stream_output(self, input_data: DatabaseQueryInput):
        """流式输出查询结果"""
        # 在实际实现中,这里会连接到真实的数据库
        # 示例:连接到 PostgreSQL、MySQL 等

        # 模拟流式输出
        for i in range(min(input_data.limit, 10)):
            yield DatabaseQueryRow(
                columns={
                    "id": i,
                    "name": f"Row {i}",
                    "value": i * 100
                }
            )
            await asyncio.sleep(0.1)  # 模拟网络延迟

    def supports_streaming(self) -> bool:
        return True

    def name(self) -> str:
        return "database_query"

    def description(self) -> str:
        return "Query database with streaming results"

    def input_schema(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "SQL query to execute"
                },
                "limit": {
                    "type": "integer",
                    "description": "Maximum number of rows to return",
                    "default": 1000
                }
            },
            "required": ["query"]
        }

def build_tool(tool_spec: Dict[str, Any]) -> Tool:
    """
    工厂函数:根据规范构造工具

    Args:
        tool_spec: 工具规范
            {
                "type": "bash" | "file" | "network" | ...,
                "name": "tool_name",
                "config": {...}  # 工具特定的配置
            }

    Returns:
        构造的工具实例
    """

    tool_type = tool_spec.get("type")
    config = tool_spec.get("config", {})

    if tool_type == "bash":
        return BashTool()

    elif tool_type == "file_copy":
        return FileCopyTool()

    elif tool_type == "database_query":
        return DatabaseQueryTool(
            connection_string=config.get("connection_string")
        )

    elif tool_type == "custom":
        # 只允许从受信任配置和受控命名空间加载自定义工具
        module_path = config.get("module")
        class_name = config.get("class")
        allowed_prefixes = ("my_harness_tools.",)
        if not module_path or not module_path.startswith(allowed_prefixes):
            raise ToolExecutionError("Custom tool module is not in the trusted namespace")

        # 使用 importlib 动态加载
        import importlib
        module = importlib.import_module(module_path)
        tool_class = getattr(module, class_name)
        return tool_class(**config.get("args", {}))

    else:
        raise ValueError(f"Unknown tool type: {tool_type}")

class ToolRegistry:
    """工具注册中心,管理所有可用工具"""

    def __init__(self):
        self.tools: Dict[str, Tool] = {}
        self.tool_cache: Dict[str, Dict[str, Any]] = {}  # Schema 缓存

    def register(self, tool: Tool):
        """注册工具"""
        name = tool.name()
        self.tools[name] = tool

        # 缓存工具的 Schema
        self.tool_cache[name] = {
            "name": name,
            "description": tool.description(),
            "input_schema": tool.input_schema()
        }

        print(f"Tool registered: {name}")

    def get(self, name: str) -> Optional[Tool]:
        """获取工具"""
        return self.tools.get(name)

    def list_tools(self) -> List[Dict[str, Any]]:
        """列出所有工具的 Schema"""
        return list(self.tool_cache.values())

    def get_tool_schema(self, name: str) -> Optional[Dict[str, Any]]:
        """获取单个工具的 Schema"""
        return self.tool_cache.get(name)

    def tool_exists(self, name: str) -> bool:
        """检查工具是否存在"""
        return name in self.tools

    def unregister(self, name: str):
        """注销工具"""
        if name in self.tools:
            del self.tools[name]
            del self.tool_cache[name]
            print(f"Tool unregistered: {name}")

# 使用示例
registry = ToolRegistry()

# 注册内置工具
registry.register(BashTool())
registry.register(FileCopyTool())
registry.register(DatabaseQueryTool("postgresql://localhost/mydb"))

# 列出所有工具
for tool_schema in registry.list_tools():
    print(f"- {tool_schema['name']}: {tool_schema['description']}")

# 获取特定工具
bash_tool = registry.get("bash_exec")
if bash_tool:
    print(f"Found tool: {bash_tool.name()}")

class ToolStrategy(Enum):
    """工具使用策略"""
    ALLOW = "allow"           # 允许使用
    DENY = "deny"            # 禁止使用
    REQUIRE_APPROVAL = "require_approval"  # 需要用户批准

class ToolPolicy:
    """工具策略:定义哪些智能体可以使用哪些工具"""

    def __init__(self):
        self.policies: Dict[str, Dict[str, ToolStrategy]] = {}
        # policies[agent_id][tool_name] = strategy

    def set_policy(self, agent_id: str, tool_name: str,
                  strategy: ToolStrategy):
        """设置工具策略"""
        if agent_id not in self.policies:
            self.policies[agent_id] = {}
        self.policies[agent_id][tool_name] = strategy

    def can_use_tool(self, agent_id: str, tool_name: str) -> bool:
        """检查智能体是否可以使用工具"""
        if agent_id not in self.policies:
            return True  # 默认允许

        strategy = self.policies[agent_id].get(
            tool_name,
            ToolStrategy.ALLOW
        )

        return strategy != ToolStrategy.DENY

    def requires_approval(self, agent_id: str, tool_name: str) -> bool:
        """检查是否需要审批"""
        if agent_id not in self.policies:
            return False

        strategy = self.policies[agent_id].get(tool_name)
        return strategy == ToolStrategy.REQUIRE_APPROVAL

"""工具执行流水线"""

from dataclasses import dataclass
from typing import Optional, Dict, Any, Callable, Tuple
# 使用第4章定义的消息块类型（anthropic SDK 只提供请求侧的 ToolResultBlockParam，
# 响应侧没有 ToolResultBlock；这里沿用本书 mini_harness 的自定义类型）
from mini_harness.runtime.models import ToolUseBlock, ToolResultBlock
import asyncio
import time
import json

class ExecutionPipeline:
    """工具执行流水线"""

    def __init__(self,
                 tool_registry: ToolRegistry,
                 max_result_size: int = 1024 * 1024,
                 default_timeout: int = 30):
        self.tool_registry = tool_registry
        self.max_result_size = max_result_size
        self.default_timeout = default_timeout
        self.execution_history = []
        self.result_cache = {}

    async def execute(self,
                     tool_use: ToolUseBlock,
                     context: Any,
                     progress_callback: Optional[Callable] = None) -> ToolResultBlock:
        """执行工具的完整流水线"""
        execution_record = {
            "tool_use_id": tool_use.id,
            "tool_name": tool_use.name,
            "start_time": time.time(),
            "stages": {}
        }

        try:
            # 阶段1:工具发现
            tool, error = await self._find_tool(tool_use.name, execution_record)
            if error:
                return self._create_error_result(tool_use, error)

            # 阶段2:权限检查
            is_allowed, error = await self._check_permissions(
                tool, context, execution_record
            )
            if not is_allowed:
                return self._create_error_result(tool_use, error)

            # 阶段3:输入验证
            is_valid, error = await self._validate_input(
                tool, tool_use.input, execution_record
            )
            if not is_valid:
                return self._create_error_result(tool_use, error)

            # 阶段4:工具执行
            result, error = await self._execute_tool(
                tool, tool_use.input, execution_record, progress_callback
            )
            if error:
                return self._create_error_result(tool_use, error)

            # 阶段5:结果处理
            processed_result, error = await self._process_result(
                result, execution_record
            )
            if error:
                return self._create_error_result(tool_use, error)

            # 阶段6:本地缓存与记录交接
            await self._record_result(
                tool_use.id, processed_result, execution_record
            )

            self.execution_history.append(execution_record)

            return ToolResultBlock(
                tool_use_id=tool_use.id,
                content=str(processed_result),
                is_error=False
            )

        except Exception as e:
            execution_record["error"] = str(e)
            self.execution_history.append(execution_record)
            return self._create_error_result(tool_use, str(e))

    async def _find_tool(self, tool_name: str,
                        execution_record: Dict) -> Tuple[Optional[Tool], Optional[str]]:
        """阶段1:工具发现"""
        start = time.time()
        tool = self.tool_registry.get(tool_name)
        execution_record["stages"]["find_tool"] = {
            "duration": time.time() - start,
            "found": tool is not None
        }
        if not tool:
            return None, f"Tool '{tool_name}' not found"
        return tool, None

    async def _check_permissions(self, tool: Tool, context: Any,
                                execution_record: Dict) -> Tuple[bool, Optional[str]]:
        """阶段2:权限检查"""
        start = time.time()
        try:
            has_permission = tool.check_permissions(context)
            execution_record["stages"]["check_permissions"] = {
                "duration": time.time() - start,
                "result": has_permission
            }
            if not has_permission:
                return False, f"Permission denied for tool '{tool.name()}'"
            return True, None
        except Exception as e:
            return False, f"Permission check error: {str(e)}"

    async def _validate_input(self, tool: Tool, input_data: Dict[str, Any],
                             execution_record: Dict) -> Tuple[bool, Optional[str]]:
        """阶段3:输入验证"""
        start = time.time()
        schema = tool.input_schema()
        errors = []

        required = schema.get("required", [])
        for field in required:
            if field not in input_data:
                errors.append(f"Missing required field: {field}")

        properties = schema.get("properties", {})
        for field_name, field_value in input_data.items():
            if field_name not in properties:
                errors.append(f"Unknown field: {field_name}")
                continue
            field_schema = properties[field_name]
            if not self._check_type(field_value, field_schema):
                expected_type = field_schema.get("type", "unknown")
                errors.append(
                    f"Field '{field_name}' validation failed: "
                    f"expected {expected_type}"
                )

        execution_record["stages"]["validate_input"] = {
            "duration": time.time() - start,
            "valid": len(errors) == 0,
            "errors": errors
        }

        if errors:
            return False, "; ".join(errors)
        return True, None

    def _check_type(self, value: Any, schema: Dict[str, Any]) -> bool:
        """基于JSON Schema的类型验证"""
        import jsonschema
        try:
            jsonschema.validate(value, schema)
            return True
        except jsonschema.ValidationError:
            return False

    async def _execute_tool(self,
                           tool: Tool,
                           input_data: Dict[str, Any],
                           execution_record: Dict,
                           progress_callback: Optional[Callable] = None
                           ) -> Tuple[Any, Optional[str]]:
        """阶段4:工具执行(带超时控制)"""
        start = time.time()
        try:
            result = await asyncio.wait_for(
                tool.call(input_data),
                timeout=self.default_timeout
            )
            execution_record["stages"]["execute_tool"] = {
                "duration": time.time() - start,
                "success": True
            }
            return result, None
        except asyncio.TimeoutError:
            error = f"Tool execution timeout ({self.default_timeout}s)"
            execution_record["stages"]["execute_tool"] = {
                "duration": time.time() - start,
                "success": False,
                "error": "timeout"
            }
            return None, error
        except Exception as e:
            execution_record["stages"]["execute_tool"] = {
                "duration": time.time() - start,
                "success": False,
                "error": type(e).__name__
            }
            return None, f"Tool execution error: {str(e)}"

    async def _process_result(self, result: Any,
                             execution_record: Dict) -> Tuple[str, Optional[str]]:
        """阶段5:结果处理(序列化和截断)"""
        start = time.time()
        try:
            if isinstance(result, str):
                result_str = result
            elif isinstance(result, (dict, list)):
                result_str = json.dumps(result, indent=2)
            else:
                result_str = str(result)

            if len(result_str) > self.max_result_size:
                original_size = len(result_str)
                result_str = result_str[:self.max_result_size]
                result_str += f"\n... [Output truncated from {original_size} bytes]"
                execution_record["stages"]["process_result"] = {
                    "duration": time.time() - start,
                    "truncated": True,
                    "original_size": original_size,
                }
            else:
                execution_record["stages"]["process_result"] = {
                    "duration": time.time() - start,
                    "truncated": False,
                    "size": len(result_str)
                }
            return result_str, None
        except Exception as e:
            return None, f"Result processing error: {str(e)}"

    async def _record_result(self, tool_use_id: str, result: str,
                             execution_record: Dict):
        """阶段6:本地结果缓存，不写入记忆系统"""
        start = time.time()
        self.result_cache[tool_use_id] = {
            "content": result,
            "timestamp": time.time()
        }
        execution_record["stages"]["record_result"] = {
            "duration": time.time() - start,
            "cached": True
        }

    def _create_error_result(self, tool_use: ToolUseBlock,
                            error: str) -> ToolResultBlock:
        """创建错误结果"""
        return ToolResultBlock(
            tool_use_id=tool_use.id,
            content=error,
            is_error=True,
            error_type="ExecutionError"
        )

    def get_execution_history(self, tool_use_id: Optional[str] = None):
        """获取执行历史"""
        if tool_use_id:
            return [r for r in self.execution_history
                   if r["tool_use_id"] == tool_use_id]
        return self.execution_history

    def clear_cache(self, max_age: int = 3600):
        """清理缓存(删除超过max_age秒的项)"""
        now = time.time()
        to_delete = [
            key for key, value in self.result_cache.items()
            if now - value["timestamp"] > max_age
        ]
        for key in to_delete:
            del self.result_cache[key]
        return len(to_delete)

from typing import Optional, Dict
import asyncio
import time

class ProgressMonitor:
    """进度监控器"""

    def __init__(self, tool_use_id: str):
        self.tool_use_id = tool_use_id
        self.progress_events = []

    def report_progress(self, step: int, total_steps: int,
                       status: str, estimated_time_remaining: Optional[float] = None):
        """报告进度"""
        event = {
            "tool_use_id": self.tool_use_id,
            "timestamp": time.time(),
            "step": step,
            "total_steps": total_steps,
            "percentage": (step / total_steps * 100) if total_steps > 0 else 0,
            "status": status,
            "estimated_time_remaining": estimated_time_remaining
        }
        self.progress_events.append(event)
        return event

    def get_current_progress(self) -> Optional[Dict]:
        """获取当前进度"""
        return self.progress_events[-1] if self.progress_events else None

# 在工具执行中使用
async def long_running_tool_call(progress_monitor: ProgressMonitor):
    """长时间运行的工具调用示例"""
    total_items = 1000

    for i in range(total_items):
        # 执行工作
        await asyncio.sleep(0.01)  # 模拟工作

        # 报告进度
        if i % 100 == 0:
            progress_monitor.report_progress(
                step=i,
                total_steps=total_items,
                status=f"Processing item {i}/{total_items}",
                estimated_time_remaining=(total_items - i) * 0.01
            )

    return f"Processed {total_items} items"

from typing import List, Dict, Any
import asyncio

class StreamingToolExecutor:
    """并发工具执行器"""

    def __init__(self, pipeline: ExecutionPipeline,
                 max_concurrent: int = 5):
        self.pipeline = pipeline
        self.max_concurrent = max_concurrent

    async def execute_tools(self,
                           tool_uses: List[ToolUseBlock],
                           context: Any) -> Dict[str, ToolResultBlock]:
        """
        并发执行多个工具

        Args:
            tool_uses: 工具调用列表
            context: 执行上下文

        Returns:
            {tool_use_id: tool_result_block}
        """

        semaphore = asyncio.Semaphore(self.max_concurrent)

        async def execute_with_semaphore(tool_use):
            async with semaphore:
                return await self.pipeline.execute(tool_use, context)

        results = await asyncio.gather(
            *[execute_with_semaphore(tu) for tu in tool_uses],
            return_exceptions=True
        )

        # 组织结果：gather 的返回顺序与 tool_uses 一一对应；
        # return_exceptions=True 时异常项是裸 Exception 对象，不能直接解包
        result_dict = {}
        for tool_use, result in zip(tool_uses, results):
            if isinstance(result, Exception):
                result_dict[tool_use.id] = ToolResultBlock(
                    tool_use_id=tool_use.id,
                    content=str(result),
                    is_error=True,
                    error_type=type(result).__name__
                )
            else:
                result_dict[tool_use.id] = result

        return result_dict

class BashTool(Tool):
    """执行 bash 命令"""
    def name(self) -> str: return "bash_exec"

class PythonTool(Tool):
    """执行 Python 代码"""
    def name(self) -> str: return "python_exec"

class FileReadTool(Tool):
    """读取文件内容"""
    def name(self) -> str: return "file_read"

class FileWriteTool(Tool):
    """写入文件"""
    def name(self) -> str: return "file_write"

class FileListTool(Tool):
    """列出目录内容"""
    def name(self) -> str: return "file_list"

class HTTPTool(Tool):
    """发送 HTTP 请求"""
    def name(self) -> str: return "http_request"

class WebSearchTool(Tool):
    """网络搜索"""
    def name(self) -> str: return "web_search"

class APIDelegationTool(Tool):
    """调用第三方 API"""
    def name(self) -> str: return "api_call"

class SubAgentTool(Tool):
    """创建和委托子任务给另一个 Agent"""
    def name(self) -> str: return "delegate_to_agent"

class ParallelAgentsTool(Tool):
    """并行执行多个 Agent"""
    def name(self) -> str: return "parallel_agents"

class DatabaseQueryTool(Tool):
    """数据库查询"""
    def name(self) -> str: return "database_query"

class ImageProcessingTool(Tool):
    """图像处理"""
    def name(self) -> str: return "image_process"

class CodeAnalysisTool(Tool):
    """代码分析(静态或动态)"""
    def name(self) -> str: return "code_analyze"

class DataVisualizationTool(Tool):
    """数据可视化"""
    def name(self) -> str: return "plot_data"

class ToolAccessLevel(Enum):
    """工具访问级别"""
    SYSTEM = "system"      # 仅系统可用
    ADMIN = "admin"        # 管理员可用
    TRUSTED = "trusted"    # 信任的用户可用
    PUBLIC = "public"      # 所有人可用

class ToolPolicy:
    """工具策略定义"""

    def __init__(self):
        self.access_levels: Dict[str, ToolAccessLevel] = {}
        self.resource_limits: Dict[str, Dict[str, Any]] = {}

    def set_access_level(self, tool_name: str,
                        level: ToolAccessLevel):
        """设置工具访问级别"""
        self.access_levels[tool_name] = level

    def set_resource_limit(self, tool_name: str,
                          max_execution_time: int = None,
                          max_memory: int = None,
                          max_disk_usage: int = None):
        """设置工具的资源限制"""
        self.resource_limits[tool_name] = {
            "max_execution_time": max_execution_time,
            "max_memory": max_memory,
            "max_disk_usage": max_disk_usage
        }

    def can_execute(self, tool_name: str,
                   user_role: str) -> bool:
        """检查用户是否可以执行工具"""
        level = self.access_levels.get(
            tool_name,
            ToolAccessLevel.PUBLIC
        )

        if level == ToolAccessLevel.SYSTEM:
            return user_role == "system"
        elif level == ToolAccessLevel.ADMIN:
            return user_role in ["system", "admin"]
        elif level == ToolAccessLevel.TRUSTED:
            return user_role in ["system", "admin", "trusted"]
        else:
            return True

import asyncio

def timeout_wrapper(tool: Tool, timeout: int = 30) -> Tool:
    """为工具添加超时装饰器"""
    class TimeoutTool(Tool):
        async def call(self, input_data):
            try:
                return await asyncio.wait_for(
                    tool.call(input_data),
                    timeout=timeout
                )
            except asyncio.TimeoutError:
                raise ToolExecutionError(f"Tool timeout after {timeout}s")

        def name(self) -> str:
            return tool.name()

        def description(self) -> str:
            return tool.description()

        # 其他方法委托...

    return TimeoutTool()

def retry_wrapper(tool: Tool, max_retries: int = 3) -> Tool:
    """为工具添加重试装饰器"""
    class RetryTool(Tool):
        async def call(self, input_data):
            for attempt in range(max_retries):
                try:
                    return await tool.call(input_data)
                except Exception as e:
                    if attempt == max_retries - 1:
                        raise
                    await asyncio.sleep(2 ** attempt)  # 指数退避

        # 其他方法...

    return RetryTool()

from typing import List, Dict

class CompositeToolChain(Tool):
    """工具链:顺序执行多个工具"""

    def __init__(self, tools: List[Tool],
                 input_mapping: Dict[str, List[str]]):
        """
        input_mapping: 定义每个工具的输入如何来自前一个工具的输出
        （上下文键名格式为 "<工具名>_output"，与 call() 中的写入格式一致）
        例如:{
            "tool_2": ["tool_1_output.field_a"],
            "tool_3": ["tool_2_output.field_b", "user_input.param"]
        }
        """
        self.tools = tools
        self.input_mapping = input_mapping

    async def call(self, input_data):
        """依次执行工具"""
        context = {"user_input": input_data}

        for tool in self.tools:
            tool_name = tool.name()

            # 解析这个工具的输入
            tool_input = self._resolve_input(
                self.input_mapping.get(tool_name, []),
                context
            )

            # 执行工具
            result = await tool.call(tool_input)
            context[f"{tool_name}_output"] = result

        return context[f"{self.tools[-1].name()}_output"]

    def _resolve_input(self, mapping: List[str],
                      context: Dict) -> Dict:
        """从上下文中解析输入"""
        result = {}
        for ref in mapping:
            parts = ref.split(".")
            value = context
            for part in parts:
                value = value.get(part, {})
            # 提取字段名
            result[parts[-1]] = value
        return result

class ConditionalTool(Tool):
    """条件执行:根据条件选择执行哪个工具"""

    def __init__(self, condition_fn, tool_if_true: Tool,
                 tool_if_false: Tool):
        self.condition_fn = condition_fn
        self.tool_if_true = tool_if_true
        self.tool_if_false = tool_if_false

    async def call(self, input_data):
        """根据条件选择执行"""
        if self.condition_fn(input_data):
            return await self.tool_if_true.call(input_data)
        else:
            return await self.tool_if_false.call(input_data)

    def name(self) -> str:
        return "conditional_tool"

@dataclass
class ToolCapability:
    """工具的能力描述"""
    name: str
    description: str
    input_schema: Dict[str, Any]
    output_schema: Dict[str, Any]
    access_level: ToolAccessLevel
    required_permissions: List[str]
    estimated_execution_time: float  # 秒
    estimated_resource_usage: Dict[str, int]  # 内存、CPU 等
    examples: List[Dict[str, Any]]  # 使用示例
    limitations: List[str]  # 工具的限制
    dependencies: List[str]  # 外部依赖

class ToolLoader:
    """延迟加载机制"""

    def __init__(self):
        self.loaded_tools = {}
        self.tool_metadata = {}

    def should_defer_loading(self, tool_name: str,
                            context: Dict) -> bool:
        """判断是否延迟加载工具"""
        # 示例:某些工具只在特定上下文才加载
        if tool_name == "database_query":
            return context.get("user_role") != "admin"
        return False

    async def load_tool_on_demand(self, tool_name: str,
                                 context: Dict) -> Optional[Tool]:
        """按需加载工具"""
        if tool_name in self.loaded_tools:
            return self.loaded_tools[tool_name]

        if self.should_defer_loading(tool_name, context):
            return None

        # 动态导入和加载
        try:
            spec = self.tool_metadata.get(tool_name)
            if not spec:
                return None

            module_path = spec["module"]
            if not module_path.startswith("my_harness_tools."):
                raise ValueError("Tool spec must use a trusted module namespace")

            module = await self._async_import(module_path)
            tool_class = getattr(module, spec["class"])
            tool = tool_class(**spec.get("args", {}))

            self.loaded_tools[tool_name] = tool
            return tool

        except Exception as e:
            print(f"Failed to load tool '{tool_name}': {e}")
            return None

    async def _async_import(self, module_path: str):
        """异步导入模块"""
        import importlib
        return importlib.import_module(module_path)

class SkillPriority(Enum):
    """技能加载优先级"""
    SYSTEM = 1          # 固定系统技能
    SESSION = 2         # 会话上下文技能
    USER = 3           # 用户自定义
    DOMAIN = 4         # 领域专业技能
    BUILTIN = 5        # 通用工具
    DISCOVERY = 6      # 动态发现

class SkillLoader:
    """6 级优先级技能加载器"""

    def __init__(self, storage):
        self.storage = storage
        self.loaded_skills = {}

    async def load_skills_for_session(self, session_id: str):
        """为会话加载所有技能"""
        skills = []

        # 1. 系统技能
        skills.extend(await self._load_system_skills())

        # 2. 会话上下文技能
        skills.extend(await self._load_session_skills(session_id))

        # 3. 用户自定义技能
        skills.extend(await self._load_user_skills(session_id))

        # 4. 领域技能
        domain = await self._detect_domain(session_id)
        skills.extend(await self._load_domain_skills(domain))

        # 5. 内置工具
        skills.extend(await self._load_builtin_tools())

        # 6. 动态发现技能
        skills.extend(await self._discover_skills_dynamically(session_id))

        return skills

    async def _load_system_skills(self):
        """加载系统技能(始终可用)"""
        return await self.storage.get_system_skills()

    async def _load_session_skills(self, session_id: str):
        """加载会话特定的技能"""
        return await self.storage.get_session_skills(session_id)

    async def _load_user_skills(self, session_id: str):
        """加载用户定义的技能"""
        user_id = await self.storage.get_user_for_session(session_id)
        return await self.storage.get_user_skills(user_id)

    async def _load_domain_skills(self, domain: str):
        """加载特定领域的技能"""
        if not domain:
            return []
        return await self.storage.get_domain_skills(domain)

    async def _load_builtin_tools(self):
        """加载内置工具"""
        return [
            BashTool(),
            FileReadTool(),
            FileWriteTool(),
            # ... 其他内置工具
        ]

    async def _discover_skills_dynamically(self, session_id: str):
        """动态发现技能"""
        # 扫描技能目录或 MCP
        discovered = await self.storage.discover_available_skills()
        return discovered

    async def _detect_domain(self, session_id: str) -> str:
        """根据会话内容检测领域"""
        recent_messages = await self.storage.get_recent_messages(session_id, limit=5)
        # 使用 NLP 检测领域
        text = " ".join(m.get_text() for m in recent_messages)
        return self._classify_domain(text)

    def _classify_domain(self, text: str) -> str:
        """简单的领域分类"""
        keywords = {
            "data": ["sql", "database", "query", "table"],
            "web": ["html", "css", "javascript", "react"],
            "system": ["bash", "shell", "linux", "windows"],
            "code": ["python", "java", "javascript"]
        }

        text_lower = text.lower()
        for domain, kws in keywords.items():
            if any(kw in text_lower for kw in kws):
                return domain
        return "general"