feat: key management and scheduled task tools

agent/tools/scheduler/scheduler_service.py

@@ -0,0 +1,192 @@
+"""
+Background scheduler service for executing scheduled tasks
+"""
+
+import time
+import threading
+from datetime import datetime, timedelta
+from typing import Callable, Optional
+from croniter import croniter
+from common.log import logger
+
+
+class SchedulerService:
+    """
+    Background service that executes scheduled tasks
+    """
+    
+    def __init__(self, task_store, execute_callback: Callable):
+        """
+        Initialize scheduler service
+        
+        Args:
+            task_store: TaskStore instance
+            execute_callback: Function to call when executing a task
+        """
+        self.task_store = task_store
+        self.execute_callback = execute_callback
+        self.running = False
+        self.thread = None
+        self._lock = threading.Lock()
+    
+    def start(self):
+        """Start the scheduler service"""
+        with self._lock:
+            if self.running:
+                logger.warning("[Scheduler] Service already running")
+                return
+            
+            self.running = True
+            self.thread = threading.Thread(target=self._run_loop, daemon=True)
+            self.thread.start()
+            logger.info("[Scheduler] Service started")
+    
+    def stop(self):
+        """Stop the scheduler service"""
+        with self._lock:
+            if not self.running:
+                return
+            
+            self.running = False
+            if self.thread:
+                self.thread.join(timeout=5)
+            logger.info("[Scheduler] Service stopped")
+    
+    def _run_loop(self):
+        """Main scheduler loop"""
+        logger.info("[Scheduler] Scheduler loop started")
+        
+        while self.running:
+            try:
+                self._check_and_execute_tasks()
+            except Exception as e:
+                logger.error(f"[Scheduler] Error in scheduler loop: {e}")
+            
+            # Sleep for 30 seconds between checks
+            time.sleep(30)
+    
+    def _check_and_execute_tasks(self):
+        """Check for due tasks and execute them"""
+        now = datetime.now()
+        tasks = self.task_store.list_tasks(enabled_only=True)
+        
+        for task in tasks:
+            try:
+                # Check if task is due
+                if self._is_task_due(task, now):
+                    logger.info(f"[Scheduler] Executing task: {task['id']} - {task['name']}")
+                    self._execute_task(task)
+                    
+                    # Update next run time
+                    next_run = self._calculate_next_run(task, now)
+                    if next_run:
+                        self.task_store.update_task(task['id'], {
+                            "next_run_at": next_run.isoformat(),
+                            "last_run_at": now.isoformat()
+                        })
+                    else:
+                        # One-time task, disable it
+                        self.task_store.update_task(task['id'], {
+                            "enabled": False,
+                            "last_run_at": now.isoformat()
+                        })
+                        logger.info(f"[Scheduler] One-time task completed and disabled: {task['id']}")
+            except Exception as e:
+                logger.error(f"[Scheduler] Error processing task {task.get('id')}: {e}")
+    
+    def _is_task_due(self, task: dict, now: datetime) -> bool:
+        """
+        Check if a task is due to run
+        
+        Args:
+            task: Task dictionary
+            now: Current datetime
+            
+        Returns:
+            True if task should run now
+        """
+        next_run_str = task.get("next_run_at")
+        if not next_run_str:
+            # Calculate initial next_run_at
+            next_run = self._calculate_next_run(task, now)
+            if next_run:
+                self.task_store.update_task(task['id'], {
+                    "next_run_at": next_run.isoformat()
+                })
+                return False
+            return False
+        
+        try:
+            next_run = datetime.fromisoformat(next_run_str)
+            return now >= next_run
+        except:
+            return False
+    
+    def _calculate_next_run(self, task: dict, from_time: datetime) -> Optional[datetime]:
+        """
+        Calculate next run time for a task
+        
+        Args:
+            task: Task dictionary
+            from_time: Calculate from this time
+            
+        Returns:
+            Next run datetime or None for one-time tasks
+        """
+        schedule = task.get("schedule", {})
+        schedule_type = schedule.get("type")
+        
+        if schedule_type == "cron":
+            # Cron expression
+            expression = schedule.get("expression")
+            if not expression:
+                return None
+            
+            try:
+                cron = croniter(expression, from_time)
+                return cron.get_next(datetime)
+            except Exception as e:
+                logger.error(f"[Scheduler] Invalid cron expression '{expression}': {e}")
+                return None
+        
+        elif schedule_type == "interval":
+            # Interval in seconds
+            seconds = schedule.get("seconds", 0)
+            if seconds <= 0:
+                return None
+            return from_time + timedelta(seconds=seconds)
+        
+        elif schedule_type == "once":
+            # One-time task at specific time
+            run_at_str = schedule.get("run_at")
+            if not run_at_str:
+                return None
+            
+            try:
+                run_at = datetime.fromisoformat(run_at_str)
+                # Only return if in the future
+                if run_at > from_time:
+                    return run_at
+            except:
+                pass
+            return None
+        
+        return None
+    
+    def _execute_task(self, task: dict):
+        """
+        Execute a task
+        
+        Args:
+            task: Task dictionary
+        """
+        try:
+            # Call the execute callback
+            self.execute_callback(task)
+        except Exception as e:
+            logger.error(f"[Scheduler] Error executing task {task['id']}: {e}")
+            # Update task with error
+            self.task_store.update_task(task['id'], {
+                "last_error": str(e),
+                "last_error_at": datetime.now().isoformat()
+            })