vent_agent_2.0/extensions/registry.py

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

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

import config
from tools import registry as tools
from extensions.services import ServiceSpec, ordered_names, service_view

_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", "main"}
_RESERVED_MIDDLEWARE = {
    "DynamicTimeMiddleware", "SummarizationMiddleware", "TodoListMiddleware",
    "FilesystemMiddleware", "SubAgentMiddleware", "PatchToolCallsMiddleware",
    "SkillsMiddleware", "MemoryMiddleware", "HumanInTheLoopMiddleware",
    "AnthropicPromptCachingMiddleware",
}


def _names(values: Iterable[str], kind: str) -> tuple[str, ...]:
    if isinstance(values, str):
        raise ValueError(f"{kind}必须是名称列表")
    result = tuple(_name(value, kind) for value in values)
    if len(set(result)) != len(result):
        raise ValueError(f"{kind}不能重复")
    return result


@dataclass(frozen=True)
class MiddlewareSpec:
    name: str
    factory: Callable[..., Any]
    targets: tuple[str, ...]
    requires: tuple[str, ...]
    before: tuple[str, ...]
    after: tuple[str, ...]


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()


def _call_sync(hook: Callable[..., Any], name: str, *args) -> Any:
    result = hook(*args)
    if inspect.isawaitable(result):
        if inspect.iscoroutine(result):
            result.close()
        raise ValueError(f"{name} 必须同步执行；异步初始化请放入服务工厂")
    return result


@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.disabled_modules: list[str] = []
        self.services: dict[str, ServiceSpec] = {}
        self.middleware: dict[str, MiddlewareSpec] = {}
        self.tool_services: dict[str, tuple[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,
                      requires: Iterable[str] = ()) -> 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.tool_services[spec.name] = _names(requires, "工具服务依赖")
        self.tools[spec.name] = spec
        return spec.name

    def register_service(self, name: str, factory: Callable[..., Any], *,
                         requires: Iterable[str] = ()) -> None:
        """声明服务。factory(dependencies) 返回 async context manager。"""
        self._editable()
        _name(name, "服务")
        if name in self.services:
            raise ValueError(f"服务名称冲突：{name}")
        if not callable(factory) or inspect.iscoroutinefunction(factory):
            raise ValueError(f"服务 {name} 的 factory 必须返回异步上下文管理器，请使用 asynccontextmanager")
        self.services[name] = ServiceSpec(name, factory, _names(requires, "服务依赖"))

    def register_middleware(self, name: str, factory: Callable[..., Any], *,
                            targets: Iterable[str] = ("main",), requires: Iterable[str] = (),
                            before: Iterable[str] = (), after: Iterable[str] = ()) -> None:
        """声明无副作用的同步工厂；每次编译为每个目标生成独立中间件。"""
        self._editable()
        _name(name, "中间件")
        if name in self.middleware or name in _RESERVED_MIDDLEWARE:
            raise ValueError(f"中间件名称冲突：{name}")
        if not callable(factory) or inspect.iscoroutinefunction(factory):
            raise ValueError(f"中间件 {name} 的 factory 必须是同步函数")
        targets = _names(targets, "中间件目标")
        if not targets:
            raise ValueError(f"中间件 {name} 至少需要一个目标")
        self.middleware[name] = MiddlewareSpec(name, factory, targets,
            _names(requires, "中间件服务依赖"), _names(before, "before"), _names(after, "after"))

    def middleware_order(self) -> list[str]:
        deps = {name: set(spec.after) for name, spec in self.middleware.items()}
        for name, spec in self.middleware.items():
            for ref in (*spec.before, *spec.after):
                if ref not in deps:
                    raise ValueError(f"中间件 {name} 引用未注册项：{ref}")
                if not set(spec.targets).intersection(self.middleware[ref].targets):
                    raise ValueError(f"中间件 {name} 与 {ref} 没有共同目标，不能指定顺序")
            for ref in spec.before:
                deps[ref].add(name)
        return ordered_names(deps, "中间件")

    def build_middleware(self, target: str) -> list[Any]:
        from langchain.agents.middleware.types import AgentMiddleware
        result = []
        for name in self.middleware_order():
            spec = self.middleware[name]
            if target not in spec.targets:
                continue
            instance = spec.factory(service_view(spec.requires))
            if not isinstance(instance, AgentMiddleware) or instance.name != name:
                if inspect.iscoroutine(instance):
                    instance.close()
                raise ValueError(f"中间件 {name} 的工厂必须返回同名 AgentMiddleware 实例")
            result.append(instance)
        return result

    def validate_components(self) -> None:
        ordered_names({name: set(spec.requires) for name, spec in self.services.items()}, "服务")
        for name, requires in self.tool_services.items():
            self._check_services(f"工具 {name}", requires)
        for name, spec in self.middleware.items():
            self._check_services(f"中间件 {name}", spec.requires)
            missing = set(spec.targets) - {"main", *self.agents}
            if missing:
                raise ValueError(f"中间件 {name} 引用了未注册目标：{sorted(missing)}")
        self.middleware_order()

    def _check_services(self, owner: str, names: Iterable[str]) -> None:
        missing = set(names) - self.services.keys()
        if missing:
            raise ValueError(f"{owner} 引用了未注册服务：{sorted(missing)}")

    def check_middleware_factories(self) -> None:
        # 在服务就绪后验证工厂；编译时重新创建，不能在工厂里开连接或启动任务。
        for target in ("main", *sorted(self.agents)):
            self.build_middleware(target)

    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]:
        """在所有模块声明完成后检查跨模块工具引用及审批，允许前向引用。"""
        self.validate_components()
        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)
                enabled = getattr(module, "enabled", None)
                if enabled is not None:
                    if not callable(enabled):
                        raise ValueError("enabled 必须是返回布尔值的函数")
                    active = _call_sync(enabled, "enabled")
                    if type(active) is not bool:
                        raise ValueError("enabled() 必须返回布尔值")
                    if not active:
                        catalog.disabled_modules.append(module_name)
                        continue
                setup = getattr(module, "setup", None)
                if not callable(setup):
                    raise ValueError("缺少 setup(registry) 函数")
                validate_config = getattr(module, "validate_config", None)
                if validate_config is not None:
                    if not callable(validate_config) or inspect.iscoroutinefunction(validate_config):
                        raise ValueError("validate_config 必须是同步校验函数")
                    _call_sync(validate_config, "validate_config")
                _call_sync(setup, "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