167 lines
6.1 KiB
Python
167 lines
6.1 KiB
Python
# Copyright (c) 2026 PowerOn AG
|
|
# All rights reserved.
|
|
"""
|
|
Event manager for SSE streaming (Layer L0 - shared).
|
|
Manages event queues for Server-Sent Events (SSE) streaming across features.
|
|
Generic pub/sub infrastructure with zero internal dependencies.
|
|
"""
|
|
|
|
import logging
|
|
import asyncio
|
|
from typing import Dict, Optional, Any
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class EventManager:
|
|
"""
|
|
Manages event queues for SSE streaming.
|
|
Each workflow has its own async queue for events.
|
|
"""
|
|
|
|
def __init__(self):
|
|
"""Initialize the event manager."""
|
|
self._queues: Dict[str, asyncio.Queue] = {}
|
|
self._cleanup_tasks: Dict[str, asyncio.Task] = {}
|
|
self._agent_tasks: Dict[str, asyncio.Task] = {}
|
|
self._cancelled: Dict[str, bool] = {}
|
|
|
|
def create_queue(self, workflow_id: str) -> asyncio.Queue:
|
|
"""Create an event queue for a workflow."""
|
|
if workflow_id in self._cleanup_tasks:
|
|
self._cleanup_tasks[workflow_id].cancel()
|
|
del self._cleanup_tasks[workflow_id]
|
|
logger.debug(f"Cancelled pending cleanup for workflow {workflow_id}")
|
|
|
|
if workflow_id not in self._queues:
|
|
self._queues[workflow_id] = asyncio.Queue()
|
|
logger.debug(f"Created event queue for workflow {workflow_id}")
|
|
else:
|
|
old = self._queues[workflow_id]
|
|
while not old.empty():
|
|
try:
|
|
old.get_nowait()
|
|
except asyncio.QueueEmpty:
|
|
break
|
|
logger.debug(f"Reusing event queue for workflow {workflow_id} (drained stale events)")
|
|
return self._queues[workflow_id]
|
|
|
|
def get_queue(self, workflow_id: str) -> Optional[asyncio.Queue]:
|
|
"""Get the event queue for a workflow."""
|
|
return self._queues.get(workflow_id)
|
|
|
|
def has_queue(self, workflow_id: str) -> bool:
|
|
"""Check if a queue exists for a workflow."""
|
|
return workflow_id in self._queues
|
|
|
|
def register_agent_task(self, workflow_id: str, task: asyncio.Task) -> None:
|
|
"""Register the asyncio Task running the agent for a workflow."""
|
|
self._agent_tasks[workflow_id] = task
|
|
self._cancelled.pop(workflow_id, None)
|
|
|
|
def is_cancelled(self, workflow_id: str) -> bool:
|
|
"""Check if a workflow has been cancelled."""
|
|
return self._cancelled.get(workflow_id, False)
|
|
|
|
async def cancel_agent(self, workflow_id: str) -> bool:
|
|
"""Cancel the running agent task for a workflow. Returns True if cancelled."""
|
|
self._cancelled[workflow_id] = True
|
|
task = self._agent_tasks.pop(workflow_id, None)
|
|
if task and not task.done():
|
|
task.cancel()
|
|
logger.info(f"Cancelled agent task for workflow {workflow_id}")
|
|
return True
|
|
logger.debug(f"No running agent task found for workflow {workflow_id}")
|
|
return False
|
|
|
|
def _unregister_agent_task(self, workflow_id: str) -> None:
|
|
"""Remove the agent task reference after completion."""
|
|
self._agent_tasks.pop(workflow_id, None)
|
|
self._cancelled.pop(workflow_id, None)
|
|
|
|
async def emit_event(
|
|
self,
|
|
context_id: str,
|
|
event_type: str,
|
|
data: Dict[str, Any],
|
|
event_category: str = "chat",
|
|
message: Optional[str] = None,
|
|
step: Optional[str] = None
|
|
) -> None:
|
|
"""Emit an event to the queue for a workflow."""
|
|
queue = self._queues.get(context_id)
|
|
if not queue:
|
|
return
|
|
|
|
event = {
|
|
"type": event_type,
|
|
"data": data,
|
|
"category": event_category,
|
|
"message": message,
|
|
"step": step
|
|
}
|
|
|
|
try:
|
|
await queue.put(event)
|
|
if event_type not in ("chunk",):
|
|
logger.debug(f"Emitted {event_type} event for workflow {context_id}")
|
|
except Exception as e:
|
|
logger.error(f"Error emitting event for workflow {context_id}: {e}", exc_info=True)
|
|
|
|
async def cleanup(self, workflow_id: str, delay: float = 60.0) -> None:
|
|
"""Schedule cleanup of a queue after a delay."""
|
|
if workflow_id in self._cleanup_tasks:
|
|
self._cleanup_tasks[workflow_id].cancel()
|
|
|
|
async def _cleanup():
|
|
try:
|
|
await asyncio.sleep(delay)
|
|
if workflow_id in self._queues:
|
|
queue = self._queues[workflow_id]
|
|
while not queue.empty():
|
|
try:
|
|
queue.get_nowait()
|
|
except asyncio.QueueEmpty:
|
|
break
|
|
del self._queues[workflow_id]
|
|
logger.info(f"Cleaned up event queue for workflow {workflow_id}")
|
|
except asyncio.CancelledError:
|
|
logger.debug(f"Cleanup cancelled for workflow {workflow_id}")
|
|
except Exception as e:
|
|
logger.error(f"Error during cleanup for workflow {workflow_id}: {e}", exc_info=True)
|
|
finally:
|
|
if workflow_id in self._cleanup_tasks:
|
|
del self._cleanup_tasks[workflow_id]
|
|
|
|
task = asyncio.create_task(_cleanup())
|
|
self._cleanup_tasks[workflow_id] = task
|
|
|
|
def shutdown(self) -> None:
|
|
"""Cancel all pending cleanup and agent tasks for fast process exit."""
|
|
for _wfId, q in list(self._queues.items()):
|
|
try:
|
|
q.put_nowait(None)
|
|
except Exception:
|
|
pass
|
|
for wfId, task in list(self._cleanup_tasks.items()):
|
|
if not task.done():
|
|
task.cancel()
|
|
self._cleanup_tasks.clear()
|
|
for wfId, task in list(self._agent_tasks.items()):
|
|
if not task.done():
|
|
task.cancel()
|
|
self._agent_tasks.clear()
|
|
self._queues.clear()
|
|
logger.info("EventManager shutdown: all tasks cancelled, queues drained")
|
|
|
|
|
|
# Global event manager instance
|
|
_event_manager: Optional[EventManager] = None
|
|
|
|
|
|
def get_event_manager() -> EventManager:
|
|
"""Get the global event manager instance."""
|
|
global _event_manager
|
|
if _event_manager is None:
|
|
_event_manager = EventManager()
|
|
return _event_manager
|