changes

api_endpoint/websocket_connection_pool.py

@@ -0,0 +1,382 @@
+"""
+WebSocket Connection Pool - Persistent Connection Manager
+
+This module provides a singleton WebSocket connection that persists across
+multiple agent runs. This is a more robust architecture than creating a new
+connection for each run.
+
+Benefits:
+- No timeout issues (connection stays alive indefinitely)
+- No reconnection overhead (connect once)
+- Supports parallel agent runs (multiple sessions share one socket)
+- Proper shutdown handling (SIGTERM/SIGINT)
+- Thread-safe concurrent sends
+"""
+
+import asyncio
+import signal
+import websockets
+from typing import Optional, Dict, Any
+import json
+import atexit
+import sys
+import threading
+from datetime import datetime
+
+
+class WebSocketConnectionPool:
+    """
+    Singleton WebSocket connection manager.
+    
+    Maintains a single persistent connection to the logging server
+    that all agent sessions can use. Handles graceful shutdown.
+    
+    Usage:
+        # Get singleton instance
+        pool = WebSocketConnectionPool()
+        
+        # Connect (idempotent - safe to call multiple times)
+        await pool.connect()
+        
+        # Send events (thread-safe, multiple sessions can call concurrently)
+        await pool.send_event("query", session_id, {...})
+        
+        # Shutdown handled automatically on SIGTERM/SIGINT
+    """
+    
+    _instance: Optional['WebSocketConnectionPool'] = None
+    
+    def __new__(cls):
+        """Ensure only one instance exists (singleton pattern)."""
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+            cls._instance._initialized = False
+        return cls._instance
+    
+    def __init__(self):
+        """Initialize the connection pool (only once)."""
+        if getattr(self, '_initialized', False):
+            return
+            
+        self.websocket: Optional[websockets.WebSocketClientProtocol] = None
+        self.server_url: str = "ws://localhost:8000/ws"
+        self.connected: bool = False
+        # Store reference to loop for signal handlers
+        # Agent code should never close event loops when using persistent connections
+        self.loop: Optional[asyncio.AbstractEventLoop] = None
+        # Locks are created lazily when event loop exists
+        self._send_lock: Optional[asyncio.Lock] = None
+        self._connect_lock: Optional[asyncio.Lock] = None
+        self._locks_loop: Optional[asyncio.AbstractEventLoop] = None  # Track which loop created locks
+        self._init_lock = threading.Lock()  # Thread-safe lock initialization
+        self._shutdown_in_progress = False
+        self._initialized = True
+        
+        # Register shutdown handlers for graceful cleanup
+        # These ensure WebSocket is closed properly on exit
+        signal.signal(signal.SIGTERM, self._signal_handler)
+        signal.signal(signal.SIGINT, self._signal_handler)
+        atexit.register(self._cleanup_sync)
+        
+        print("🔌 WebSocket connection pool initialized")
+    
+    def _ensure_locks(self):
+        """
+        Lazy initialization of asyncio locks with thread safety and loop tracking.
+        
+        Locks must be created when an event loop exists, not at import time.
+        If the event loop changes between runs, locks must be recreated because
+        asyncio.Lock objects are bound to the loop that created them.
+        
+        This is called before any async operation that needs locks.
+        Uses a threading.Lock to prevent race conditions during initialization.
+        """
+        with self._init_lock:  # Thread-safe initialization
+            try:
+                current_loop = asyncio.get_event_loop()
+            except RuntimeError:
+                # No event loop in current thread
+                return
+            
+            # Recreate locks if:
+            # 1. Locks don't exist yet, OR
+            # 2. Event loop has changed (locks are bound to the loop that created them)
+            if self._locks_loop is not current_loop or self._send_lock is None:
+                self._send_lock = asyncio.Lock()
+                self._connect_lock = asyncio.Lock()
+                self._locks_loop = current_loop
+    
+    async def connect(self, server_url: str = "ws://localhost:8000/ws") -> bool:
+        """
+        Connect to WebSocket server.
+        
+        This is idempotent - safe to call multiple times. If already connected,
+        does nothing. If connection failed previously, will retry.
+        
+        Args:
+            server_url: WebSocket server URL (default: ws://localhost:8000/ws)
+            
+        Returns:
+            bool: True if connected successfully, False otherwise
+        """
+        # Ensure locks exist (lazy initialization)
+        self._ensure_locks()
+        
+        async with self._connect_lock:
+            # Always update loop reference to current loop (even if already connected)
+            # This ensures signal handlers and cleanup use the correct loop
+            self.loop = asyncio.get_event_loop()
+            
+            # Already connected - nothing to do
+            if self.connected and self.websocket:
+                return True
+            
+            try:
+                self.server_url = server_url
+                
+                # Establish persistent WebSocket connection
+                # No ping/pong needed since connection stays open indefinitely
+                self.websocket = await websockets.connect(
+                    server_url,
+                    ping_interval=None,  # Disable ping/pong (not needed for persistent connection)
+                    max_size=10 * 1024 * 1024,  # 10MB max message size for large tool results
+                    open_timeout=10,  # 10s timeout for initial connection
+                    close_timeout=5   # 5s timeout for close handshake
+                )
+                
+                self.connected = True
+                
+                print(f"✅ Connected to logging server (persistent): {server_url}")
+                return True
+                
+            except Exception as e:
+                print(f"⚠️ Failed to connect to logging server: {e}")
+                self.connected = False
+                self.websocket = None
+                return False
+    
+    async def send_event(
+        self,
+        event_type: str,
+        session_id: str,
+        data: Dict[str, Any],
+        retry: bool = True
+    ) -> bool:
+        """
+        Send event to logging server (thread-safe).
+        
+        Multiple agent runs can call this concurrently. The send lock ensures
+        only one message is sent at a time (WebSocket protocol requirement).
+        
+        Args:
+            event_type: Type of event (query, api_call, response, tool_call, tool_result, error, complete)
+            session_id: Unique session identifier
+            data: Event-specific data dictionary
+            retry: Whether to retry connection if disconnected (default: True)
+            
+        Returns:
+            bool: True if sent successfully, False otherwise
+        """
+        # Try to connect if not connected (or reconnect if disconnected)
+        if not self.connected or not self.websocket:
+            if retry:
+                await self.connect()
+            if not self.connected:
+                return False  # Give up if connection fails
+        
+        # Ensure locks exist (lazy initialization)
+        self._ensure_locks()
+        
+        # Lock to prevent concurrent sends (WebSocket requires sequential sends)
+        async with self._send_lock:
+            try:
+                # Create standardized message format
+                message = {
+                    "session_id": session_id,
+                    "event_type": event_type,
+                    "data": data,
+                    "timestamp": datetime.now().isoformat()
+                }
+                
+                # Send message as JSON
+                await self.websocket.send(json.dumps(message))
+                
+                # Wait for server acknowledgment (with timeout)
+                # This confirms the server received and processed the event
+                try:
+                    response = await asyncio.wait_for(
+                        self.websocket.recv(),
+                        timeout=2.0  # Increased to 2s for busy servers
+                    )
+                    # Successfully received acknowledgment
+                    return True
+                    
+                except asyncio.TimeoutError:
+                    # No response within timeout - that's OK, message likely sent
+                    # Server might be busy processing
+                    return True
+                    
+            except websockets.exceptions.ConnectionClosed:
+                print(f"⚠️ WebSocket connection closed unexpectedly")
+                self.connected = False
+                
+                # Try to reconnect and resend (one retry)
+                if retry:
+                    print("🔄 Attempting to reconnect...")
+                    if await self.connect():
+                        # Recursively call with retry=False to avoid infinite loop
+                        return await self.send_event(event_type, session_id, data, retry=False)
+                
+                return False
+                
+            except Exception as e:
+                print(f"⚠️ Error sending event: {e}")
+                self.connected = False
+                return False
+    
+    async def disconnect(self):
+        """
+        Gracefully close the WebSocket connection.
+        
+        Called on shutdown (SIGTERM/SIGINT/exit). Ensures proper cleanup.
+        """
+        if self._shutdown_in_progress:
+            return  # Already shutting down
+        
+        self._shutdown_in_progress = True
+        
+        if self.websocket and self.connected:
+            try:
+                await self.websocket.close()
+                self.connected = False
+                print("✅ WebSocket connection pool closed gracefully")
+            except Exception as e:
+                print(f"⚠️ Error closing WebSocket: {e}")
+        
+        self._shutdown_in_progress = False
+    
+    def _signal_handler(self, signum, frame):
+        """
+        Handle SIGTERM/SIGINT signals for graceful shutdown.
+        
+        When user presses Ctrl+C or system sends SIGTERM, this ensures
+        the WebSocket is closed properly before exit.
+        """
+        print(f"\n🛑 Received signal {signum}, closing WebSocket connection pool...")
+        
+        # Check if we have a valid loop and are connected
+        if self.loop and not self.loop.is_closed() and self.connected and not self._shutdown_in_progress:
+            try:
+                # If loop is not running, we can wait for disconnect
+                if not self.loop.is_running():
+                    self.loop.run_until_complete(self.disconnect())
+                else:
+                    # Loop is running, can't wait for task - just mark disconnected
+                    # The disconnect task would be cancelled when we exit anyway
+                    self.connected = False
+                    print("⚠️ Loop is running, marking disconnected without waiting")
+            except Exception as e:
+                print(f"⚠️ Error during signal handler cleanup: {e}")
+        
+        # Exit gracefully
+        sys.exit(0)
+    
+    def _cleanup_sync(self):
+        """
+        Cleanup at exit (atexit handler).
+        
+        This is a fallback in case signal handlers don't fire.
+        Called when Python interpreter shuts down normally.
+        """
+        if self.loop and not self.loop.is_closed() and self.connected and not self._shutdown_in_progress:
+            try:
+                # Try to run disconnect synchronously
+                self.loop.run_until_complete(self.disconnect())
+            except Exception:
+                # Ignore errors during exit cleanup
+                pass
+    
+    def is_connected(self) -> bool:
+        """Check if currently connected to server."""
+        return self.connected and self.websocket is not None
+    
+    def get_stats(self) -> Dict[str, Any]:
+        """Get connection statistics for debugging."""
+        return {
+            "connected": self.connected,
+            "server_url": self.server_url,
+            "shutdown_in_progress": self._shutdown_in_progress,
+            "has_websocket": self.websocket is not None,
+            "has_loop": self.loop is not None
+        }
+
+
+# Global singleton instance
+# Import this in other modules: from websocket_connection_pool import ws_pool
+ws_pool = WebSocketConnectionPool()
+
+
+# Convenience functions for direct usage
+async def connect(server_url: str = "ws://localhost:8000/ws") -> bool:
+    """Connect to logging server (convenience function)."""
+    return await ws_pool.connect(server_url)
+
+
+async def send_event(event_type: str, session_id: str, data: Dict[str, Any]) -> bool:
+    """Send event to logging server (convenience function)."""
+    return await ws_pool.send_event(event_type, session_id, data)
+
+
+async def disconnect():
+    """Disconnect from logging server (convenience function)."""
+    await ws_pool.disconnect()
+
+
+def is_connected() -> bool:
+    """Check if connected to logging server (convenience function)."""
+    return ws_pool.is_connected()
+
+
+# ============================================================================
+# SYNCHRONOUS API FOR AGENT LAYER
+# ============================================================================
+# These functions provide a clean abstraction that hides event loop management
+# from the agent layer. Agent code should ONLY use these functions.
+
+def connect_sync(server_url: str = "ws://localhost:8000/ws") -> bool:
+    """
+    Synchronous connect - handles event loop internally.
+    
+    Agent code should use this instead of directly managing event loops.
+    This ensures the connection pool maintains full control over its lifecycle.
+    """
+    try:
+        loop = asyncio.get_event_loop()
+        if loop.is_closed():
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+    except RuntimeError:
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+    
+    return loop.run_until_complete(ws_pool.connect(server_url))
+
+
+def send_event_sync(event_type: str, session_id: str, data: Dict[str, Any]) -> bool:
+    """
+    Synchronous send event - handles event loop internally.
+    
+    Agent code should use this instead of managing event loops.
+    This ensures the connection pool maintains full control over its lifecycle.
+    """
+    try:
+        loop = asyncio.get_event_loop()
+        if loop.is_closed():
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+    except RuntimeError:
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+    
+    return loop.run_until_complete(ws_pool.send_event(event_type, session_id, data))
+