from __future__ import annotations import asyncio from collections import deque from collections.abc import Iterable from pydantic import BaseModel from .event import EventEnvelope, InterruptData, PluginErrorData from .plugin import EventHandler, Plugin class PluginManager: def __init__(self, *, event_limit: int = 1_000) -> None: if event_limit < 0: raise ValueError("event limit must not be negative") self._plugins: dict[str, Plugin] = {} self._subscribers: dict[str, list[tuple[Plugin, EventHandler]]] = {} self._contracts: dict[str, type[BaseModel]] = { "error": PluginErrorData, "interrupt": InterruptData, } self._events: deque[EventEnvelope[BaseModel]] = deque(maxlen=event_limit) self._dispatch_tasks: set[asyncio.Task[None]] = set() def register_event(self, name: str, contract: type[BaseModel]) -> None: if not name or not issubclass(contract, BaseModel): raise ValueError("event registrations require a name and Pydantic contract") self._contracts[name] = contract async def enable_plugin(self, plugin: Plugin) -> None: if plugin.name in self._plugins and self._plugins[plugin.name] is not plugin: raise ValueError(f"plugin already registered: {plugin.name}") if plugin.enabled: return plugin.manager = self if not plugin._initialized: await plugin.initialize() plugin._initialized = True plugin.enabled = True self._plugins[plugin.name] = plugin for name, handler in plugin.iter_event_handlers().items(): self._subscribers.setdefault(name, []).append((plugin, handler)) plugin._start_worker() async def disable_plugin(self, name: str) -> None: plugin = self._plugins.get(name) if plugin is None or not plugin.enabled: return plugin.enabled = False self._subscribers = { event: [(item, handler) for item, handler in subscribers if item is not plugin] for event, subscribers in self._subscribers.items() } await plugin._stop_worker() async def close(self) -> None: for name in tuple(self._plugins): await self.disable_plugin(name) def get_plugins(self) -> tuple[Plugin, ...]: return tuple(self._plugins.values()) def get_events(self) -> tuple[EventEnvelope[BaseModel], ...]: return tuple(self._events) async def emit( self, name: str, data: BaseModel, *, source: str = "space", target: str | None = None, capabilities: Iterable[str] = (), only_once: bool = False, ) -> EventEnvelope[BaseModel]: contract = self._contracts.get(name) if contract is not None and not isinstance(data, contract): raise TypeError(f"{name} requires {contract.__name__}") required = tuple(capabilities) envelope = EventEnvelope[BaseModel]( name=name, data=data, source=source, target=target, capabilities=list(required) or None, only_once=only_once, ) self._events.append(envelope) recipients = [ (plugin, handler) for plugin, handler in self._subscribers.get(name, []) if plugin.enabled and (plugin.receive_self_events or plugin.name != source.split("/", maxsplit=1)[0]) and (target is None or plugin.name == target) ] if only_once: recipients = recipients[:1] for plugin, handler in recipients: task = asyncio.create_task(plugin._enqueue(envelope, handler)) self._dispatch_tasks.add(task) task.add_done_callback(self._dispatch_tasks.discard) return envelope async def wait_for_idle(self, *, timeout: float = 5) -> None: async with asyncio.timeout(timeout): while self._dispatch_tasks or any(not plugin.queue.empty() for plugin in self._plugins.values()): if self._dispatch_tasks: await asyncio.gather(*tuple(self._dispatch_tasks)) await asyncio.gather(*(plugin.queue.join() for plugin in self._plugins.values() if plugin.enabled)) await asyncio.sleep(0) async def report_plugin_error(self, plugin: Plugin, envelope: EventEnvelope[BaseModel], error: Exception) -> None: if envelope.name != "error": await self.emit( "error", PluginErrorData( plugin=plugin.name, event_name=envelope.name, error_type=type(error).__name__, message=str(error), ), source=plugin.name, )