File size: 9,944 Bytes
92c4ae6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 | """
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()
|