vent_agent_2.0/extensions/registry.py

"""确定性发现、声明和校验业务扩展；通过校验后一次性发布工具目录。"""
from __future__ import annotations

from dataclasses import dataclass, replace
import importlib
import re
from pathlib import Path
import threading
from typing import Any, Callable, Iterable

import config
from tools import registry as tools

_NAME = re.compile(r"[A-Za-z_][A-Za-z0-9_-]{0,63}\Z")
_RESERVED_TOOLS = {"write_todos", "ls", "read_file", "write_file", "edit_file", "glob", "grep", "execute", "task"}
_RESERVED_AGENTS = {"deep-agent", "general-purpose"}


def _name(value: str, kind: str) -> str:
    if not isinstance(value, str) or not _NAME.fullmatch(value):
        raise ValueError(f"{kind}名称非法：{value!r}（使用字母、数字、下划线或连字符，最多 64 字符）")
    return value


def _text(value: str, kind: str) -> str:
    if not isinstance(value, str) or not value.strip():
        raise ValueError(f"{kind}不能为空")
    return value.strip()


@dataclass(frozen=True)
class AgentSpec:
    name: str
    description: str
    prompt: str
    tools: tuple[str, ...]
    approval_tools: tuple[str, ...]


class ExtensionRegistry:
    """setup() 收到的注册器；不在导入业务模块时修改全局状态。"""

    def __init__(self):
        self.tools: dict[str, tools.ToolSpec] = {}
        self.agents: dict[str, AgentSpec] = {}
        self.modules: list[str] = []
        self._sealed = False

    def _editable(self):
        if self._sealed:
            raise ValueError("扩展已加载；新增或修改扩展后请重启服务")

    def register_tool(self, func: Callable[..., Any], *, group: str | None = None,
                      name: str | None = None, label: str | None = None,
                      summary: str | None = None, requires_approval: bool | None = None) -> str:
        """注册函数；参数覆盖 define_tool 元数据，也可以不使用装饰器。"""
        self._editable()
        spec = tools.tool_spec(func, group=group, name=name, label=label,
                               summary=summary, requires_approval=requires_approval)
        _name(spec.name, "工具")
        if spec.name in _RESERVED_TOOLS:
            raise ValueError(f"工具名与框架内置工具冲突：{spec.name}")
        if spec.name in self.tools:
            raise ValueError(f"工具名冲突：{spec.name}")
        self.tools[spec.name] = spec
        return spec.name

    def _tool_name(self, value: str | Callable[..., Any]) -> str:
        if isinstance(value, str):
            return _name(value, "工具")
        # 支持 register_agent(tools=[函数])；已显式注册的函数复用其注册名。
        matches = [spec.name for spec in self.tools.values() if spec.func is value]
        if len(matches) > 1:
            raise ValueError("函数注册了多个名称，请使用明确的工具名引用")
        return matches[0] if matches else self.register_tool(value)

    def register_agent(self, *, name: str, description: str, prompt: str,
                       tools: Iterable[str | Callable[..., Any]],
                       approval_tools: Iterable[str | Callable[..., Any]] = ()) -> None:
        """注册子 Agent；描述自动进入主 Agent 的委派说明，prompt 是子 Agent 的系统提示词。"""
        self._editable()
        _name(name, "Agent")
        if name in _RESERVED_AGENTS or name in self.agents:
            raise ValueError(f"Agent 名称冲突：{name}")
        if isinstance(tools, str) or isinstance(approval_tools, str):
            raise ValueError(f"Agent {name} 的 tools / approval_tools 必须是列表")
        tool_names = tuple(self._tool_name(tool) for tool in tools)
        approvals = []
        for tool in approval_tools:
            if isinstance(tool, str):
                approvals.append(_name(tool, "工具"))
            else:
                matches = [spec.name for spec in self.tools.values() if spec.func is tool]
                if len(matches) != 1:
                    raise ValueError(f"Agent {name} 的审批工具未注册或名称不明确")
                approvals.append(matches[0])
        self.agents[name] = AgentSpec(name, _text(description, f"Agent {name} 描述"),
                                      _text(prompt, f"Agent {name} 提示词"), tool_names, tuple(approvals))

    def validate(self, extra_approvals: Iterable[str] = ()) -> set[str]:
        """在所有模块声明完成后检查跨模块工具引用及审批，允许前向引用。"""
        for spec in self.tools.values():
            tools.spec_as_tool(spec)  # 提前验证函数 schema，错误不延迟到第一轮对话
        approvals = {spec.name for spec in self.tools.values() if spec.requires_approval}
        for agent in self.agents.values():
            if len(set(agent.tools)) != len(agent.tools):
                raise ValueError(f"Agent {agent.name} 重复引用工具")
            missing = set(agent.tools) - self.tools.keys()
            if missing:
                raise ValueError(f"Agent {agent.name} 引用了未注册工具：{sorted(missing)}")
            outside = set(agent.approval_tools) - set(agent.tools)
            if outside:
                raise ValueError(f"Agent {agent.name} 的审批工具未绑定：{sorted(outside)}")
            approvals.update(agent.approval_tools)
        approvals.update(extra_approvals)
        missing = approvals - self.tools.keys()
        if missing:
            raise ValueError(f"审批配置引用了未注册工具：{sorted(missing)}")
        return approvals

    def subagents(self, approvals: set[str]) -> list[dict[str, Any]]:
        return [{"name": agent.name, "description": agent.description,
                 "system_prompt": agent.prompt,
                 "tools": tools.get_tools_by_name(agent.tools),
                 "interrupt_on": {name: True for name in agent.tools if name in approvals}}
                for agent in self.agents.values()]


_lock = threading.RLock()
_loaded: ExtensionRegistry | None = None


def discover_modules() -> list[str]:
    """只扫描代码包下一级的 extension.py；不扫描用户上传或运行数据。"""
    base = Path(__file__).resolve().parent
    modules = []
    for entry in sorted(base.iterdir()):
        if entry.name.startswith(("_", ".")) or not entry.is_dir():
            continue
        path = entry / "extension.py"
        if not path.exists():
            continue
        if entry.is_symlink() or path.is_symlink() or not entry.name.isidentifier():
            raise ValueError(f"扩展目录非法或为符号链接：{entry.name}")
        modules.append(f"extensions.{entry.name}.extension")
    return modules


def extra_approvals() -> list[str]:
    return [name.strip() for name in config.get_str("HITL_TOOLS", "").split(",") if name.strip()]


def load_extensions() -> ExtensionRegistry:
    global _loaded
    with _lock:
        if _loaded is not None:
            return _loaded
        catalog = ExtensionRegistry()
        for module_name in discover_modules():
            try:
                module = importlib.import_module(module_name)
                setup = getattr(module, "setup", None)
                if not callable(setup):
                    raise ValueError("缺少 setup(registry) 函数")
                setup(catalog)
                catalog.modules.append(module_name)
            except Exception as exc:
                raise ValueError(f"扩展 {module_name} 加载失败：{exc}") from exc
        approvals = catalog.validate(extra_approvals())
        # 主提示词与恢复流程依赖这两个系统工具，缺失时必须在启动阶段报告。
        missing = {"request_plan_approval", "ask_user"} - catalog.tools.keys()
        if missing or not catalog.tools["request_plan_approval"].requires_approval:
            raise ValueError("扩展必须注册 ask_user 和需要审批的 request_plan_approval")
        # 工具清单展示的审批标记与实际执行规则保持一致。
        catalog.tools = {name: replace(spec, requires_approval=name in approvals)
                         for name, spec in catalog.tools.items()}
        tools.install_specs(catalog.tools.values())
        catalog._sealed = True
        _loaded = catalog
        return catalog


def approval_names() -> set[str]:
    catalog = load_extensions()
    # 兼容通过旧 register() 入口增加的工具，并严格校验 HITL_TOOLS。
    names = {spec.name for spec in catalog.tools.values() if spec.requires_approval}
    names.update(name for agent in catalog.agents.values() for name in agent.approval_tools)
    names.update(name for name in tools.get_tool_names() if tools.get_spec(name).requires_approval)
    names.update(extra_approvals())
    missing = names - set(tools.get_tool_names())
    if missing:
        raise ValueError(f"审批配置引用了未注册工具：{sorted(missing)}")
    return names


def reset_for_tests() -> None:
    global _loaded
    with _lock:
        _loaded = None