448 lines
12 KiB
Python
448 lines
12 KiB
Python
"""
|
|
Tool Scheduler
|
|
|
|
Scheduled tool execution system with support for one-time, interval, and cron-based scheduling.
|
|
"""
|
|
|
|
import asyncio
|
|
import uuid
|
|
from datetime import datetime, timedelta
|
|
from enum import Enum
|
|
from typing import Any, Callable, Dict, Optional
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
|
|
class ScheduleType(str, Enum):
|
|
"""Schedule type enumeration"""
|
|
|
|
ONCE = "once" # Single execution
|
|
INTERVAL = "interval" # Fixed interval
|
|
CRON = "cron" # Cron expression
|
|
|
|
|
|
class ScheduledTask(BaseModel):
|
|
"""Scheduled task model"""
|
|
|
|
id: str = Field(default_factory=lambda: str(uuid.uuid4())[:8])
|
|
name: str
|
|
schedule_type: ScheduleType
|
|
schedule_value: str # datetime string / interval seconds / cron expression
|
|
tool_name: str
|
|
parameters: Dict[str, Any] = Field(default_factory=dict)
|
|
enabled: bool = True
|
|
last_run: Optional[datetime] = None
|
|
next_run: Optional[datetime] = None
|
|
run_count: int = 0
|
|
created_at: datetime = Field(default_factory=datetime.utcnow)
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"id": self.id,
|
|
"name": self.name,
|
|
"type": self.schedule_type.value,
|
|
"schedule_value": self.schedule_value,
|
|
"tool_name": self.tool_name,
|
|
"enabled": self.enabled,
|
|
"last_run": self.last_run.isoformat() if self.last_run else None,
|
|
"next_run": self.next_run.isoformat() if self.next_run else None,
|
|
"run_count": self.run_count,
|
|
}
|
|
|
|
|
|
class ToolScheduler:
|
|
"""Tool scheduler for automated execution"""
|
|
|
|
def __init__(self):
|
|
self._tasks: Dict[str, ScheduledTask] = {}
|
|
self._running: bool = False
|
|
self._loop_task: Optional[asyncio.Task] = None
|
|
self._executor: Optional[Callable] = None
|
|
|
|
def set_executor(self, executor: Callable) -> None:
|
|
"""Set the tool executor function
|
|
|
|
Args:
|
|
executor: Async callable that takes (tool_name, command, parameters)
|
|
"""
|
|
self._executor = executor
|
|
|
|
async def schedule(
|
|
self,
|
|
name: str,
|
|
schedule_type: ScheduleType,
|
|
schedule_value: str,
|
|
tool_name: str,
|
|
parameters: Optional[Dict[str, Any]] = None,
|
|
) -> str:
|
|
"""Create a scheduled task
|
|
|
|
Args:
|
|
name: Task name
|
|
schedule_type: Type of schedule
|
|
schedule_value: Schedule value (datetime/interval/cron)
|
|
tool_name: Tool to execute
|
|
parameters: Tool parameters
|
|
|
|
Returns:
|
|
Task ID
|
|
"""
|
|
task = ScheduledTask(
|
|
name=name,
|
|
schedule_type=schedule_type,
|
|
schedule_value=schedule_value,
|
|
tool_name=tool_name,
|
|
parameters=parameters or {},
|
|
)
|
|
|
|
task.next_run = self._calculate_next_run(task)
|
|
self._tasks[task.id] = task
|
|
|
|
# Start scheduler loop if not running
|
|
if not self._running:
|
|
await self.start()
|
|
|
|
return task.id
|
|
|
|
async def schedule_once(
|
|
self,
|
|
name: str,
|
|
run_at: str, # ISO datetime string
|
|
tool_name: str,
|
|
parameters: Optional[Dict[str, Any]] = None,
|
|
) -> str:
|
|
"""Schedule a one-time task
|
|
|
|
Args:
|
|
name: Task name
|
|
run_at: ISO datetime string
|
|
tool_name: Tool to execute
|
|
parameters: Tool parameters
|
|
|
|
Returns:
|
|
Task ID
|
|
"""
|
|
return await self.schedule(
|
|
name=name,
|
|
schedule_type=ScheduleType.ONCE,
|
|
schedule_value=run_at,
|
|
tool_name=tool_name,
|
|
parameters=parameters,
|
|
)
|
|
|
|
async def schedule_interval(
|
|
self,
|
|
name: str,
|
|
interval_seconds: int,
|
|
tool_name: str,
|
|
parameters: Optional[Dict[str, Any]] = None,
|
|
) -> str:
|
|
"""Schedule an interval-based recurring task
|
|
|
|
Args:
|
|
name: Task name
|
|
interval_seconds: Interval in seconds
|
|
tool_name: Tool to execute
|
|
parameters: Tool parameters
|
|
|
|
Returns:
|
|
Task ID
|
|
"""
|
|
return await self.schedule(
|
|
name=name,
|
|
schedule_type=ScheduleType.INTERVAL,
|
|
schedule_value=str(interval_seconds),
|
|
tool_name=tool_name,
|
|
parameters=parameters,
|
|
)
|
|
|
|
async def schedule_cron(
|
|
self,
|
|
name: str,
|
|
cron_expression: str,
|
|
tool_name: str,
|
|
parameters: Optional[Dict[str, Any]] = None,
|
|
) -> str:
|
|
"""Schedule a cron-based recurring task
|
|
|
|
Args:
|
|
name: Task name
|
|
cron_expression: Cron expression (5 fields: min hour day month weekday)
|
|
tool_name: Tool to execute
|
|
parameters: Tool parameters
|
|
|
|
Returns:
|
|
Task ID
|
|
"""
|
|
return await self.schedule(
|
|
name=name,
|
|
schedule_type=ScheduleType.CRON,
|
|
schedule_value=cron_expression,
|
|
tool_name=tool_name,
|
|
parameters=parameters,
|
|
)
|
|
|
|
def _calculate_next_run(self, task: ScheduledTask) -> Optional[datetime]:
|
|
"""Calculate the next run time for a task
|
|
|
|
Args:
|
|
task: The scheduled task
|
|
|
|
Returns:
|
|
Next run datetime or None if schedule type is invalid
|
|
"""
|
|
now = datetime.utcnow()
|
|
|
|
if task.schedule_type == ScheduleType.ONCE:
|
|
try:
|
|
return datetime.fromisoformat(task.schedule_value)
|
|
except ValueError:
|
|
return None
|
|
|
|
elif task.schedule_type == ScheduleType.INTERVAL:
|
|
try:
|
|
seconds = int(task.schedule_value)
|
|
return now + timedelta(seconds=seconds)
|
|
except ValueError:
|
|
return None
|
|
|
|
elif task.schedule_type == ScheduleType.CRON:
|
|
return self._parse_cron_next(task.schedule_value, now)
|
|
|
|
return None
|
|
|
|
def _parse_cron_next(self, cron_expr: str, base_time: datetime) -> datetime:
|
|
"""Parse cron expression and calculate next run
|
|
|
|
Args:
|
|
cron_expr: Cron expression (min hour day month weekday)
|
|
base_time: Reference datetime
|
|
|
|
Returns:
|
|
Next matching datetime
|
|
"""
|
|
try:
|
|
parts = cron_expr.split()
|
|
if len(parts) != 5:
|
|
return base_time + timedelta(hours=1)
|
|
|
|
minute, hour, day, month, weekday = parts
|
|
|
|
# Simple next-run calculation
|
|
# For production, use croniter library
|
|
next_time = base_time.replace(second=0, microsecond=0) + timedelta(minutes=1)
|
|
|
|
# Basic validation and advancement
|
|
max_iterations = 366 * 24 * 60 # 1 year of minutes max
|
|
for _ in range(max_iterations):
|
|
if self._cron_matches(next_time, minute, hour, day, month, weekday):
|
|
return next_time
|
|
next_time += timedelta(minutes=1)
|
|
|
|
return next_time
|
|
|
|
except Exception:
|
|
return base_time + timedelta(hours=1)
|
|
|
|
def _cron_matches(
|
|
self,
|
|
dt: datetime,
|
|
minute: str,
|
|
hour: str,
|
|
day: str,
|
|
month: str,
|
|
weekday: str,
|
|
) -> bool:
|
|
"""Check if datetime matches cron fields
|
|
|
|
Args:
|
|
dt: Datetime to check
|
|
minute: Minute field
|
|
hour: Hour field
|
|
day: Day of month field
|
|
month: Month field
|
|
weekday: Day of week field
|
|
|
|
Returns:
|
|
True if matches
|
|
"""
|
|
return (
|
|
self._cron_field_matches(dt.minute, minute)
|
|
and self._cron_field_matches(dt.hour, hour)
|
|
and self._cron_field_matches(dt.day, day)
|
|
and self._cron_field_matches(dt.month, month)
|
|
and self._cron_field_matches(dt.weekday(), weekday)
|
|
)
|
|
|
|
def _cron_field_matches(self, value: int, field: str) -> bool:
|
|
"""Check if a value matches a cron field
|
|
|
|
Args:
|
|
value: The actual value
|
|
field: Cron field (number, *, */n, n,m, n-m)
|
|
|
|
Returns:
|
|
True if matches
|
|
"""
|
|
if field == "*":
|
|
return True
|
|
|
|
if field.startswith("*/"):
|
|
try:
|
|
step = int(field[2:])
|
|
return value % step == 0
|
|
except ValueError:
|
|
return False
|
|
|
|
if "," in field:
|
|
return str(value) in field.split(",")
|
|
|
|
if "-" in field:
|
|
try:
|
|
start, end = field.split("-")
|
|
return start <= str(value) <= end
|
|
except ValueError:
|
|
return False
|
|
|
|
try:
|
|
return int(field) == value
|
|
except ValueError:
|
|
return False
|
|
|
|
async def start(self) -> None:
|
|
"""Start the scheduler loop"""
|
|
if self._running:
|
|
return
|
|
self._running = True
|
|
self._loop_task = asyncio.create_task(self._run_loop())
|
|
|
|
async def stop(self) -> None:
|
|
"""Stop the scheduler loop"""
|
|
self._running = False
|
|
if self._loop_task:
|
|
self._loop_task.cancel()
|
|
try:
|
|
await self._loop_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
self._loop_task = None
|
|
|
|
async def _run_loop(self) -> None:
|
|
"""Main scheduler loop"""
|
|
while self._running:
|
|
now = datetime.utcnow()
|
|
|
|
for task in list(self._tasks.values()):
|
|
if not task.enabled:
|
|
continue
|
|
|
|
if task.next_run and task.next_run <= now:
|
|
await self._execute_task(task)
|
|
|
|
await asyncio.sleep(1) # Check every second
|
|
|
|
async def _execute_task(self, task: ScheduledTask) -> None:
|
|
"""Execute a scheduled task
|
|
|
|
Args:
|
|
task: The task to execute
|
|
"""
|
|
if self._executor is None:
|
|
return
|
|
|
|
try:
|
|
await self._executor(
|
|
tool_name=task.tool_name,
|
|
command=task.parameters.get("command", ""),
|
|
parameters=task.parameters,
|
|
)
|
|
except Exception:
|
|
pass # Log error in production
|
|
|
|
# Update task state
|
|
task.last_run = datetime.utcnow()
|
|
task.run_count += 1
|
|
|
|
# Calculate next run or disable one-time tasks
|
|
if task.schedule_type != ScheduleType.ONCE:
|
|
task.next_run = self._calculate_next_run(task)
|
|
else:
|
|
task.enabled = False
|
|
|
|
async def cancel(self, task_id: str) -> bool:
|
|
"""Cancel a scheduled task
|
|
|
|
Args:
|
|
task_id: Task ID
|
|
|
|
Returns:
|
|
True if cancelled
|
|
"""
|
|
if task_id in self._tasks:
|
|
del self._tasks[task_id]
|
|
return True
|
|
return False
|
|
|
|
async def enable(self, task_id: str) -> bool:
|
|
"""Enable a scheduled task
|
|
|
|
Args:
|
|
task_id: Task ID
|
|
|
|
Returns:
|
|
True if enabled
|
|
"""
|
|
task = self._tasks.get(task_id)
|
|
if task:
|
|
task.enabled = True
|
|
task.next_run = self._calculate_next_run(task)
|
|
return True
|
|
return False
|
|
|
|
async def disable(self, task_id: str) -> bool:
|
|
"""Disable a scheduled task
|
|
|
|
Args:
|
|
task_id: Task ID
|
|
|
|
Returns:
|
|
True if disabled
|
|
"""
|
|
task = self._tasks.get(task_id)
|
|
if task:
|
|
task.enabled = False
|
|
return True
|
|
return False
|
|
|
|
async def list_tasks(self) -> list:
|
|
"""List all scheduled tasks
|
|
|
|
Returns:
|
|
List of task info dicts
|
|
"""
|
|
return [task.to_dict() for task in self._tasks.values()]
|
|
|
|
async def get_task(self, task_id: str) -> Optional[dict]:
|
|
"""Get a specific task
|
|
|
|
Args:
|
|
task_id: Task ID
|
|
|
|
Returns:
|
|
Task dict or None
|
|
"""
|
|
task = self._tasks.get(task_id)
|
|
return task.to_dict() if task else None
|
|
|
|
|
|
# Global scheduler instance
|
|
_scheduler: Optional[ToolScheduler] = None
|
|
|
|
|
|
def get_scheduler() -> ToolScheduler:
|
|
"""Get the global tool scheduler instance"""
|
|
global _scheduler
|
|
if _scheduler is None:
|
|
_scheduler = ToolScheduler()
|
|
return _scheduler
|