solace-agent-mesh 1.11.2__py3-none-any.whl

solace_agent_mesh/gateway/http_sse/services/background_task_monitor.py

@@ -0,0 +1,186 @@
+"""
+Background Task Monitor Service.
+Monitors background tasks and enforces timeouts.
+"""
+
+import logging
+from typing import Callable
+from sqlalchemy.orm import Session as DBSession
+
+from ....gateway.http_sse.repository.task_repository import TaskRepository
+from ....gateway.http_sse.services.task_service import TaskService
+from ....gateway.http_sse.shared import now_epoch_ms
+
+log = logging.getLogger(__name__)
+
+
+class BackgroundTaskMonitor:
+    """
+    Monitors background tasks and enforces timeouts.
+    Runs as a periodic background job in the gateway.
+    """
+
+    def __init__(
+        self,
+        session_factory: Callable[[], DBSession],
+        task_service: TaskService,
+        default_timeout_ms: int = 3600000,  # 1 hour default
+    ):
+        """
+        Initialize the background task monitor.
+        
+        Args:
+            session_factory: Factory function to create database sessions
+            task_service: TaskService instance for cancelling tasks
+            default_timeout_ms: Default timeout in milliseconds if task doesn't specify one
+        """
+        self.session_factory = session_factory
+        self.task_service = task_service
+        self.default_timeout_ms = default_timeout_ms
+        self.log_identifier = "[BackgroundTaskMonitor]"
+        log.info(f"{self.log_identifier} Initialized with default timeout {default_timeout_ms}ms")
+
+    async def check_timeouts(self) -> dict:
+        """
+        Check for timed-out background tasks and cancel them.
+        
+        Returns:
+            Dictionary with statistics about the check
+        """
+        log.debug(f"{self.log_identifier} Starting timeout check")
+        
+        if not self.session_factory:
+            log.warning(f"{self.log_identifier} No database session factory, skipping timeout check")
+            return {"checked": 0, "timed_out": 0, "cancelled": 0}
+        
+        db = self.session_factory()
+        try:
+            repo = TaskRepository()
+            current_time = now_epoch_ms()
+            
+            # Get all running background tasks
+            running_tasks = repo.find_background_tasks_by_status(db, status=None)
+            running_tasks = [
+                task for task in running_tasks
+                if task.status in [None, "running", "pending"] and task.end_time is None
+            ]
+            
+            log.debug(f"{self.log_identifier} Found {len(running_tasks)} running background tasks")
+            
+            timed_out_tasks = []
+            
+            for task in running_tasks:
+                # Determine timeout for this task
+                task_timeout = task.max_execution_time_ms or self.default_timeout_ms
+                
+                # Check if task has exceeded timeout based on last activity
+                if task.last_activity_time:
+                    time_since_activity = current_time - task.last_activity_time
+                    if time_since_activity > task_timeout:
+                        timed_out_tasks.append({
+                            "task": task,
+                            "timeout_ms": task_timeout,
+                            "time_since_activity": time_since_activity
+                        })
+                        log.warning(
+                            f"{self.log_identifier} Task {task.id} exceeded timeout: "
+                            f"{time_since_activity}ms since last activity (timeout: {task_timeout}ms)"
+                        )
+            
+            # Cancel timed-out tasks
+            cancelled_count = 0
+            for item in timed_out_tasks:
+                task = item["task"]
+                try:
+                    # Extract agent name from task metadata or use a default
+                    # In a real implementation, we'd need to store this in the task record
+                    log.info(f"{self.log_identifier} Cancelling timed-out task {task.id}")
+                    
+                    # Update task status in database
+                    task.status = "timeout"
+                    task.end_time = current_time
+                    repo.save_task(db, task)
+                    
+                    # Note: We can't easily cancel the actual agent task from here
+                    # without knowing the agent name. This would require storing
+                    # the agent name in the task record or having a different
+                    # cancellation mechanism.
+                    
+                    cancelled_count += 1
+                    log.info(f"{self.log_identifier} Marked task {task.id} as timed out")
+                    
+                except Exception as e:
+                    log.error(
+                        f"{self.log_identifier} Failed to cancel task {task.id}: {e}",
+                        exc_info=True
+                    )
+            
+            db.commit()
+            
+            stats = {
+                "checked": len(running_tasks),
+                "timed_out": len(timed_out_tasks),
+                "cancelled": cancelled_count,
+                "timestamp": current_time
+            }
+            
+            if timed_out_tasks:
+                log.info(
+                    f"{self.log_identifier} Timeout check complete: "
+                    f"checked={stats['checked']}, timed_out={stats['timed_out']}, "
+                    f"cancelled={stats['cancelled']}"
+                )
+            else:
+                log.debug(f"{self.log_identifier} Timeout check complete: no timeouts detected")
+            
+            return stats
+            
+        except Exception as e:
+            log.error(
+                f"{self.log_identifier} Error during timeout check: {e}",
+                exc_info=True
+            )
+            db.rollback()
+            return {"checked": 0, "timed_out": 0, "cancelled": 0, "error": str(e)}
+        finally:
+            db.close()
+
+    async def get_background_task_stats(self) -> dict:
+        """
+        Get statistics about background tasks.
+        
+        Returns:
+            Dictionary with task statistics
+        """
+        if not self.session_factory:
+            return {"error": "No database configured"}
+        
+        db = self.session_factory()
+        try:
+            repo = TaskRepository()
+            
+            # Get all background tasks
+            all_background = repo.find_background_tasks_by_status(db, status=None)
+            
+            # Count by status
+            running = len([t for t in all_background if t.status in [None, "running", "pending"] and t.end_time is None])
+            completed = len([t for t in all_background if t.status == "completed"])
+            failed = len([t for t in all_background if t.status in ["failed", "error"]])
+            cancelled = len([t for t in all_background if t.status == "cancelled"])
+            timeout = len([t for t in all_background if t.status == "timeout"])
+            
+            return {
+                "total": len(all_background),
+                "running": running,
+                "completed": completed,
+                "failed": failed,
+                "cancelled": cancelled,
+                "timeout": timeout,
+                "timestamp": now_epoch_ms()
+            }
+            
+        except Exception as e:
+            log.error(f"{self.log_identifier} Error getting stats: {e}", exc_info=True)
+            return {"error": str(e)}
+        finally:
+            db.close()

solace_agent_mesh/gateway/http_sse/services/data_retention_service.py

@@ -0,0 +1,273 @@
+"""
+Service for managing automatic cleanup of old data based on retention policies.
+"""
+
+import logging
+import time
+from typing import Any, Callable, Dict
+
+from sqlalchemy.orm import Session as DBSession
+
+from ..repository.feedback_repository import FeedbackRepository
+from ..repository.task_repository import TaskRepository
+from ..shared import now_epoch_ms
+
+log = logging.getLogger(__name__)
+
+class DataRetentionService:
+    """
+    Service for automatically cleaning up old tasks, task events, and feedback
+    based on configurable retention policies.
+    """
+
+    # Validation constants
+    MIN_RETENTION_DAYS = 1
+    MIN_CLEANUP_INTERVAL_HOURS = 1
+    MIN_BATCH_SIZE = 1
+    MAX_BATCH_SIZE = 10000
+
+    def __init__(
+        self, session_factory: Callable[[], DBSession] | None, config: Dict[str, Any]
+    ):
+        """
+        Initialize the DataRetentionService.
+
+        Args:
+            session_factory: Factory function to create database sessions
+            config: Configuration dictionary with retention settings
+        """
+        self.session_factory = session_factory
+        self.config = config
+        self.log_identifier = "[DataRetentionService]"
+
+        # Validate and store configuration
+        self._validate_config()
+
+        log.info(
+            "%s Initialized with task_retention=%d days, feedback_retention=%d days, "
+            "cleanup_interval=%d hours, batch_size=%d",
+            self.log_identifier,
+            self.config.get("task_retention_days"),
+            self.config.get("feedback_retention_days"),
+            self.config.get("cleanup_interval_hours"),
+            self.config.get("batch_size"),
+        )
+
+    def _validate_config(self) -> None:
+        """
+        Validates configuration values and applies safe defaults if needed.
+        Logs warnings for invalid values.
+        """
+        # Validate task retention days
+        task_retention = self.config.get("task_retention_days", 90)
+        if task_retention < self.MIN_RETENTION_DAYS:
+            log.warning(
+                "%s task_retention_days (%d) is below minimum (%d days). Using minimum.",
+                self.log_identifier,
+                task_retention,
+                self.MIN_RETENTION_DAYS,
+            )
+            self.config["task_retention_days"] = self.MIN_RETENTION_DAYS
+        else:
+            self.config["task_retention_days"] = task_retention
+
+        # Validate feedback retention days
+        feedback_retention = self.config.get("feedback_retention_days", 90)
+        if feedback_retention < self.MIN_RETENTION_DAYS:
+            log.warning(
+                "%s feedback_retention_days (%d) is below minimum (%d days). Using minimum.",
+                self.log_identifier,
+                feedback_retention,
+                self.MIN_RETENTION_DAYS,
+            )
+            self.config["feedback_retention_days"] = self.MIN_RETENTION_DAYS
+        else:
+            self.config["feedback_retention_days"] = feedback_retention
+
+        # Validate cleanup interval
+        cleanup_interval = self.config.get("cleanup_interval_hours", 24)
+        if cleanup_interval < self.MIN_CLEANUP_INTERVAL_HOURS:
+            log.warning(
+                "%s cleanup_interval_hours (%d) is below minimum (%d hours). Using minimum.",
+                self.log_identifier,
+                cleanup_interval,
+                self.MIN_CLEANUP_INTERVAL_HOURS,
+            )
+            self.config["cleanup_interval_hours"] = self.MIN_CLEANUP_INTERVAL_HOURS
+        else:
+            self.config["cleanup_interval_hours"] = cleanup_interval
+
+        # Validate batch size
+        batch_size = self.config.get("batch_size", 1000)
+        if batch_size < self.MIN_BATCH_SIZE:
+            log.warning(
+                "%s batch_size (%d) is below minimum (%d). Using minimum.",
+                self.log_identifier,
+                batch_size,
+                self.MIN_BATCH_SIZE,
+            )
+            self.config["batch_size"] = self.MIN_BATCH_SIZE
+        elif batch_size > self.MAX_BATCH_SIZE:
+            log.warning(
+                "%s batch_size (%d) exceeds maximum (%d). Using maximum.",
+                self.log_identifier,
+                batch_size,
+                self.MAX_BATCH_SIZE,
+            )
+            self.config["batch_size"] = self.MAX_BATCH_SIZE
+        else:
+            self.config["batch_size"] = batch_size
+
+    def cleanup_old_data(self) -> None:
+        """
+        Main orchestration method for cleaning up old data.
+        Calls cleanup methods for tasks and feedback.
+        """
+        if not self.config.get("enabled", True):
+            log.warning(
+                "%s Data retention cleanup is disabled via configuration.",
+                self.log_identifier,
+            )
+            return
+
+        if not self.session_factory:
+            log.warning(
+                "%s No database session factory available. Skipping cleanup.",
+                self.log_identifier,
+            )
+            return
+
+        log.info("%s Starting data retention cleanup...", self.log_identifier)
+        start_time = time.time()
+
+        try:
+            # Cleanup old tasks
+            task_retention_days = self.config.get("task_retention_days")
+            tasks_deleted = self._cleanup_old_tasks(task_retention_days)
+
+            # Cleanup old feedback
+            feedback_retention_days = self.config.get("feedback_retention_days")
+            feedback_deleted = self._cleanup_old_feedback(feedback_retention_days)
+
+            elapsed_time = time.time() - start_time
+            log.info(
+                "%s Cleanup completed. Tasks deleted: %d, Feedback deleted: %d, Time taken: %.2f seconds",
+                self.log_identifier,
+                tasks_deleted,
+                feedback_deleted,
+                elapsed_time,
+            )
+
+        except Exception as e:
+            log.error(
+                "%s Error during data retention cleanup: %s",
+                self.log_identifier,
+                e,
+                exc_info=True,
+            )
+
+    def _cleanup_old_tasks(self, retention_days: int) -> int:
+        """
+        Deletes tasks (and their events via cascade) older than the retention period.
+
+        Args:
+            retention_days: Number of days to retain tasks
+
+        Returns:
+            Total number of tasks deleted
+        """
+        log.info(
+            "%s Cleaning up tasks older than %d days...",
+            self.log_identifier,
+            retention_days,
+        )
+
+        # Calculate cutoff time in milliseconds
+        cutoff_time_ms = now_epoch_ms() - (retention_days * 24 * 60 * 60 * 1000)
+        batch_size = self.config.get("batch_size")
+
+        db = self.session_factory()
+        try:
+            repo = TaskRepository()
+            total_deleted = repo.delete_tasks_older_than(db, cutoff_time_ms, batch_size)
+
+            if total_deleted == 0:
+                log.info(
+                    "%s No tasks found older than %d days.",
+                    self.log_identifier,
+                    retention_days,
+                )
+            else:
+                log.info(
+                    "%s Deleted %d tasks older than %d days.",
+                    self.log_identifier,
+                    total_deleted,
+                    retention_days,
+                )
+
+            return total_deleted
+
+        except Exception as e:
+            log.error(
+                "%s Error cleaning up old tasks: %s",
+                self.log_identifier,
+                e,
+                exc_info=True,
+            )
+            db.rollback()
+            return 0
+        finally:
+            db.close()
+
+    def _cleanup_old_feedback(self, retention_days: int) -> int:
+        """
+        Deletes feedback records older than the retention period.
+
+        Args:
+            retention_days: Number of days to retain feedback
+
+        Returns:
+            Total number of feedback records deleted
+        """
+        log.info(
+            "%s Cleaning up feedback older than %d days...",
+            self.log_identifier,
+            retention_days,
+        )
+
+        # Calculate cutoff time in milliseconds
+        cutoff_time_ms = now_epoch_ms() - (retention_days * 24 * 60 * 60 * 1000)
+        batch_size = self.config.get("batch_size")
+
+        db = self.session_factory()
+        try:
+            repo = FeedbackRepository()
+            total_deleted = repo.delete_feedback_older_than(db, cutoff_time_ms, batch_size)
+
+            if total_deleted == 0:
+                log.info(
+                    "%s No feedback found older than %d days.",
+                    self.log_identifier,
+                    retention_days,
+                )
+            else:
+                log.info(
+                    "%s Deleted %d feedback records older than %d days.",
+                    self.log_identifier,
+                    total_deleted,
+                    retention_days,
+                )
+
+            return total_deleted
+
+        except Exception as e:
+            log.error(
+                "%s Error cleaning up old feedback: %s",
+                self.log_identifier,
+                e,
+                exc_info=True,
+            )
+            db.rollback()
+            return 0
+        finally:
+            db.close()