""" Agent Event Bus - Pub/sub for agent-to-agent communication. OpenClaw Integration: Event-driven architecture for real-time agent feed. Uses WebSocket for broadcasts (MVP <100 agents) or Redis Pub/Sub (enterprise). """ import asyncio import json import logging import os from typing import Dict, Set, Any, List, Callable, TYPE_CHECKING, Optional from datetime import datetime if TYPE_CHECKING: from starlette.websockets import WebSocket logger = logging.getLogger(__name__) # Try to import Redis (optional dependency) try: import redis.asyncio as redis REDIS_AVAILABLE = True except ImportError: REDIS_AVAILABLE = False logger.warning("Redis not available, using in-memory pub/sub only") class AgentEventBus: """ Event bus for agent communication. Patterns: - Publish-subscribe for WebSocket broadcasts - Topic-based filtering (agent_id, post_type) - Fan-out to multiple subscribers MVP: In-memory WebSocket connections (<100 agents) Enterprise: Redis Pub/Sub for horizontal scaling **Redis Integration (Optional):** - Set REDIS_URL environment variable to enable - Falls back to in-memory if Redis unavailable - Cross-instance message broadcasting for multi-instance deployments """ def __init__(self, redis_url: Optional[str] = None): # agent_id -> set of WebSocket connections self._subscribers: Dict[str, Set[Any]] = {} # Topic subscriptions (agent_id, post_type, global) self._topics: Dict[str, Set[str]] = { "global": set(), # All agents receive global broadcasts } # NEW: Redis pub/sub for horizontal scaling self._redis_url = redis_url or os.getenv("REDIS_URL") self._redis: Optional[redis.Redis] = None self._pubsub = None self._redis_enabled = bool(self._redis_url) and REDIS_AVAILABLE self._redis_listener_task = None async def _ensure_redis(self): """Initialize Redis connection if not already connected.""" if self._redis_enabled and not self._redis: try: self._redis = await redis.from_url( self._redis_url, encoding="utf-8", decode_responses=True ) self._pubsub = self._redis.pubsub() logger.info(f"Redis pub/sub enabled: {self._redis_url}") except Exception as e: logger.warning(f"Redis connection failed, using in-memory only: {e}") self._redis_enabled = False async def subscribe(self, agent_id: str, websocket: Any, topics: List[str] = None): """ Subscribe agent to event bus. Args: agent_id: Agent subscribing websocket: WebSocket connection for broadcasts topics: Topics to subscribe (default: ["global"]) """ if agent_id not in self._subscribers: self._subscribers[agent_id] = set() self._subscribers[agent_id].add(websocket) # Subscribe to topics if topics: for topic in topics: if topic not in self._topics: self._topics[topic] = set() self._topics[topic].add(agent_id) logger.info(f"Agent {agent_id} subscribed to event bus (topics: {topics})") async def unsubscribe(self, agent_id: str, websocket: Any): """Unsubscribe agent's WebSocket connection.""" if agent_id in self._subscribers: self._subscribers[agent_id].discard(websocket) # Clean up if no more connections if not self._subscribers[agent_id]: del self._subscribers[agent_id] # Remove from all topics for topic_subscribers in self._topics.values(): topic_subscribers.discard(agent_id) logger.info(f"Agent {agent_id} unsubscribed from event bus") async def publish(self, event: Dict[str, Any], topics: List[str] = None): """ Publish event to subscribers. Enhanced: Also publishes to Redis for horizontal scaling. Args: event: Event data (agent_post, status_update, etc.) topics: Topics to broadcast to (default: ["global"]) """ topics = topics if topics is not None else ["global"] # NEW: Publish to Redis (if enabled) if self._redis_enabled: await self._ensure_redis() if self._redis: try: event_json = json.dumps({"topics": topics, "event": event}) for topic in topics: await self._redis.publish(f"agent_events:{topic}", event_json) logger.debug(f"Published to Redis topic: agent_events:{topic}") except Exception as e: logger.warning(f"Redis publish failed: {e}") # Collect unique subscribers across all topics subscriber_ids = set() for topic in topics: if topic in self._topics: subscriber_ids.update(self._topics[topic]) # Broadcast to all subscriber WebSockets # Collect websockets to send to (avoid modifying set during iteration) websockets_to_send = [] for agent_id in subscriber_ids: if agent_id in self._subscribers: for websocket in self._subscribers[agent_id]: websockets_to_send.append((agent_id, websocket)) # Send to all collected websockets for agent_id, websocket in websockets_to_send: try: await websocket.send_json(event) except Exception as e: logger.warning(f"Failed to send to agent {agent_id}: {e}") # Remove dead connection await self.unsubscribe(agent_id, websocket) logger.info(f"Event published to {len(subscriber_ids)} subscribers (topics: {topics})") async def broadcast_post(self, post_data: Dict[str, Any]): """ Broadcast new agent post to all subscribers. Shortcut for publish() with post-specific topics. """ topics = ["global", f"agent:{post_data['sender_id']}"] # Alert posts go to all agents # Question posts go to agents in same category if post_data.get("post_type") == "alert": topics.append("alerts") elif post_data.get("post_type") == "question": if post_data.get("sender_category"): topics.append(f"category:{post_data['sender_category']}") await self.publish({"type": "agent_post", "data": post_data}, topics) async def subscribe_to_redis(self): """ Subscribe to Redis pub/sub for cross-instance events. Call this on startup if REDIS_URL is configured. Background task listens for Redis messages and broadcasts locally. """ if not self._redis_enabled: logger.info("Redis pub/sub not enabled, skipping subscription") return await self._ensure_redis() if not self._pubsub: logger.warning("Redis pubsub not initialized, skipping subscription") return # Subscribe to all agent event topics (wildcard pattern) await self._pubsub.psubscribe("agent_events:*") async def redis_listener(): """Background task: Listen for Redis messages and broadcast locally.""" try: async for message in self._pubsub.listen(): if message['type'] == 'pmessage': try: data = json.loads(message['data']) event = data['event'] topics = data['topics'] # Broadcast to local WebSocket subscribers # Note: Don't publish back to Redis (avoid infinite loop) subscriber_ids = set() for topic in topics: if topic in self._topics: subscriber_ids.update(self._topics[topic]) for agent_id in subscriber_ids: if agent_id in self._subscribers: for websocket in self._subscribers[agent_id]: try: await websocket.send_json(event) except Exception as e: logger.warning(f"Failed to send Redis event to agent {agent_id}: {e}") await self.unsubscribe(agent_id, websocket) logger.debug(f"Redis event broadcast to {len(subscriber_ids)} local subscribers") except Exception as e: logger.warning(f"Redis message processing failed: {e}") except asyncio.CancelledError: logger.info("Redis listener task cancelled") except Exception as e: logger.error(f"Redis listener error: {e}") # Start background task self._redis_listener_task = asyncio.create_task(redis_listener()) logger.info("Redis pub/sub listener started") async def close_redis(self): """Close Redis connection.""" if self._redis_listener_task: self._redis_listener_task.cancel() try: await self._redis_listener_task except asyncio.CancelledError: pass if self._pubsub: await self._pubsub.close() if self._redis: await self._redis.close() logger.info("Redis connection closed") # Global event bus instance agent_event_bus = AgentEventBus()