aegra-api 0.1.0__py3-none-any.whl

aegra_api/services/broker.py

@@ -0,0 +1,141 @@
+"""Event broker for managing run-specific event queues"""
+
+import asyncio
+import contextlib
+from collections.abc import AsyncIterator
+from typing import Any
+
+import structlog
+
+from aegra_api.services.base_broker import BaseBrokerManager, BaseRunBroker
+
+logger = structlog.getLogger(__name__)
+
+
+class RunBroker(BaseRunBroker):
+    """Manages event queuing and distribution for a specific run"""
+
+    def __init__(self, run_id: str):
+        self.run_id = run_id
+        self.queue: asyncio.Queue[tuple[str, Any]] = asyncio.Queue()
+        self.finished = asyncio.Event()
+        self._created_at = asyncio.get_event_loop().time()
+
+    async def put(self, event_id: str, payload: Any) -> None:
+        """Put an event into the broker queue"""
+        if self.finished.is_set():
+            logger.warning(f"Attempted to put event {event_id} into finished broker for run {self.run_id}")
+            return
+
+        await self.queue.put((event_id, payload))
+
+        # Check if this is an end event
+        if isinstance(payload, tuple) and len(payload) >= 1 and payload[0] == "end":
+            self.mark_finished()
+
+    async def aiter(self) -> AsyncIterator[tuple[str, Any]]:
+        """Async iterator yielding (event_id, payload) pairs"""
+        while True:
+            try:
+                # Use timeout to check if run is finished
+                event_id, payload = await asyncio.wait_for(self.queue.get(), timeout=0.1)
+                yield event_id, payload
+
+                # Check if this is an end event
+                if isinstance(payload, tuple) and len(payload) >= 1 and payload[0] == "end":
+                    break
+
+            except TimeoutError:
+                # Check if run is finished and queue is empty
+                if self.finished.is_set() and self.queue.empty():
+                    break
+                continue
+
+    def mark_finished(self) -> None:
+        """Mark this broker as finished"""
+        self.finished.set()
+        logger.debug(f"Broker for run {self.run_id} marked as finished")
+
+    def is_finished(self) -> bool:
+        """Check if this broker is finished"""
+        return self.finished.is_set()
+
+    def is_empty(self) -> bool:
+        """Check if the queue is empty"""
+        return self.queue.empty()
+
+    def get_age(self) -> float:
+        """Get the age of this broker in seconds"""
+        return asyncio.get_event_loop().time() - self._created_at
+
+
+class BrokerManager(BaseBrokerManager):
+    """Manages multiple RunBroker instances"""
+
+    def __init__(self) -> None:
+        self._brokers: dict[str, RunBroker] = {}
+        self._cleanup_task: asyncio.Task | None = None
+
+    def get_or_create_broker(self, run_id: str) -> RunBroker:
+        """Get or create a broker for a run"""
+        if run_id not in self._brokers:
+            self._brokers[run_id] = RunBroker(run_id)
+            logger.debug(f"Created new broker for run {run_id}")
+        return self._brokers[run_id]
+
+    def get_broker(self, run_id: str) -> RunBroker | None:
+        """Get an existing broker or None"""
+        return self._brokers.get(run_id)
+
+    def cleanup_broker(self, run_id: str) -> None:
+        """Clean up a broker for a run"""
+        if run_id in self._brokers:
+            self._brokers[run_id].mark_finished()
+            # Don't immediately delete in case there are still consumers
+            logger.debug(f"Marked broker for run {run_id} for cleanup")
+
+    def remove_broker(self, run_id: str) -> None:
+        """Remove a broker completely"""
+        if run_id in self._brokers:
+            self._brokers[run_id].mark_finished()
+            del self._brokers[run_id]
+            logger.debug(f"Removed broker for run {run_id}")
+
+    async def start_cleanup_task(self) -> None:
+        """Start background cleanup task for old brokers"""
+        if self._cleanup_task is None or self._cleanup_task.done():
+            self._cleanup_task = asyncio.create_task(self._cleanup_old_brokers())
+
+    async def stop_cleanup_task(self) -> None:
+        """Stop background cleanup task"""
+        if self._cleanup_task and not self._cleanup_task.done():
+            self._cleanup_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._cleanup_task
+
+    async def _cleanup_old_brokers(self) -> None:
+        """Background task to clean up old finished brokers"""
+        while True:
+            try:
+                await asyncio.sleep(300)  # Check every 5 minutes
+
+                asyncio.get_event_loop().time()
+                to_remove = []
+
+                for run_id, broker in self._brokers.items():
+                    # Remove brokers that are finished and older than 1 hour
+                    if broker.is_finished() and broker.is_empty() and broker.get_age() > 3600:
+                        to_remove.append(run_id)
+
+                for run_id in to_remove:
+                    self.remove_broker(run_id)
+                    logger.info(f"Cleaned up old broker for run {run_id}")
+
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                logger.error(f"Error in broker cleanup task: {e}")
+
+
+# Global broker manager instance
+broker_manager = BrokerManager()

aegra_api/services/event_converter.py

@@ -0,0 +1,157 @@
+"""Event converter for SSE streaming"""
+
+from typing import Any
+
+from aegra_api.core.sse import (
+    create_debug_event,
+    create_end_event,
+    create_error_event,
+    create_messages_event,
+    create_metadata_event,
+    format_sse_message,
+)
+
+
+class EventConverter:
+    """Converts events to SSE format"""
+
+    def __init__(self):
+        """Initialize event converter"""
+        self.subgraphs: bool = False
+
+    def set_subgraphs(self, subgraphs: bool) -> None:
+        """Set whether subgraphs mode is enabled for namespace extraction"""
+        self.subgraphs = subgraphs
+
+    def convert_raw_to_sse(self, event_id: str, raw_event: Any) -> str | None:
+        """Convert raw event to SSE format"""
+        stream_mode, payload, namespace = self._parse_raw_event(raw_event)
+        return self._create_sse_event(stream_mode, payload, event_id, namespace)
+
+    def convert_stored_to_sse(self, stored_event, run_id: str = None) -> str | None:
+        """Convert stored event to SSE format"""
+        event_type = stored_event.event
+        data = stored_event.data
+        event_id = stored_event.id
+
+        # Handle special cases with custom logic
+        if event_type == "messages":
+            message_chunk = data.get("message_chunk")
+            metadata = data.get("metadata")
+            if message_chunk is None:
+                return None
+            message_data = (message_chunk, metadata) if metadata is not None else message_chunk
+            return create_messages_event(message_data, event_id=event_id)
+        elif event_type == "metadata":
+            return create_metadata_event(run_id, event_id)
+        elif event_type == "debug":
+            return create_debug_event(data.get("debug"), event_id)
+        elif event_type == "end":
+            return create_end_event(event_id)
+        elif event_type == "error":
+            return create_error_event(data.get("error"), event_id)
+        else:
+            # Handle all other event types generically (values, state, logs, tasks, etc.)
+            # Extract payload - try common patterns
+            payload = data.get(event_type) or data.get("chunk") or data
+            return format_sse_message(event_type, payload, event_id)
+
+    def _parse_raw_event(self, raw_event: Any) -> tuple[str, Any, list[str] | None]:
+        """
+        Parse raw event into (stream_mode, payload, namespace).
+
+        When subgraphs=True, 3-tuple format is (namespace, mode, chunk).
+        When subgraphs=False, 3-tuple format is (node_path, mode, chunk) for legacy support.
+        """
+        namespace = None
+
+        if isinstance(raw_event, tuple):
+            if len(raw_event) == 2:
+                # Standard format: (mode, chunk)
+                return raw_event[0], raw_event[1], None
+            elif len(raw_event) == 3:
+                if self.subgraphs:
+                    # Subgraphs format: (namespace, mode, chunk)
+                    namespace, mode, chunk = raw_event
+                    # Normalize namespace to list format
+                    if namespace is None or (isinstance(namespace, (list, tuple)) and not namespace):
+                        # Handle None or empty tuple/list - no namespace prefix
+                        namespace_list = None
+                    elif isinstance(namespace, (list, tuple)):
+                        # Convert tuple/list to list of strings
+                        namespace_list = [str(item) for item in namespace]
+                    elif isinstance(namespace, str):
+                        # Handle string namespace (shouldn't happen but be safe)
+                        namespace_list = [namespace] if namespace else None
+                    else:
+                        # Fallback - shouldn't reach here
+                        namespace_list = [str(namespace)]
+                    return mode, chunk, namespace_list
+                else:
+                    # Legacy format: (node_path, mode, chunk)
+                    return raw_event[1], raw_event[2], None
+
+        # Non-tuple events are values mode
+        return "values", raw_event, None
+
+    def _create_sse_event(
+        self,
+        stream_mode: str,
+        payload: Any,
+        event_id: str,
+        namespace: list[str] | None = None,
+    ) -> str | None:
+        """
+        Create SSE event based on stream mode.
+
+        Args:
+            stream_mode: The stream mode (e.g., "messages", "values")
+            payload: The event payload
+            event_id: The event ID
+            namespace: Optional namespace for subgraph events (e.g., ["subagent_name"])
+
+        Returns:
+            SSE-formatted event string or None
+        """
+        # Prefix event type with namespace if subgraphs enabled
+        if self.subgraphs and namespace:
+            event_type = f"{stream_mode}|{'|'.join(namespace)}"
+        else:
+            event_type = stream_mode
+
+        # Handle updates events (rarely reached - updates are filtered in graph_streaming)
+        if stream_mode == "updates":
+            if isinstance(payload, dict) and "__interrupt__" in payload:
+                # Convert interrupt updates to values events
+                if self.subgraphs and namespace:
+                    event_type = f"values|{'|'.join(namespace)}"
+                else:
+                    event_type = "values"
+                return format_sse_message(event_type, payload, event_id)
+            else:
+                # Non-interrupt updates (pass through as-is)
+                return format_sse_message(event_type, payload, event_id)
+
+        # Handle specific message event types (Studio compatibility and standard messages)
+        if stream_mode in (
+            "messages/metadata",
+            "messages/partial",
+            "messages/complete",
+        ):
+            # Studio-specific message events - pass through as-is
+            return format_sse_message(stream_mode, payload, event_id)
+        elif stream_mode == "messages" or event_type.startswith("messages"):
+            return create_messages_event(payload, event_type=event_type, event_id=event_id)
+        elif stream_mode == "values" or event_type.startswith("values"):
+            # For values events, use format_sse_message directly to support namespaces
+            return format_sse_message(event_type, payload, event_id)
+        elif stream_mode == "debug":
+            return create_debug_event(payload, event_id)
+        elif stream_mode == "end":
+            return create_end_event(event_id)
+        elif stream_mode == "error":
+            return create_error_event(payload, event_id)
+        else:
+            # Generic handler for all other event types (state, logs, tasks, events, etc.)
+            # This automatically supports any new event types without code changes
+            return format_sse_message(event_type, payload, event_id)

aegra_api/services/event_store.py

@@ -0,0 +1,196 @@
+"""Persistent event store for SSE replay functionality (Postgres-backed)."""
+
+import asyncio
+import contextlib
+import json
+from datetime import UTC, datetime
+
+import structlog
+from psycopg.types.json import Jsonb
+
+from aegra_api.core.database import db_manager
+from aegra_api.core.serializers import GeneralSerializer
+from aegra_api.core.sse import SSEEvent
+
+logger = structlog.get_logger(__name__)
+
+
+class EventStore:
+    """Postgres-backed event store for SSE replay functionality"""
+
+    CLEANUP_INTERVAL = 300  # seconds
+
+    def __init__(self) -> None:
+        self._cleanup_task: asyncio.Task | None = None
+
+    async def start_cleanup_task(self) -> None:
+        if self._cleanup_task is None or self._cleanup_task.done():
+            self._cleanup_task = asyncio.create_task(self._cleanup_loop())
+
+    async def stop_cleanup_task(self) -> None:
+        if self._cleanup_task and not self._cleanup_task.done():
+            self._cleanup_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._cleanup_task
+
+    async def store_event(self, run_id: str, event: SSEEvent) -> None:
+        """Persist an event with sequence extracted from id suffix.
+
+        We expect event.id format: f"{run_id}_event_{seq}".
+        """
+        try:
+            seq = int(str(event.id).split("_event_")[-1])
+        except Exception:
+            seq = 0
+
+        # USE SHARED POOL
+        if not db_manager.lg_pool:
+            logger.error("Database pool not initialized!")
+            return
+
+        async with db_manager.lg_pool.connection() as conn, conn.cursor() as cur:
+            await cur.execute(
+                """
+                INSERT INTO run_events (id, run_id, seq, event, data, created_at)
+                VALUES (%(id)s, %(run_id)s, %(seq)s, %(event)s, %(data)s, NOW())
+                ON CONFLICT (id) DO NOTHING
+                """,
+                {
+                    "id": event.id,
+                    "run_id": run_id,
+                    "seq": seq,
+                    "event": event.event,
+                    "data": Jsonb(event.data),
+                },
+            )
+
+    async def get_events_since(self, run_id: str, last_event_id: str) -> list[SSEEvent]:
+        """Fetch all events for run after last_event_id sequence."""
+        try:
+            last_seq = int(str(last_event_id).split("_event_")[-1])
+        except Exception:
+            last_seq = -1
+
+        if not db_manager.lg_pool:
+            return []
+
+        async with db_manager.lg_pool.connection() as conn, conn.cursor() as cur:
+            await cur.execute(
+                """
+                SELECT id, event, data, created_at
+                FROM run_events
+                WHERE run_id = %(run_id)s AND seq > %(last_seq)s
+                ORDER BY seq ASC
+                """,
+                {"run_id": run_id, "last_seq": last_seq},
+            )
+            rows = await cur.fetchall()
+
+        return [SSEEvent(id=r["id"], event=r["event"], data=r["data"], timestamp=r["created_at"]) for r in rows]
+
+    async def get_all_events(self, run_id: str) -> list[SSEEvent]:
+        if not db_manager.lg_pool:
+            return []
+
+        async with db_manager.lg_pool.connection() as conn, conn.cursor() as cur:
+            await cur.execute(
+                """
+                SELECT id, event, data, created_at
+                FROM run_events
+                WHERE run_id = %(run_id)s
+                ORDER BY seq ASC
+                """,
+                {"run_id": run_id},
+            )
+            rows = await cur.fetchall()
+
+        return [SSEEvent(id=r["id"], event=r["event"], data=r["data"], timestamp=r["created_at"]) for r in rows]
+
+    async def cleanup_events(self, run_id: str) -> None:
+        if not db_manager.lg_pool:
+            return
+
+        async with db_manager.lg_pool.connection() as conn, conn.cursor() as cur:
+            await cur.execute(
+                "DELETE FROM run_events WHERE run_id = %(run_id)s",
+                {"run_id": run_id},
+            )
+
+    async def get_run_info(self, run_id: str) -> dict | None:
+        if not db_manager.lg_pool:
+            return
+
+        async with db_manager.lg_pool.connection() as conn, conn.cursor() as cur:
+            # 1. Fetch sequence range
+            await cur.execute(
+                """
+                SELECT MIN(seq) AS first_seq, MAX(seq) AS last_seq
+                FROM run_events
+                WHERE run_id = %(run_id)s
+                """,
+                {"run_id": run_id},
+            )
+            row = await cur.fetchone()
+
+            if not row or row["last_seq"] is None:
+                return None
+
+            # 2. Fetch last event
+            await cur.execute(
+                """
+                SELECT id, created_at
+                FROM run_events
+                WHERE run_id = %(run_id)s AND seq = %(last_seq)s
+                LIMIT 1
+                """,
+                {"run_id": run_id, "last_seq": row["last_seq"]},
+            )
+            last = await cur.fetchone()
+
+        return {
+            "run_id": run_id,
+            "event_count": int(row["last_seq"]) - int(row["first_seq"]) + 1 if row["first_seq"] is not None else 0,
+            "first_event_time": None,
+            "last_event_time": last["created_at"] if last else None,
+            "last_event_id": last["id"] if last else None,
+        }
+
+    async def _cleanup_loop(self) -> None:
+        while True:
+            try:
+                await asyncio.sleep(self.CLEANUP_INTERVAL)
+                await self._cleanup_old_runs()
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                logger.error(f"Error in event store cleanup: {e}")
+
+    async def _cleanup_old_runs(self) -> None:
+        # Retain events for 1 hour by default
+        if not db_manager.lg_pool:
+            return
+
+        try:
+            async with db_manager.lg_pool.connection() as conn, conn.cursor() as cur:
+                await cur.execute("DELETE FROM run_events WHERE created_at < NOW() - INTERVAL '1 hour'")
+        except Exception as e:
+            logger.error(f"Failed to cleanup old runs: {e}")
+
+
+# Global event store instance
+event_store = EventStore()
+
+
+async def store_sse_event(run_id: str, event_id: str, event_type: str, data: dict) -> SSEEvent:
+    """Store SSE event with proper serialization"""
+    serializer = GeneralSerializer()
+
+    # Ensure JSONB-safe data by serializing complex objects
+    try:
+        safe_data = json.loads(json.dumps(data, default=serializer.serialize))
+    except Exception:
+        # Fallback to stringifying as a last resort to avoid crashing the run
+        safe_data = {"raw": str(data)}
+    event = SSEEvent(id=event_id, event=event_type, data=safe_data, timestamp=datetime.now(UTC))
+    await event_store.store_event(run_id, event)
+    return event