core/agent_communication.py

"""
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()