solace-agent-mesh 1.3.3__py3-none-any.whl → 1.4.0__py3-none-any.whl

solace_agent_mesh/gateway/http_sse/component.py

@@ -17,12 +17,14 @@ from fastapi import Request as FastAPIRequest
 from solace_ai_connector.common.log import log
 from solace_ai_connector.components.inputs_outputs.broker_input import BrokerInput
 from solace_ai_connector.flow.app import App as SACApp
+from solace_ai_connector.common.event import Event, EventType
 
 from ...common.agent_registry import AgentRegistry
 from ...core_a2a.service import CoreA2AService
 from ...gateway.base.component import BaseGatewayComponent
 from ...gateway.http_sse.session_manager import SessionManager
 from ...gateway.http_sse.sse_manager import SSEManager
+from .sse_event_buffer import SSEEventBuffer
 from .components import VisualizationForwarderComponent
 
 try:
@@ -118,8 +120,23 @@ class WebUIBackendComponent(BaseGatewayComponent):
             raise ValueError(f"Configuration retrieval error: {e}") from e
 
         sse_max_queue_size = self.get_config("sse_max_queue_size", 200)
+        sse_buffer_max_age_seconds = self.get_config("sse_buffer_max_age_seconds", 600)
 
-        self.sse_manager = SSEManager(max_queue_size=sse_max_queue_size)
+        self.sse_event_buffer = SSEEventBuffer(
+            max_queue_size=sse_max_queue_size,
+            max_age_seconds=sse_buffer_max_age_seconds,
+        )
+        self.sse_manager = SSEManager(
+            max_queue_size=sse_max_queue_size, event_buffer=self.sse_event_buffer
+        )
+
+        self._sse_cleanup_timer_id = f"sse_cleanup_{self.gateway_id}"
+        cleanup_interval_sec = self.get_config("sse_buffer_cleanup_interval_seconds", 300)
+        self.add_timer(
+            delay_ms=cleanup_interval_sec * 1000,
+            timer_id=self._sse_cleanup_timer_id,
+            interval_ms=cleanup_interval_sec * 1000,
+        )
 
         session_config = self._resolve_session_config()
         if session_config.get("type") == "sql":
@@ -167,6 +184,15 @@ class WebUIBackendComponent(BaseGatewayComponent):
 
         log.info("%s Web UI Backend Component initialized.", self.log_identifier)
 
+    def process_event(self, event: Event):
+        if event.event_type == EventType.TIMER:
+            if event.data.get("timer_id") == self._sse_cleanup_timer_id:
+                log.debug("%s SSE buffer cleanup timer triggered.", self.log_identifier)
+                self.sse_event_buffer.cleanup_stale_buffers()
+                return
+
+        super().process_event(event)
+
     def _get_visualization_lock(self) -> asyncio.Lock:
         """Get or create a visualization lock for the current event loop."""
         try:
@@ -1025,6 +1051,7 @@ class WebUIBackendComponent(BaseGatewayComponent):
     def cleanup(self):
         """Gracefully shuts down the component and the FastAPI server."""
         log.info("%s Cleaning up Web UI Backend Component...", self.log_identifier)
+        self.cancel_timer(self._sse_cleanup_timer_id)
         log.info("%s Cleaning up visualization resources...", self.log_identifier)
         if self._visualization_message_queue:
             self._visualization_message_queue.put(None)
@@ -1056,6 +1083,35 @@ class WebUIBackendComponent(BaseGatewayComponent):
         self._cleanup_visualization_locks()
         log.info("%s Visualization resources cleaned up.", self.log_identifier)
 
+        super().cleanup()
+
+        if self.fastapi_thread and self.fastapi_thread.is_alive():
+            log.info(
+                "%s Waiting for FastAPI server thread to exit...", self.log_identifier
+            )
+            self.fastapi_thread.join(timeout=10)
+            if self.fastapi_thread.is_alive():
+                log.warning(
+                    "%s FastAPI server thread did not exit gracefully.",
+                    self.log_identifier,
+                )
+
+        if self.sse_manager:
+            log.info(
+                "%s Closing active SSE connections (best effort)...",
+                self.log_identifier,
+            )
+            try:
+                asyncio.run(self.sse_manager.close_all())
+            except Exception as sse_close_err:
+                log.error(
+                    "%s Error closing SSE connections during cleanup: %s",
+                    self.log_identifier,
+                    sse_close_err,
+                )
+
+        log.info("%s Web UI Backend Component cleanup finished.", self.log_identifier)
+
     def _infer_visualization_event_details(
         self, topic: str, payload: dict[str, Any]
     ) -> dict[str, Any]:
@@ -1302,35 +1358,6 @@ class WebUIBackendComponent(BaseGatewayComponent):
             )
         return agents
 
-        super().cleanup()
-
-        if self.fastapi_thread and self.fastapi_thread.is_alive():
-            log.info(
-                "%s Waiting for FastAPI server thread to exit...", self.log_identifier
-            )
-            self.fastapi_thread.join(timeout=10)
-            if self.fastapi_thread.is_alive():
-                log.warning(
-                    "%s FastAPI server thread did not exit gracefully.",
-                    self.log_identifier,
-                )
-
-        if self.sse_manager:
-            log.info(
-                "%s Closing active SSE connections (best effort)...",
-                self.log_identifier,
-            )
-            try:
-                asyncio.run(self.sse_manager.close_all())
-            except Exception as sse_close_err:
-                log.error(
-                    "%s Error closing SSE connections during cleanup: %s",
-                    self.log_identifier,
-                    sse_close_err,
-                )
-
-        log.info("%s Web UI Backend Component cleanup finished.", self.log_identifier)
-
     def get_agent_registry(self) -> AgentRegistry:
         return self.agent_registry
 

solace_agent_mesh/gateway/http_sse/sse_event_buffer.py

@@ -0,0 +1,87 @@
+"""
+A thread-safe buffer for holding early SSE events before a client connects.
+"""
+
+import datetime
+import threading
+from typing import Any, Dict, List, Optional, Tuple
+
+from solace_ai_connector.common.log import log
+
+
+class SSEEventBuffer:
+    """Manages buffering and cleanup of SSE events for tasks without active listeners."""
+
+    def __init__(self, max_queue_size: int, max_age_seconds: int):
+        self._pending_events: Dict[
+            str, Tuple[datetime.datetime, List[Dict[str, Any]]]
+        ] = {}
+        self._lock = threading.Lock()
+        self._max_queue_size = max_queue_size
+        self._max_age_seconds = max_age_seconds
+        self.log_identifier = "[SSEEventBuffer]"
+        log.debug(
+            "%s Initialized with max_age:%ds, max_size:%d",
+            self.log_identifier,
+            self._max_age_seconds,
+            self._max_queue_size,
+        )
+
+    def buffer_event(self, task_id: str, event: Dict[str, Any]):
+        """Buffers an event for a given task ID."""
+        with self._lock:
+            if task_id not in self._pending_events:
+                self._pending_events[task_id] = (
+                    datetime.datetime.now(datetime.timezone.utc),
+                    [],
+                )
+
+            if len(self._pending_events[task_id][1]) < self._max_queue_size:
+                self._pending_events[task_id][1].append(event)
+            else:
+                log.warning(
+                    "%s Buffer full for Task ID: %s. Event dropped.",
+                    self.log_identifier,
+                    task_id,
+                )
+
+    def get_and_remove_buffer(self, task_id: str) -> Optional[List[Dict[str, Any]]]:
+        """Atomically retrieves and removes the event buffer for a task."""
+        with self._lock:
+            buffer_tuple = self._pending_events.pop(task_id, None)
+            if buffer_tuple:
+                log.debug(
+                    "%s Flushing %d events for Task ID: %s",
+                    self.log_identifier,
+                    len(buffer_tuple[1]),
+                    task_id,
+                )
+                return buffer_tuple[1]
+            return None
+
+    def remove_buffer(self, task_id: str):
+        """Explicitly removes a buffer for a task, e.g., on finalization."""
+        with self._lock:
+            if self._pending_events.pop(task_id, None):
+                log.debug(
+                    "%s Removed buffer for task %s.", self.log_identifier, task_id
+                )
+
+    def cleanup_stale_buffers(self):
+        """Removes all pending event buffers older than the max age."""
+        with self._lock:
+            now = datetime.datetime.now(datetime.timezone.utc)
+            stale_tasks = [
+                task_id
+                for task_id, (timestamp, _) in self._pending_events.items()
+                if (now - timestamp).total_seconds() > self._max_age_seconds
+            ]
+
+            if stale_tasks:
+                log.debug(
+                    "%s Cleaning up %d stale event buffers.",
+                    self.log_identifier,
+                    len(stale_tasks),
+                )
+                for task_id in stale_tasks:
+                    del self._pending_events[task_id]

solace_agent_mesh/gateway/http_sse/sse_manager.py

@@ -11,6 +11,8 @@ import math
 
 from solace_ai_connector.common.log import log
 
+from .sse_event_buffer import SSEEventBuffer
+
 
 class SSEManager:
     """
@@ -18,8 +20,9 @@ class SSEManager:
     Uses asyncio Queues for buffering events per connection.
     """
 
-    def __init__(self, max_queue_size: int = 200):
+    def __init__(self, max_queue_size: int, event_buffer: SSEEventBuffer):
         self._connections: Dict[str, List[asyncio.Queue]] = {}
+        self._event_buffer = event_buffer
         self._locks: Dict[asyncio.AbstractEventLoop, asyncio.Lock] = {}
         self._locks_lock = threading.Lock()
         self.log_identifier = "[SSEManager]"
@@ -64,7 +67,6 @@ class SSEManager:
         else:
             return str(obj)
 
-
     async def create_sse_connection(self, task_id: str) -> asyncio.Queue:
         """
         Creates a new queue for an SSE connection subscribing to a task.
@@ -81,8 +83,15 @@ class SSEManager:
                 self._connections[task_id] = []
 
             connection_queue = asyncio.Queue(maxsize=self._max_queue_size)
+
+            # Flush any pending events from the buffer to the new connection
+            buffered_events = self._event_buffer.get_and_remove_buffer(task_id)
+            if buffered_events:
+                for event in buffered_events:
+                    await connection_queue.put(event)
+
             self._connections[task_id].append(connection_queue)
-            log.info(
+            log.debug(
                 "%s Created SSE connection queue for Task ID: %s. Total queues for task: %d",
                 self.log_identifier,
                 task_id,
@@ -105,7 +114,7 @@ class SSEManager:
             if task_id in self._connections:
                 try:
                     self._connections[task_id].remove(connection_queue)
-                    log.info(
+                    log.debug(
                         "%s Removed SSE connection queue for Task ID: %s. Remaining queues: %d",
                         self.log_identifier,
                         task_id,
@@ -113,7 +122,7 @@ class SSEManager:
                     )
                     if not self._connections[task_id]:
                         del self._connections[task_id]
-                        log.info(
+                        log.debug(
                             "%s Removed Task ID entry: %s as no connections remain.",
                             self.log_identifier,
                             task_id,
@@ -145,19 +154,11 @@ class SSEManager:
         """
         lock = self._get_lock()
         async with lock:
-            if task_id not in self._connections:
-                log.debug(
-                    "%s No active SSE connections for Task ID: %s. Event not sent.",
-                    self.log_identifier,
-                    task_id,
-                )
-                return
+            queues = self._connections.get(task_id)
 
-            queues_to_remove = []
             try:
                 serialized_data = json.dumps(
-                    self._sanitize_json(event_data),
-                    allow_nan=False
+                    self._sanitize_json(event_data), allow_nan=False
                 )
             except Exception as json_err:
                 log.error(
@@ -169,6 +170,16 @@ class SSEManager:
                 return
 
             sse_payload = {"event": event_type, "data": serialized_data}
+
+            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
+
             log.debug(
                 "%s Prepared SSE payload for Task ID %s: %s",
                 self.log_identifier,
@@ -176,6 +187,7 @@ class SSEManager:
                 sse_payload,
             )
 
+            queues_to_remove = []
             for connection_queue in list(self._connections.get(task_id, [])):
                 try:
                     await asyncio.wait_for(
@@ -224,7 +236,7 @@ class SSEManager:
 
                 if not current_queues:
                     del self._connections[task_id]
-                    log.info(
+                    log.debug(
                         "%s Removed Task ID entry: %s after cleaning queues.",
                         self.log_identifier,
                         task_id,
@@ -235,7 +247,7 @@ class SSEManager:
         Signals a specific SSE connection queue to close by putting None.
         Also removes the queue from the manager.
         """
-        log.info(
+        log.debug(
             "%s Closing specific SSE connection queue for Task ID: %s",
             self.log_identifier,
             task_id,
@@ -267,13 +279,17 @@ class SSEManager:
     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:
             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)
-                log.info(
-                    "%s Closing %d SSE connections for Task ID: %s",
+                log.debug(
+                    "%s Closing %d SSE connections for Task ID: %s and cleaning up buffer.",
                     self.log_identifier,
                     len(queues_to_close),
                     task_id,
@@ -300,14 +316,19 @@ class SSEManager:
                             task_id,
                             e,
                         )
-                log.info(
+
+                # 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.
                 log.debug(
-                    "%s No connections found to close for Task ID: %s",
+                    "%s No active connections found for Task ID: %s. Leaving event buffer intact.",
                     self.log_identifier,
                     task_id,
                 )
@@ -329,7 +350,7 @@ class SSEManager:
         self.cleanup_old_locks()
         lock = self._get_lock()
         async with lock:
-            log.info("%s Closing all active SSE connections...", self.log_identifier)
+            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:
@@ -341,7 +362,7 @@ class SSEManager:
                             await asyncio.wait_for(q.put(None), timeout=0.1)
                         except Exception:
                             pass
-            log.info(
+            log.debug(
                 "%s Closed %d connections for tasks: %s",
                 self.log_identifier,
                 closed_count,

{solace_agent_mesh-1.3.3.dist-info → solace_agent_mesh-1.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: solace-agent-mesh
-Version: 1.3.3
+Version: 1.4.0
 Summary: Solace Agent Mesh is an open-source framework for building event-driven, multi-agent AI systems where specialized agents collaborate on complex tasks.
 Project-URL: Homepage, https://github.com/SolaceLabs/solace-agent-mesh
 Project-URL: Repository, https://github.com/SolaceLabs/solace-agent-mesh