solace-agent-mesh 1.6.1__py3-none-any.whl → 1.13.2__py3-none-any.whl

solace_agent_mesh/gateway/http_sse/sse_manager.py

@@ -5,7 +5,7 @@ Manages Server-Sent Event (SSE) connections for streaming task updates.
 import logging
 import asyncio
 import threading
-from typing import Dict, List, Any
+from typing import Dict, List, Any, Callable, Optional
 import json
 import datetime
 import math
@@ -13,43 +13,28 @@ import math
 from .sse_event_buffer import SSEEventBuffer
 
 log = logging.getLogger(__name__)
+trace_logger = logging.getLogger("sam_trace")
+
 
 class SSEManager:
     """
     Manages active SSE connections and distributes events based on task ID.
     Uses asyncio Queues for buffering events per connection.
+
+    Note: This manager uses a threading.Lock to ensure thread-safety across
+    different event loops (e.g., FastAPI event loop and SAC component event loop).
     """
 
-    def __init__(self, max_queue_size: int, event_buffer: SSEEventBuffer):
+    def __init__(self, max_queue_size: int, event_buffer: SSEEventBuffer, session_factory: Optional[Callable] = None):
         self._connections: Dict[str, List[asyncio.Queue]] = {}
         self._event_buffer = event_buffer
-        self._locks: Dict[asyncio.AbstractEventLoop, asyncio.Lock] = {}
-        self._locks_lock = threading.Lock()
+        # Use a single threading lock for cross-event-loop synchronization
+        self._lock = threading.Lock()
         self.log_identifier = "[SSEManager]"
         self._max_queue_size = max_queue_size
-
-    def _get_lock(self) -> asyncio.Lock:
-        """Get or create a lock for the current event loop."""
-        try:
-            current_loop = asyncio.get_running_loop()
-        except RuntimeError:
-            log.error(
-                "%s _get_lock must be called from within an async context.",
-                self.log_identifier,
-            )
-            raise RuntimeError(
-                "SSEManager methods must be called from within an async context"
-            )
-
-        with self._locks_lock:
-            if current_loop not in self._locks:
-                self._locks[current_loop] = asyncio.Lock()
-                log.debug(
-                    "%s Created new lock for event loop %s",
-                    self.log_identifier,
-                    id(current_loop),
-                )
-            return self._locks[current_loop]
+        self._session_factory = session_factory
+        self._background_task_cache: Dict[str, bool] = {}  # Cache to avoid repeated DB queries
+        self._tasks_with_prior_connection: set = set()  # Track tasks that have had at least one SSE connection
 
     def _sanitize_json(self, obj):
         if isinstance(obj, dict):
@@ -77,27 +62,40 @@ class SSEManager:
         Returns:
             An asyncio.Queue that the SSE endpoint can consume from.
         """
-        lock = self._get_lock()
-        async with lock:
+        connection_queue = asyncio.Queue(maxsize=self._max_queue_size)
+        buffered_events = None
+
+        # Use threading lock for cross-event-loop synchronization
+        with self._lock:
             if task_id not in self._connections:
                 self._connections[task_id] = []
 
-            connection_queue = asyncio.Queue(maxsize=self._max_queue_size)
-
-            # Flush any pending events from the buffer to the new connection
+            # Get buffered events atomically with adding the queue to connections
             buffered_events = self._event_buffer.get_and_remove_buffer(task_id)
-            if buffered_events:
-                for event in buffered_events:
-                    await connection_queue.put(event)
 
+            # Add queue to connections BEFORE releasing lock to ensure
+            # no events are buffered after we've retrieved the buffer
             self._connections[task_id].append(connection_queue)
+
+            # Mark this task as having had at least one connection
+            # This is used to distinguish "no connection yet" from "had connection but disconnected"
+            self._tasks_with_prior_connection.add(task_id)
+
             log.debug(
                 "%s Created SSE connection queue for Task ID: %s. Total queues for task: %d",
                 self.log_identifier,
                 task_id,
                 len(self._connections[task_id]),
             )
-            return connection_queue
+
+        # Put buffered events into queue AFTER releasing the lock
+        # This is safe because the queue is already registered in _connections,
+        # so any new events will go directly to the queue via send_event
+        if buffered_events:
+            for event in buffered_events:
+                await connection_queue.put(event)
+
+        return connection_queue
 
     async def remove_sse_connection(
         self, task_id: str, connection_queue: asyncio.Queue
@@ -109,8 +107,7 @@ class SSEManager:
             task_id: The ID of the task.
             connection_queue: The specific queue instance to remove.
         """
-        lock = self._get_lock()
-        async with lock:
+        with self._lock:
             if task_id in self._connections:
                 try:
                     self._connections[task_id].remove(connection_queue)
@@ -140,6 +137,49 @@ class SSEManager:
                     task_id,
                 )
 
+    def _is_background_task(self, task_id: str) -> bool:
+        """
+        Check if a task is a background task by querying the database.
+        Uses caching to avoid repeated queries.
+        
+        Args:
+            task_id: The ID of the task to check
+            
+        Returns:
+            True if the task is a background task, False otherwise
+        """
+        # Check cache first
+        if task_id in self._background_task_cache:
+            return self._background_task_cache[task_id]
+        
+        # If no session factory, assume not a background task
+        if not self._session_factory:
+            return False
+        
+        try:
+            from .repository.task_repository import TaskRepository
+            
+            db = self._session_factory()
+            try:
+                repo = TaskRepository()
+                task = repo.find_by_id(db, task_id)
+                is_background = task and task.background_execution_enabled
+                
+                # Cache the result
+                self._background_task_cache[task_id] = is_background
+                
+                return is_background
+            finally:
+                db.close()
+        except Exception as e:
+            log.warning(
+                "%s Failed to check if task %s is a background task: %s",
+                self.log_identifier,
+                task_id,
+                e,
+            )
+            return False
+
     async def send_event(
         self, task_id: str, event_data: Dict[str, Any], event_type: str = "message"
     ):
@@ -152,95 +192,130 @@ class SSEManager:
             event_data: The dictionary representing the A2A event (e.g., TaskStatusUpdateEvent).
             event_type: The type of the SSE event (default: "message").
         """
-        lock = self._get_lock()
-        async with lock:
-            queues = self._connections.get(task_id)
+        # Serialize data outside the lock
+        try:
+            serialized_data = json.dumps(
+                self._sanitize_json(event_data), allow_nan=False
+            )
+        except Exception as json_err:
+            log.error(
+                "%s Failed to JSON serialize event data for Task ID %s: %s",
+                self.log_identifier,
+                task_id,
+                json_err,
+            )
+            return
 
-            try:
-                serialized_data = json.dumps(
-                    self._sanitize_json(event_data), allow_nan=False
-                )
-            except Exception as json_err:
-                log.error(
-                    "%s Failed to JSON serialize event data for Task ID %s: %s",
-                    self.log_identifier,
-                    task_id,
-                    json_err,
-                )
-                return
+        sse_payload = {"event": event_type, "data": serialized_data}
 
-            sse_payload = {"event": event_type, "data": serialized_data}
+        # Get queues and decide action under the lock
+        queues_copy = None
+
+        with self._lock:
+            queues = self._connections.get(task_id)
 
             if not queues:
-                log.debug(
-                    "%s No active SSE connections for Task ID: %s. Buffering event.",
-                    self.log_identifier,
-                    task_id,
-                )
-                self._event_buffer.buffer_event(task_id, sse_payload)
-                return
+                # Check if this is a background task (outside lock would be better,
+                # but we need the decision to be atomic with the buffering)
+                is_background_task = self._is_background_task(task_id)
 
-            log.debug(
-                "%s Prepared SSE payload for Task ID %s: %s",
-                self.log_identifier,
-                task_id,
-                sse_payload,
-            )
+                # Check if this task has ever had a connection
+                has_had_connection = task_id in self._tasks_with_prior_connection
 
-            queues_to_remove = []
-            for connection_queue in list(self._connections.get(task_id, [])):
-                try:
-                    await asyncio.wait_for(
-                        connection_queue.put(sse_payload), timeout=0.1
-                    )
+                # Only drop events for background tasks that have HAD a connection before
+                # If no connection has ever been made, we must buffer so the first client gets the events
+                if is_background_task and has_had_connection:
+                    # For background tasks where client disconnected, drop events to prevent buffer overflow
                     log.debug(
-                        "%s Queued event for Task ID: %s to one connection.",
+                        "%s No active SSE connections for background task %s (had prior connection). Dropping event to prevent buffer overflow.",
                         self.log_identifier,
                         task_id,
                     )
-                except asyncio.QueueFull:
-                    log.warning(
-                        "%s SSE connection queue full for Task ID: %s. Event dropped for one connection.",
-                        self.log_identifier,
-                        task_id,
-                    )
-                    queues_to_remove.append(connection_queue)
-                except asyncio.TimeoutError:
-                    log.warning(
-                        "%s Timeout putting event onto SSE queue for Task ID: %s. Event dropped for one connection.",
-                        self.log_identifier,
-                        task_id,
-                    )
-                    queues_to_remove.append(connection_queue)
-                except Exception as e:
-                    log.error(
-                        "%s Error putting event onto queue for Task ID %s: %s",
+                else:
+                    log.debug(
+                        "%s No active SSE connections for Task ID: %s. Buffering event.",
                         self.log_identifier,
                         task_id,
-                        e,
                     )
-                    queues_to_remove.append(connection_queue)
-
-            if queues_to_remove and task_id in self._connections:
-                current_queues = self._connections[task_id]
-                for q in queues_to_remove:
-                    try:
-                        current_queues.remove(q)
-                        log.warning(
-                            "%s Removed potentially broken/full SSE queue for Task ID: %s",
+                    self._event_buffer.buffer_event(task_id, sse_payload)
+                return
+            else:
+                # Make a copy of queues to iterate outside the lock
+                queues_copy = list(queues)
+
+        # Log the payload outside the lock
+        if trace_logger.isEnabledFor(logging.DEBUG):
+            trace_logger.debug(
+                "%s Prepared SSE payload for Task ID %s: %s",
+                self.log_identifier,
+                task_id,
+                sse_payload,
+            )
+        else:
+            log.debug(
+                "%s Prepared SSE payload for Task ID %s",
+                self.log_identifier,
+                task_id,
+            )
+
+        # Send to queues outside the lock (async operations)
+        queues_to_remove = []
+        for connection_queue in queues_copy:
+            try:
+                await asyncio.wait_for(
+                    connection_queue.put(sse_payload), timeout=0.1
+                )
+                log.debug(
+                    "%s Queued event for Task ID: %s to one connection.",
+                    self.log_identifier,
+                    task_id,
+                )
+            except asyncio.QueueFull:
+                log.warning(
+                    "%s SSE connection queue full for Task ID: %s. Event dropped for one connection.",
+                    self.log_identifier,
+                    task_id,
+                )
+                queues_to_remove.append(connection_queue)
+            except asyncio.TimeoutError:
+                log.warning(
+                    "%s Timeout putting event onto SSE queue for Task ID: %s. Event dropped for one connection.",
+                    self.log_identifier,
+                    task_id,
+                )
+                queues_to_remove.append(connection_queue)
+            except Exception as e:
+                log.error(
+                    "%s Error putting event onto queue for Task ID %s: %s",
+                    self.log_identifier,
+                    task_id,
+                    e,
+                )
+                queues_to_remove.append(connection_queue)
+
+        # Remove broken queues under the lock
+        if queues_to_remove:
+            with self._lock:
+                if task_id in self._connections:
+                    current_queues = self._connections[task_id]
+                    for q in queues_to_remove:
+                        try:
+                            current_queues.remove(q)
+                            log.warning(
+                                "%s Removed potentially broken/full SSE queue for Task ID: %s",
+                                self.log_identifier,
+                                task_id,
+                            )
+                        except ValueError:
+                            pass
+
+                    if not current_queues:
+                        del self._connections[task_id]
+                        log.debug(
+                            "%s Removed Task ID entry: %s after cleaning queues.",
                             self.log_identifier,
                             task_id,
                         )
-                    except ValueError:
-                        pass
-
-                if not current_queues:
-                    del self._connections[task_id]
-                    log.debug(
-                        "%s Removed Task ID entry: %s after cleaning queues.",
-                        self.log_identifier,
-                        task_id,
-                    )
 
     async def close_connection(self, task_id: str, connection_queue: asyncio.Queue):
         """
@@ -276,54 +351,58 @@ class SSEManager:
         finally:
             await self.remove_sse_connection(task_id, connection_queue)
 
+    async def drain_buffer_for_background_task(self, task_id: str):
+        """
+        Drains the event buffer for a background task when a client disconnects.
+        This prevents buffer overflow warnings when background tasks continue
+        generating events with no active consumers.
+        
+        Args:
+            task_id: The ID of the background task
+        """
+        log.info(
+            "%s Draining event buffer for background task: %s",
+            self.log_identifier,
+            task_id,
+        )
+        
+        # Remove any buffered events to prevent overflow
+        buffered_events = self._event_buffer.get_and_remove_buffer(task_id)
+        if buffered_events:
+            log.info(
+                "%s Drained %d buffered events for background task: %s",
+                self.log_identifier,
+                len(buffered_events),
+                task_id,
+            )
+        else:
+            log.debug(
+                "%s No buffered events to drain for background task: %s",
+                self.log_identifier,
+                task_id,
+            )
+
     async def close_all_for_task(self, task_id: str):
         """
         Closes all SSE connections associated with a specific task.
         If a connection existed, it also cleans up the event buffer.
         If no connection ever existed, the buffer is left for a late-connecting client.
         """
-        lock = self._get_lock()
-        async with lock:
+        queues_to_close = None
+        should_remove_buffer = False
+
+        with self._lock:
             if task_id in self._connections:
                 # This is the "normal" case: a client is or was connected.
                 # It's safe to clean up everything.
                 queues_to_close = self._connections.pop(task_id)
+                should_remove_buffer = True
                 log.debug(
                     "%s Closing %d SSE connections for Task ID: %s and cleaning up buffer.",
                     self.log_identifier,
                     len(queues_to_close),
                     task_id,
                 )
-                for q in queues_to_close:
-                    try:
-                        await asyncio.wait_for(q.put(None), timeout=0.1)
-                    except asyncio.QueueFull:
-                        log.warning(
-                            "%s Could not put None (close signal) on full queue during close_all for Task ID: %s.",
-                            self.log_identifier,
-                            task_id,
-                        )
-                    except asyncio.TimeoutError:
-                        log.warning(
-                            "%s Timeout putting None (close signal) on queue during close_all for Task ID: %s.",
-                            self.log_identifier,
-                            task_id,
-                        )
-                    except Exception as e:
-                        log.error(
-                            "%s Error putting None (close signal) on queue during close_all for Task ID %s: %s",
-                            self.log_identifier,
-                            task_id,
-                            e,
-                        )
-
-                # Since a connection existed, the buffer is no longer needed.
-                self._event_buffer.remove_buffer(task_id)
-                log.debug(
-                    "%s Removed Task ID entry: %s and signaled queues to close.",
-                    self.log_identifier,
-                    task_id,
-                )
             else:
                 # This is the "race condition" case: no client has connected yet.
                 # We MUST leave the buffer intact for the late-connecting client.
@@ -333,39 +412,80 @@ class SSEManager:
                     task_id,
                 )
 
-    def cleanup_old_locks(self):
-        """Remove locks for closed event loops to prevent memory leaks."""
-        with self._locks_lock:
-            closed_loops = [loop for loop in self._locks if loop.is_closed()]
-            for loop in closed_loops:
-                del self._locks[loop]
+        # Close queues outside the lock (async operations)
+        if queues_to_close:
+            for q in queues_to_close:
+                try:
+                    await asyncio.wait_for(q.put(None), timeout=0.1)
+                except asyncio.QueueFull:
+                    log.warning(
+                        "%s Could not put None (close signal) on full queue during close_all for Task ID: %s.",
+                        self.log_identifier,
+                        task_id,
+                    )
+                except asyncio.TimeoutError:
+                    log.warning(
+                        "%s Timeout putting None (close signal) on queue during close_all for Task ID: %s.",
+                        self.log_identifier,
+                        task_id,
+                    )
+                except Exception as e:
+                    log.error(
+                        "%s Error putting None (close signal) on queue during close_all for Task ID %s: %s",
+                        self.log_identifier,
+                        task_id,
+                        e,
+                    )
+
+            # Since a connection existed, the buffer is no longer needed.
+            # This is safe to do without lock since we already removed the task from _connections
+            if should_remove_buffer:
+                self._event_buffer.remove_buffer(task_id)
+
+                # Clean up the connection tracking
+                with self._lock:
+                    self._tasks_with_prior_connection.discard(task_id)
+
                 log.debug(
-                    "%s Cleaned up lock for closed event loop %s",
+                    "%s Removed Task ID entry: %s and signaled queues to close.",
                     self.log_identifier,
-                    id(loop),
+                    task_id,
                 )
 
+    def cleanup_old_locks(self):
+        """Legacy method - no longer needed with single threading lock.
+        Kept for API compatibility but does nothing."""
+        pass
+
     async def close_all(self):
         """Closes all active SSE connections managed by this instance."""
         self.cleanup_old_locks()
-        lock = self._get_lock()
-        async with lock:
+
+        # Collect all queues to close under the lock
+        all_queues_to_close = []
+        all_task_ids = []
+
+        with self._lock:
             log.debug("%s Closing all active SSE connections...", self.log_identifier)
             all_task_ids = list(self._connections.keys())
-            closed_count = 0
             for task_id in all_task_ids:
                 if task_id in self._connections:
                     queues = self._connections.pop(task_id)
-                    closed_count += len(queues)
-                    for q in queues:
-                        try:
-                            await asyncio.wait_for(q.put(None), timeout=0.1)
-                        except Exception:
-                            pass
-            log.debug(
-                "%s Closed %d connections for tasks: %s",
-                self.log_identifier,
-                closed_count,
-                all_task_ids,
-            )
+                    all_queues_to_close.extend(queues)
             self._connections.clear()
+            self._tasks_with_prior_connection.clear()
+
+        # Close queues outside the lock (async operations)
+        closed_count = len(all_queues_to_close)
+        for q in all_queues_to_close:
+            try:
+                await asyncio.wait_for(q.put(None), timeout=0.1)
+            except Exception:
+                pass
+
+        log.debug(
+            "%s Closed %d connections for tasks: %s",
+            self.log_identifier,
+            closed_count,
+            all_task_ids,
+        )