coding-agent/skills/agent-builder/references/tool-templates.py

"""
工具模板 - 复制后可按你的 Agent 定制。

每个工具都需要：
1. 定义（给模型看的 JSON schema）
2. 实现（Python 函数）
"""

from pathlib import Path
import subprocess

WORKDIR = Path.cwd()


# =============================================================================
# 工具定义（放入 TOOLS 列表）
# =============================================================================

BASH_TOOL = {
    "name": "bash",
    "description": "运行 shell 命令。适用于 ls、find、grep、git、npm、python 等。",
    "input_schema": {
        "type": "object",
        "properties": {
            "command": {
                "type": "string",
                "description": "要执行的 shell 命令"
            }
        },
        "required": ["command"],
    },
}

READ_FILE_TOOL = {
    "name": "read_file",
    "description": "读取文件内容，返回 UTF-8 文本。",
    "input_schema": {
        "type": "object",
        "properties": {
            "path": {
                "type": "string",
                "description": "文件的相对路径"
            },
            "limit": {
                "type": "integer",
                "description": "最多读取多少行（默认全部读取）"
            },
        },
        "required": ["path"],
    },
}

WRITE_FILE_TOOL = {
    "name": "write_file",
    "description": "向文件写入内容。需要时会创建父目录。",
    "input_schema": {
        "type": "object",
        "properties": {
            "path": {
                "type": "string",
                "description": "文件的相对路径"
            },
            "content": {
                "type": "string",
                "description": "要写入的内容"
            },
        },
        "required": ["path", "content"],
    },
}

EDIT_FILE_TOOL = {
    "name": "edit_file",
    "description": "在文件中替换一次完全匹配的文本，适合小范围精确修改。",
    "input_schema": {
        "type": "object",
        "properties": {
            "path": {
                "type": "string",
                "description": "文件的相对路径"
            },
            "old_text": {
                "type": "string",
                "description": "要查找的精确文本（必须完全匹配）"
            },
            "new_text": {
                "type": "string",
                "description": "替换后的文本"
            },
        },
        "required": ["path", "old_text", "new_text"],
    },
}

TODO_WRITE_TOOL = {
    "name": "TodoWrite",
    "description": "更新任务清单，用于规划并跟踪进度。",
    "input_schema": {
        "type": "object",
        "properties": {
            "items": {
                "type": "array",
                "description": "完整任务列表",
                "items": {
                    "type": "object",
                    "properties": {
                        "content": {"type": "string", "description": "任务描述"},
                        "status": {"type": "string", "enum": ["pending", "in_progress", "completed"]},
                        "activeForm": {"type": "string", "description": "当前进行态描述，例如“正在读取文件”"},
                    },
                    "required": ["content", "status", "activeForm"],
                },
            }
        },
        "required": ["items"],
    },
}

TASK_TOOL_TEMPLATE = """
# 根据 Agent 类型动态生成
TASK_TOOL = {
    "name": "Task",
    "description": f"为聚焦子任务启动一个子 Agent。\\n\\nAgent 类型：\\n{get_agent_descriptions()}",
    "input_schema": {
        "type": "object",
        "properties": {
            "description": {"type": "string", "description": "简短任务名（3-5 个词）"},
            "prompt": {"type": "string", "description": "详细指令"},
            "agent_type": {"type": "string", "enum": list(AGENT_TYPES.keys())},
        },
        "required": ["description", "prompt", "agent_type"],
    },
}
"""


# =============================================================================
# 工具实现
# =============================================================================

def safe_path(p: str) -> Path:
    """
    安全检查：确保路径留在工作区内。
    防止 ../../../etc/passwd 这类路径逃逸攻击。
    """
    path = (WORKDIR / p).resolve()
    if not path.is_relative_to(WORKDIR):
        raise ValueError(f"路径逃逸出工作区：{p}")
    return path


def run_bash(command: str) -> str:
    """
    带安全检查地执行 shell 命令。

    安全特性：
    - 阻止明显危险的命令
    - 60 秒超时
    - 输出截断到 50KB
    """
    dangerous = ["rm -rf /", "sudo", "shutdown", "reboot", "> /dev/"]
    if any(d in command for d in dangerous):
        return "错误：已阻止危险命令"

    try:
        result = subprocess.run(
            command,
            shell=True,
            cwd=WORKDIR,
            capture_output=True,
            text=True,
            timeout=60
        )
        output = (result.stdout + result.stderr).strip()
        return output[:50000] if output else "（无输出）"

    except subprocess.TimeoutExpired:
        return "错误：命令执行超时（60 秒）"
    except Exception as e:
        return f"错误：{e}"


def run_read_file(path: str, limit: int = None) -> str:
    """
    读取文件内容，可选行数限制。

    特性：
    - 安全路径解析
    - 大文件可限制读取行数
    - 输出截断到 50KB
    """
    try:
        text = safe_path(path).read_text()
        lines = text.splitlines()

        if limit and limit < len(lines):
            lines = lines[:limit]
            lines.append(f"...（还有 {len(text.splitlines()) - limit} 行）")

        return "\n".join(lines)[:50000]

    except Exception as e:
        return f"错误：{e}"


def run_write_file(path: str, content: str) -> str:
    """
    向文件写入内容，需要时创建父目录。

    特性：
    - 安全路径解析
    - 自动创建父目录
    - 返回写入字节数用于确认
    """
    try:
        fp = safe_path(path)
        fp.parent.mkdir(parents=True, exist_ok=True)
        fp.write_text(content)
        return f"已写入 {len(content)} 字节到 {path}"

    except Exception as e:
        return f"错误：{e}"


def run_edit_file(path: str, old_text: str, new_text: str) -> str:
    """
    在文件中替换一次完全匹配的文本（精确修改）。

    特性：
    - 精确字符串匹配（不是正则）
    - 只替换首次出现的位置（更安全）
    - 找不到文本时返回清晰错误
    """
    try:
        fp = safe_path(path)
        content = fp.read_text()

        if old_text not in content:
            return f"错误：在 {path} 中未找到目标文本"

        new_content = content.replace(old_text, new_text, 1)
        fp.write_text(new_content)
        return f"已编辑 {path}"

    except Exception as e:
        return f"错误：{e}"


# =============================================================================
# 分发器模式
# =============================================================================

def execute_tool(name: str, args: dict) -> str:
    """
    将工具调用分发到对应实现。

    这种模式便于添加新工具：
    1. 将工具定义加入 TOOLS 列表
    2. 添加实现函数
    3. 在这个分发器中增加分支
    """
    if name == "bash":
        return run_bash(args["command"])
    if name == "read_file":
        return run_read_file(args["path"], args.get("limit"))
    if name == "write_file":
        return run_write_file(args["path"], args["content"])
    if name == "edit_file":
        return run_edit_file(args["path"], args["old_text"], args["new_text"])
    # 在这里继续添加更多工具...
    return f"未知工具：{name}"