solace-agent-mesh 1.7.1__py3-none-any.whl → 1.13.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 ....shared.utils.timestamp_utils 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

@@ -10,7 +10,7 @@ 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
+from solace_agent_mesh.shared.utils.timestamp_utils import now_epoch_ms
 
 log = logging.getLogger(__name__)
 

solace_agent_mesh/gateway/http_sse/services/feedback_service.py

@@ -11,7 +11,7 @@ from sqlalchemy.orm import Session as DBSession
 
 from ..repository.entities import Feedback
 from ..repository.feedback_repository import FeedbackRepository
-from ..shared import now_epoch_ms
+from solace_agent_mesh.shared.utils.timestamp_utils import now_epoch_ms
 from ..utils.stim_utils import create_stim_from_task_data
 
 # The FeedbackPayload is defined in the router, this creates a forward reference