feat(tools): Phase T.1-T.4 complete - manifest system, registry, implementations, runtime, collaboration, scheduler
This commit is contained in:
447
backend/app/tools/scheduler.py
Normal file
447
backend/app/tools/scheduler.py
Normal file
@@ -0,0 +1,447 @@
|
||||
"""
|
||||
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
|
||||
Reference in New Issue
Block a user