vent_agent_2.0/extensions/services.py

"""扩展服务的依赖顺序、共享实例和异步资源管理。"""
from __future__ import annotations

from contextlib import AsyncExitStack
from dataclasses import dataclass
import inspect
import time

from obs import get_logger

logger = get_logger(__name__)
from types import MappingProxyType
from typing import Any, AsyncContextManager, Callable, Mapping


@dataclass(frozen=True)
class ServiceSpec:
    name: str
    factory: Callable[[Mapping[str, Any]], AsyncContextManager[Any]]
    requires: tuple[str, ...] = ()


def ordered_names(dependencies: Mapping[str, set[str]], kind: str) -> list[str]:
    """稳定拓扑排序：无依赖的节点按名称排序，缺失引用和循环立即报错。"""
    pending = {name: set(refs) for name, refs in dependencies.items()}
    for name, refs in pending.items():
        missing = refs - pending.keys()
        if missing:
            raise ValueError(f"{kind} {name} 引用未注册项：{sorted(missing)}")
    result = []
    while pending:
        ready = sorted(name for name, refs in pending.items() if not refs)
        if not ready:
            raise ValueError(f"{kind}存在循环依赖：{sorted(pending)}")
        for name in ready:
            result.append(name)
            del pending[name]
        for refs in pending.values():
            refs.difference_update(ready)
    return result


_active: ServiceContainer | None = None


class ServiceContainer:
    """由 RuntimeEnvironment 管理；工厂必须返回异步上下文管理器。

    工厂只收到其声明的依赖。资源先按依赖顺序创建，再逆序关闭。
    """

    def __init__(self, specs: Mapping[str, ServiceSpec]):
        self._specs = dict(specs)
        self._values: dict[str, Any] = {}
        self._stack = AsyncExitStack()

    async def start(self) -> None:
        global _active
        if _active is not None:
            raise RuntimeError("扩展服务已启动，请复用 RuntimeEnvironment")
        order = ordered_names({name: set(spec.requires) for name, spec in self._specs.items()}, "服务")
        for name in order:
            spec = self._specs[name]
            started = time.perf_counter()
            logger.debug("扩展服务启动", extra={"kv": {"event": "service.starting", "service": name}})
            dependencies = self.view(spec.requires)
            context = spec.factory(dependencies)
            if not hasattr(context, "__aenter__") or not hasattr(context, "__aexit__"):
                if inspect.iscoroutine(context):
                    context.close()
                raise TypeError(f"服务 {name} 的工厂必须返回异步上下文管理器")
            self._values[name] = await self._stack.enter_async_context(context)
            logger.info("扩展服务就绪", extra={"kv": {"event": "service.ready", "service": name,
                        "duration_ms": round((time.perf_counter() - started) * 1000, 1)}})
        _active = self

    def view(self, names) -> Mapping[str, Any]:
        try:
            return MappingProxyType({name: self._values[name] for name in names})
        except KeyError as exc:
            raise RuntimeError(f"扩展服务 {exc.args[0]} 尚未就绪，请在 RuntimeEnvironment 内使用") from exc

    async def close(self) -> None:
        global _active
        try:
            await self._stack.aclose()
        finally:
            self._values.clear()
            if _active is self:
                _active = None


def service_view(names) -> Mapping[str, Any]:
    if not names:
        return MappingProxyType({})
    if _active is None:
        raise RuntimeError("扩展服务尚未启动，请在 RuntimeEnvironment 内使用")
    return _active.view(names)


def get_service(name: str) -> Any:
    """工具在执行时获取服务，不要在 import 或 setup 时调用。"""
    return service_view((name,))[name]