"""Neural Sync - Event Bus System。""" from __future__ import annotations import asyncio import logging from contextlib import asynccontextmanager from typing import Any, Callable, Optional from .events_models import Event, EventType logger = logging.getLogger(__name__) class EventBus: """事件总线核心。""" def __init__(self): self._subscribers: dict[EventType, list[Callable]] = {} self._wildcard_subscribers: list[Callable] = [] self._event_history: list[Event] = [] self._max_history = 1000 self._metrics = { "published": 0, "delivered": 0, "failed": 0, } self._lock = asyncio.Lock() self._running = False @property def metrics(self) -> dict[str, int]: return self._metrics.copy() @property def history(self) -> list[dict[str, Any]]: return [event.to_dict() for event in self._event_history] async def start(self): self._running = True logger.info("EventBus started") async def stop(self): self._running = False logger.info("EventBus stopped") def subscribe( self, event_type: EventType, handler: Callable, filter_func: Optional[Callable[[Event], bool]] = None, ) -> Callable: async def wrapped_handler(event: Event): if filter_func and not filter_func(event): return try: await handler(event) self._metrics["delivered"] += 1 except Exception as exc: logger.error(f"Error in event handler for {event_type.value}: {exc}", exc_info=True) self._metrics["failed"] += 1 wrapped_handler.original = handler # type: ignore[attr-defined] async def _subscribe(): async with self._lock: if event_type not in self._subscribers: self._subscribers[event_type] = [] self._subscribers[event_type].append(wrapped_handler) asyncio.create_task(_subscribe()) def unsubscribe(): asyncio.create_task(self._unsubscribe(event_type, wrapped_handler)) return unsubscribe def subscribe_wildcard( self, handler: Callable, filter_func: Optional[Callable[[Event], bool]] = None, ) -> Callable: async def wrapped_handler(event: Event): if filter_func and not filter_func(event): return try: await handler(event) self._metrics["delivered"] += 1 except Exception as exc: logger.error(f"Error in wildcard event handler: {exc}", exc_info=True) self._metrics["failed"] += 1 wrapped_handler.original = handler # type: ignore[attr-defined] async def _subscribe(): async with self._lock: self._wildcard_subscribers.append(wrapped_handler) asyncio.create_task(_subscribe()) def unsubscribe(): asyncio.create_task(self._unsubscribe_wildcard(wrapped_handler)) return unsubscribe async def _unsubscribe(self, event_type: EventType, handler: Callable): async with self._lock: if event_type in self._subscribers: self._subscribers[event_type] = [ subscribed for subscribed in self._subscribers[event_type] if subscribed != handler ] async def _unsubscribe_wildcard(self, handler: Callable): async with self._lock: self._wildcard_subscribers = [ subscribed for subscribed in self._wildcard_subscribers if subscribed != handler ] async def publish(self, event: Event) -> bool: if not self._running: logger.warning(f"EventBus not running, discarding event: {event.type.value}") return False self._metrics["published"] += 1 self._event_history.append(event) if len(self._event_history) > self._max_history: self._event_history.pop(0) tasks = [] if event.type in self._subscribers: for handler in self._subscribers[event.type]: tasks.append(handler(event)) for handler in self._wildcard_subscribers: tasks.append(handler(event)) if tasks: results = await asyncio.gather(*tasks, return_exceptions=True) failed = sum(1 for result in results if isinstance(result, Exception)) if failed > 0: self._metrics["failed"] += failed logger.debug(f"Published event {event.type.value} to {len(tasks)} subscribers") return True async def publish_and_wait(self, event: Event, timeout: float = 5.0) -> list[Any]: if not self._running: raise RuntimeError("EventBus not running") self._metrics["published"] += 1 self._event_history.append(event) tasks = [] if event.type in self._subscribers: tasks.extend(self._subscribers[event.type]) for handler in self._wildcard_subscribers: tasks.append(handler(event)) if not tasks: return [] try: return await asyncio.wait_for( asyncio.gather(*tasks, return_exceptions=True), timeout=timeout, ) except asyncio.TimeoutError: logger.warning(f"Event {event.type.value} processing timed out after {timeout}s") raise def clear_history(self): self._event_history.clear() def get_subscriber_count(self, event_type: Optional[EventType] = None) -> int: if event_type: return len(self._subscribers.get(event_type, [])) return sum(len(subs) for subs in self._subscribers.values()) + len(self._wildcard_subscribers) _global_event_bus: Optional[EventBus] = None def get_event_bus() -> EventBus: global _global_event_bus if _global_event_bus is None: _global_event_bus = EventBus() return _global_event_bus @asynccontextmanager async def event_bus_context(): bus = get_event_bus() await bus.start() try: yield bus finally: await bus.stop() async def publish_event(event_type: EventType, data: dict[str, Any], **kwargs) -> bool: event = Event(type=event_type, data=data, **kwargs) return await get_event_bus().publish(event) def subscribe_event(event_type: EventType, handler: Callable) -> Callable: return get_event_bus().subscribe(event_type, handler) def subscribe_all_events(handler: Callable) -> Callable: return get_event_bus().subscribe_wildcard(handler) __all__ = [ "Event", "EventBus", "EventType", "event_bus_context", "get_event_bus", "publish_event", "subscribe_all_events", "subscribe_event", ]