solace-agent-mesh 0.2.4__py3-none-any.whl → 1.0.1__py3-none-any.whl

solace_agent_mesh/gateway/http_sse/session_manager.py

@@ -0,0 +1,187 @@
+"""
+Manages web user sessions and mapping to A2A Client IDs.
+"""
+
+import uuid
+from starlette.requests import Request
+from typing import Optional, Callable, Dict, Any
+
+from solace_ai_connector.common.log import log
+
+SESSION_KEY_CLIENT_ID = "a2a_client_id"
+SESSION_KEY_SESSION_ID = "a2a_session_id"
+SESSION_KEY_ACCESS_TOKEN = "access_token"
+SESSION_KEY_REFRESH_TOKEN = "refresh_token"
+SESSION_KEY_USER_ID = "user_id"
+
+
+class SessionManager:
+    """
+    Handles web user sessions using Starlette's SessionMiddleware.
+    Generates and stores unique A2A Client IDs and manages the current A2A Session ID per web session.
+    """
+
+    def __init__(self, secret_key: str, app_config: Dict[str, Any]):
+        if not secret_key:
+            raise ValueError("Session secret key cannot be empty.")
+        self.secret_key = secret_key
+        self.force_user_identity = app_config.get("force_user_identity")
+        self.default_user_identity = app_config.get("default_user_identity")
+        self._temp_code_cache = {}
+        log.info("[SessionManager] Initialized.")
+        if self.force_user_identity:
+            log.warning(
+                f"[SessionManager] Forcing user identity to: {self.force_user_identity}"
+            )
+
+    def _get_or_create_client_id(self, request: Request) -> str:
+        """
+        Retrieves the A2A Client ID. It prioritizes the authenticated user from
+        `request.state.user` and falls back to session-based or generated IDs.
+        """
+        if self.force_user_identity:
+            return self.force_user_identity
+
+        if hasattr(request.state, "user") and request.state.user:
+            user_id = request.state.user.get("id")
+            if user_id:
+                log.debug(
+                    "[SessionManager] Using authenticated user ID from request.state: %s",
+                    user_id,
+                )
+                return user_id
+            else:
+                log.warning(
+                    "[SessionManager] request.state.user exists but has no 'id' field. Falling back to other methods."
+                )
+
+        user_id = self.get_user_id(request)
+        if user_id:
+            log.debug(
+                "[SessionManager] Using authenticated user_id from session as A2A Client ID: %s",
+                user_id,
+            )
+            return user_id
+
+        client_id = request.session.get(SESSION_KEY_CLIENT_ID)
+        if not client_id:
+            if self.default_user_identity:
+                log.info(
+                    "[SessionManager] Using default_user_identity as A2A Client ID: %s",
+                    self.default_user_identity,
+                )
+                client_id = self.default_user_identity
+            else:
+                client_id = f"web-client-{uuid.uuid4().hex}"
+                log.info(
+                    "[SessionManager] Created new A2A Client ID: %s for web session.",
+                    client_id,
+                )
+            request.session[SESSION_KEY_CLIENT_ID] = client_id
+        else:
+            log.debug(
+                "[SessionManager] Using existing A2A Client ID: %s for web session.",
+                client_id,
+            )
+        return client_id
+
+    def get_a2a_client_id(self, request: Request) -> str:
+        """
+        FastAPI dependency callable to get the A2A Client ID for the current request.
+        Ensures a client ID exists in the session.
+        """
+        return self._get_or_create_client_id(request)
+
+    def get_a2a_session_id(self, request: Request) -> Optional[str]:
+        """
+        FastAPI dependency callable to get the current A2A Session ID for the current request.
+        Returns None if no session has been started for the current agent in this web session.
+        """
+        session_id = request.session.get(SESSION_KEY_SESSION_ID)
+        log.debug("[SessionManager] Retrieving A2A Session ID: %s", session_id)
+        return session_id
+
+    def start_new_a2a_session(self, request: Request) -> str:
+        """
+        Generates a new A2A Session ID, stores it in the web session, and returns it.
+        This should be called when the user explicitly starts a new chat or switches agents.
+        """
+        client_id = self._get_or_create_client_id(request)
+        new_session_id = f"web-session-{uuid.uuid4().hex}"
+        request.session[SESSION_KEY_SESSION_ID] = new_session_id
+        log.info(
+            "[SessionManager] Started new A2A Session ID: %s for Client ID: %s",
+            new_session_id,
+            client_id,
+        )
+        return new_session_id
+
+    def ensure_a2a_session(self, request: Request) -> str:
+        """
+        Ensures an A2A session ID exists, creating one if necessary.
+        Use this when an operation requires a session ID but one might not have been explicitly started yet.
+        """
+        session_id = self.get_a2a_session_id(request)
+        if not session_id:
+            session_id = self.start_new_a2a_session(request)
+            log.info(
+                "[SessionManager] No A2A Session ID found, created new one via ensure_a2a_session: %s",
+                session_id,
+            )
+        return session_id
+
+    def store_auth_tokens(
+        self, request: Request, access_token: str, refresh_token: Optional[str] = None
+    ):
+        """
+        Stores authentication tokens directly in the user's session.
+        """
+        request.session[SESSION_KEY_ACCESS_TOKEN] = access_token
+        if refresh_token:
+            request.session[SESSION_KEY_REFRESH_TOKEN] = refresh_token
+        log.info("[SessionManager] Stored auth tokens directly in session.")
+
+    def get_access_token(self, request: Request) -> Optional[str]:
+        """
+        Retrieves the access token from the web session.
+        """
+        return request.session.get(SESSION_KEY_ACCESS_TOKEN)
+
+    def get_refresh_token(self, request: Request) -> Optional[str]:
+        """
+        Retrieves the refresh token from the web session.
+        """
+        return request.session.get(SESSION_KEY_REFRESH_TOKEN)
+
+    def clear_auth_tokens(self, request: Request) -> None:
+        """
+        Clears authentication tokens from the web session.
+        """
+        request.session.pop(SESSION_KEY_ACCESS_TOKEN, None)
+        request.session.pop(SESSION_KEY_REFRESH_TOKEN, None)
+        log.info("[SessionManager] Cleared auth tokens from session")
+
+    def store_user_id(self, request: Request, user_id: str) -> None:
+        """
+        Stores the user ID in the web session.
+        """
+        request.session[SESSION_KEY_USER_ID] = user_id
+        log.info("[SessionManager] Stored user ID in session: %s", user_id)
+
+    def get_user_id(self, request: Request) -> Optional[str]:
+        """
+        Retrieves the user ID from the web session.
+        """
+        return request.session.get(SESSION_KEY_USER_ID)
+
+    def dep_get_client_id(self) -> Callable[[Request], str]:
+        """Returns a callable suitable for FastAPI Depends to get the client ID."""
+        return self.get_a2a_client_id
+
+    def dep_get_session_id(self) -> Callable[[Request], Optional[str]]:
+        """Returns a callable suitable for FastAPI Depends to get the current session ID."""
+        return self.get_a2a_session_id
+
+    def dep_ensure_session_id(self) -> Callable[[Request], str]:
+        """Returns a callable suitable for FastAPI Depends to ensure a session ID exists."""
+        return self.ensure_a2a_session

solace_agent_mesh/gateway/http_sse/sse_manager.py

@@ -0,0 +1,328 @@
+"""
+Manages Server-Sent Event (SSE) connections for streaming task updates.
+"""
+
+import asyncio
+import threading
+from typing import Dict, List, Any
+import json
+
+from solace_ai_connector.common.log import log
+
+
+class SSEManager:
+    """
+    Manages active SSE connections and distributes events based on task ID.
+    Uses asyncio Queues for buffering events per connection.
+    """
+
+    def __init__(self, max_queue_size: int = 200):
+        self._connections: Dict[str, List[asyncio.Queue]] = {}
+        self._locks: Dict[asyncio.AbstractEventLoop, asyncio.Lock] = {}
+        self._locks_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]
+
+    async def create_sse_connection(self, task_id: str) -> asyncio.Queue:
+        """
+        Creates a new queue for an SSE connection subscribing to a task.
+
+        Args:
+            task_id: The ID of the task the connection is interested in.
+
+        Returns:
+            An asyncio.Queue that the SSE endpoint can consume from.
+        """
+        lock = self._get_lock()
+        async with lock:
+            if task_id not in self._connections:
+                self._connections[task_id] = []
+
+            connection_queue = asyncio.Queue(maxsize=self._max_queue_size)
+            self._connections[task_id].append(connection_queue)
+            log.info(
+                "%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
+
+    async def remove_sse_connection(
+        self, task_id: str, connection_queue: asyncio.Queue
+    ):
+        """
+        Removes a specific SSE connection queue for a task.
+
+        Args:
+            task_id: The ID of the task.
+            connection_queue: The specific queue instance to remove.
+        """
+        lock = self._get_lock()
+        async with lock:
+            if task_id in self._connections:
+                try:
+                    self._connections[task_id].remove(connection_queue)
+                    log.info(
+                        "%s Removed SSE connection queue for Task ID: %s. Remaining queues: %d",
+                        self.log_identifier,
+                        task_id,
+                        len(self._connections[task_id]),
+                    )
+                    if not self._connections[task_id]:
+                        del self._connections[task_id]
+                        log.info(
+                            "%s Removed Task ID entry: %s as no connections remain.",
+                            self.log_identifier,
+                            task_id,
+                        )
+                except ValueError:
+                    log.debug(
+                        "%s Attempted to remove an already removed queue for Task ID: %s.",
+                        self.log_identifier,
+                        task_id,
+                    )
+            else:
+                log.warning(
+                    "%s Attempted to remove queue for non-existent Task ID: %s.",
+                    self.log_identifier,
+                    task_id,
+                )
+
+    async def send_event(
+        self, task_id: str, event_data: Dict[str, Any], event_type: str = "message"
+    ):
+        """
+        Sends an event (as a dictionary) to all active SSE connections for a specific task.
+        The event_data dictionary will be JSON serialized for the SSE 'data' field.
+
+        Args:
+            task_id: The ID of the task the event belongs to.
+            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:
+            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_to_remove = []
+            try:
+                serialized_data = json.dumps(event_data)
+            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}
+            log.debug(
+                "%s Prepared SSE payload for Task ID %s: %s",
+                self.log_identifier,
+                task_id,
+                sse_payload,
+            )
+
+            for connection_queue in list(self._connections.get(task_id, [])):
+                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)
+
+            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.log_identifier,
+                            task_id,
+                        )
+                    except ValueError:
+                        pass
+
+                if not current_queues:
+                    del self._connections[task_id]
+                    log.info(
+                        "%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):
+        """
+        Signals a specific SSE connection queue to close by putting None.
+        Also removes the queue from the manager.
+        """
+        log.info(
+            "%s Closing specific SSE connection queue for Task ID: %s",
+            self.log_identifier,
+            task_id,
+        )
+        try:
+            await asyncio.wait_for(connection_queue.put(None), timeout=0.1)
+        except asyncio.QueueFull:
+            log.warning(
+                "%s Could not put None (close signal) on full queue for Task ID: %s. Connection might not close cleanly.",
+                self.log_identifier,
+                task_id,
+            )
+        except asyncio.TimeoutError:
+            log.warning(
+                "%s Timeout putting None (close signal) on queue for Task ID: %s.",
+                self.log_identifier,
+                task_id,
+            )
+        except Exception as e:
+            log.error(
+                "%s Error putting None (close signal) on queue for Task ID %s: %s",
+                self.log_identifier,
+                task_id,
+                e,
+            )
+        finally:
+            await self.remove_sse_connection(task_id, connection_queue)
+
+    async def close_all_for_task(self, task_id: str):
+        """
+        Closes all SSE connections associated with a specific task.
+        """
+        lock = self._get_lock()
+        async with lock:
+            if task_id in self._connections:
+                queues_to_close = self._connections.pop(task_id)
+                log.info(
+                    "%s Closing %d SSE connections for Task ID: %s",
+                    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,
+                        )
+                log.info(
+                    "%s Removed Task ID entry: %s and signaled queues to close.",
+                    self.log_identifier,
+                    task_id,
+                )
+            else:
+                log.debug(
+                    "%s No connections found to close for Task ID: %s",
+                    self.log_identifier,
+                    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]
+                log.debug(
+                    "%s Cleaned up lock for closed event loop %s",
+                    self.log_identifier,
+                    id(loop),
+                )
+
+    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:
+            log.info("%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.info(
+                "%s Closed %d connections for tasks: %s",
+                self.log_identifier,
+                closed_count,
+                all_task_ids,
+            )
+            self._connections.clear()