tiphys.scheduler.service

  1"""
  2Scheduler service for executing background tasks.
  3
  4Manages the background loop, task execution, and schedule calculation.
  5"""
  6
  7import asyncio
  8import contextlib
  9from datetime import datetime, timedelta
 10from typing import TYPE_CHECKING
 11from uuid import uuid4
 12
 13import structlog
 14from croniter import croniter  # type: ignore[import-untyped]
 15
 16from tiphys.scheduler.models import ScheduledTask, ScheduleType, TaskRunResult, TaskStatus
 17from tiphys.scheduler.store import TaskStore
 18
 19MAX_STEPS = 10
 20
 21if TYPE_CHECKING:
 22    from tiphys.context import AppContext
 23
 24logger = structlog.get_logger(__name__)
 25
 26
 27class TaskScheduler:
 28    """
 29    Manages background execution of scheduled agent tasks.
 30    """
 31
 32    def __init__(self, context: "AppContext", store: TaskStore | None = None):
 33        """
 34        Initialize the scheduler.
 35
 36        Args:
 37            context: The application context.
 38            store: Optional task store override.
 39        """
 40        self.context = context
 41        self.store = store or TaskStore()
 42        self._running = False
 43        self._loop_task: asyncio.Task[None] | None = None
 44
 45    async def start(self) -> None:
 46        """Start the background scheduler loop."""
 47        if self._running:
 48            return
 49
 50        self._running = True
 51        self._loop_task = asyncio.create_task(self._run_loop())
 52        await logger.ainfo("Task scheduler started")
 53
 54    async def stop(self) -> None:
 55        """Stop the background scheduler loop."""
 56        self._running = False
 57        if self._loop_task:
 58            self._loop_task.cancel()
 59            with contextlib.suppress(asyncio.CancelledError):
 60                await self._loop_task
 61            self._loop_task = None
 62        await logger.ainfo("Task scheduler stopped")
 63
 64    async def create_task(
 65        self,
 66        name: str,
 67        group: str,
 68        prompt: str,
 69        schedule_type: ScheduleType,
 70        schedule_value: str,
 71        agent_id: str = "default",
 72    ) -> str:
 73        """
 74        Create and schedule a new task.
 75
 76        Returns:
 77            The ID of the created task.
 78        """
 79        task_id = str(uuid4())
 80        task = ScheduledTask(
 81            id=task_id,
 82            name=name,
 83            group=group,
 84            prompt=prompt,
 85            schedule_type=schedule_type,
 86            schedule_value=schedule_value,
 87            agent_id=agent_id,
 88        )
 89
 90        # Calculate first run
 91        task.next_run = self._calculate_next_run(task)
 92
 93        self.store.save(task)
 94        await logger.ainfo("Task created", task_id=task_id, name=name, next_run=task.next_run)
 95        return task_id
 96
 97    async def pause_task(self, task_id: str) -> bool:
 98        """Pause a scheduled task."""
 99        task = self.store.load(task_id)
100        if not task:
101            return False
102
103        task.status = TaskStatus.PAUSED
104        self.store.save(task)
105        await logger.ainfo("Task paused", task_id=task_id)
106        return True
107
108    async def resume_task(self, task_id: str) -> bool:
109        """Resume a paused task."""
110        task = self.store.load(task_id)
111        if not task:
112            return False
113
114        task.status = TaskStatus.ACTIVE
115        task.next_run = self._calculate_next_run(task)
116        self.store.save(task)
117        await logger.ainfo("Task resumed", task_id=task_id, next_run=task.next_run)
118        return True
119
120    def _calculate_next_run(self, task: ScheduledTask) -> datetime | None:
121        """Calculate the next run time for a task."""
122        now = datetime.utcnow()
123
124        try:
125            if task.schedule_type == ScheduleType.CRON:
126                iter = croniter(task.schedule_value, now)
127                from typing import cast
128
129                return cast(datetime, iter.get_next(datetime))
130
131            elif task.schedule_type == ScheduleType.INTERVAL:
132                seconds = int(task.schedule_value)
133                return now + timedelta(seconds=seconds)
134
135            elif task.schedule_type == ScheduleType.ONCE:
136                dt = datetime.fromisoformat(task.schedule_value)
137                return dt if dt > now else None
138        except Exception as e:
139            logger.error("Failed to calculate next run", task_id=task.id, error=str(e))
140            return None
141
142        return None
143
144    async def _run_loop(self) -> None:
145        """Background loop to check for due tasks."""
146        while self._running:
147            try:
148                due_tasks = self.store.get_due_tasks()
149                for task in due_tasks:
150                    # Run in background to not block the loop
151                    asyncio.create_task(self._execute_task(task))
152
153                # Check every 30 seconds
154                await asyncio.sleep(30)
155            except asyncio.CancelledError:
156                break
157            except Exception as e:
158                logger.error("Error in scheduler loop", error=str(e))
159                await asyncio.sleep(60)
160
161    async def _execute_task(self, task: ScheduledTask) -> None:
162        """Execute a single due task."""
163        run_id = str(uuid4())
164        start_time = datetime.utcnow()
165
166        await logger.ainfo("Executing task", task_id=task.id, name=task.name, run_id=run_id)
167
168        # Mark as running/updating next run so it's not picked up again immediately
169        task.last_run = start_time
170        task.run_count += 1
171        task.next_run = self._calculate_next_run(task)
172
173        # If it was a 'once' task, it might not have a next run
174        if task.schedule_type == ScheduleType.ONCE:
175            task.status = TaskStatus.COMPLETED
176
177        self.store.save(task)
178
179        try:
180            from tiphys.agent.runner import AgentRunner
181
182            # Build a session key for the task if not provided
183            session_key = task.session_key or f"task:{task.id}"
184
185            # Use AppContext to create runner
186            runner = AgentRunner(context=self.context)
187
188            # Execute agent
189            output = await runner.run(
190                session_key=session_key, user_message=task.prompt, max_steps=10
191            )
192
193            duration = int((datetime.utcnow() - start_time).total_seconds() * 1000)
194            TaskRunResult(
195                task_id=task.id, run_id=run_id, success=True, output=output, duration_ms=duration
196            )
197
198            await logger.ainfo("Task succeeded", task_id=task.id, duration_ms=duration)
199            # In a real system, we'd log the result to a history store
200
201        except Exception as e:
202            duration = int((datetime.utcnow() - start_time).total_seconds() * 1000)
203            logger.error("Task failed", task_id=task.id, error=str(e))
204
205            # Optionally mark task as failed if it keeps failing
206            # For now, just keep it active for the next run