Coverage for agentos/tools/task_scheduler.py: 0%

1""" 

2Lightweight Task Scheduler & Job Queue for AgentOS. 

3 

4PriorityTaskQueue — bounded priority queue with deadline-aware scheduling. 

5TaskScheduler — interval/delay/cron-like scheduling with persistence hooks. 

6WorkerPool — simple thread-pool backed by the task queue. 

7""" 

8 

9import heapq 

10import threading 

11import time 

12import uuid 

13from collections.abc import Callable 

14from dataclasses import dataclass, field 

15from enum import Enum, auto 

16from typing import Any 

17 

18# ============================================================================ 

19# Task & Result Types 

20# ============================================================================ 

21 

22 

23class TaskStatus(Enum): 

24 PENDING = auto() 

25 RUNNING = auto() 

26 COMPLETED = auto() 

27 FAILED = auto() 

28 EXPIRED = auto() 

29 CANCELLED = auto() 

30 

31 

32@dataclass(order=True) 

33class _PrioritizedTask: 

34 """Internal heap item. Lower priority value = higher urgency.""" 

35 

36 priority: int 

37 deadline: float 

38 seq: int # tiebreaker for stable ordering 

39 task: "Task" = field(compare=False) 

40 

41 

42@dataclass 

43class Task: 

44 """A schedulable task with metadata.""" 

45 

46 task_id: str 

47 func: Callable 

48 args: tuple = () 

49 kwargs: dict[str, Any] = field(default_factory=dict) 

50 priority: int = 0 

51 deadline: float | None = None 

52 status: TaskStatus = TaskStatus.PENDING 

53 result: Any = None 

54 error: str | None = None 

55 created_at: float = field(default_factory=time.monotonic) 

56 started_at: float | None = None 

57 completed_at: float | None = None 

58 

59 def run(self): 

60 self.started_at = time.monotonic() 

61 self.status = TaskStatus.RUNNING 

62 try: 

63 self.result = self.func(*self.args, **self.kwargs) 

64 self.status = TaskStatus.COMPLETED 

65 except Exception as e: 

66 self.error = str(e) 

67 self.status = TaskStatus.FAILED 

68 finally: 

69 self.completed_at = time.monotonic() 

70 

71 @property 

72 def elapsed(self) -> float | None: 

73 if self.started_at is None: 

74 return None 

75 return (self.completed_at or time.monotonic()) - self.started_at 

76 

77 

78# ============================================================================ 

79# PriorityTaskQueue 

80# ============================================================================ 

81 

82 

83class PriorityTaskQueue: 

84 """Thread-safe bounded priority queue for tasks. 

85 

86 Lower priority value = executed first. Deadlines auto-expire stale tasks. 

87 """ 

88 

89 def __init__(self, max_size: int = 1000): 

90 self._max_size = max_size 

91 self._heap: list[_PrioritizedTask] = [] 

92 self._lookup: dict[str, _PrioritizedTask] = {} 

93 self._lock = threading.RLock() 

94 self._seq: int = 0 

95 self._total_enqueued: int = 0 

96 self._total_dequeued: int = 0 

97 self._total_expired: int = 0 

98 

99 def enqueue(self, task: Task) -> bool: 

100 with self._lock: 

101 if len(self._heap) >= self._max_size: 

102 return False 

103 self._seq += 1 

104 pt = _PrioritizedTask( 

105 priority=task.priority, 

106 deadline=task.deadline or float("inf"), 

107 seq=self._seq, 

108 task=task, 

109 ) 

110 heapq.heappush(self._heap, pt) 

111 self._lookup[task.task_id] = pt 

112 self._total_enqueued += 1 

113 return True 

114 

115 def dequeue(self) -> Task | None: 

116 with self._lock: 

117 self._clean_expired() 

118 while self._heap: 

119 pt = heapq.heappop(self._heap) 

120 task = pt.task 

121 if task.status == TaskStatus.CANCELLED: 

122 self._lookup.pop(task.task_id, None) 

123 continue 

124 if task.deadline and time.monotonic() > task.deadline: 

125 task.status = TaskStatus.EXPIRED 

126 self._lookup.pop(task.task_id, None) 

127 self._total_expired += 1 

128 continue 

129 self._lookup.pop(task.task_id, None) 

130 self._total_dequeued += 1 

131 return task 

132 return None 

133 

134 def cancel(self, task_id: str) -> bool: 

135 with self._lock: 

136 pt = self._lookup.pop(task_id, None) 

137 if pt: 

138 pt.task.status = TaskStatus.CANCELLED 

139 return True 

140 return False 

141 

142 def _clean_expired(self): 

143 """Remove cancelled tasks from heap top.""" 

144 now = time.monotonic() 

145 while self._heap: 

146 pt = self._heap[0] 

147 if pt.task.status == TaskStatus.CANCELLED: 

148 heapq.heappop(self._heap) 

149 self._lookup.pop(pt.task.task_id, None) 

150 elif pt.deadline != float("inf") and now > pt.deadline: 

151 pt.task.status = TaskStatus.EXPIRED 

152 heapq.heappop(self._heap) 

153 self._lookup.pop(pt.task.task_id, None) 

154 self._total_expired += 1 

155 else: 

156 break 

157 

158 @property 

159 def size(self) -> int: 

160 with self._lock: 

161 return len(self._heap) 

162 

163 @property 

164 def stats(self) -> dict[str, Any]: 

165 with self._lock: 

166 return { 

167 "size": len(self._heap), 

168 "max_size": self._max_size, 

169 "total_enqueued": self._total_enqueued, 

170 "total_dequeued": self._total_dequeued, 

171 "total_expired": self._total_expired, 

172 } 

173 

174 

175# ============================================================================ 

176# TaskScheduler 

177# ============================================================================ 

178 

179 

180class TaskScheduler: 

181 """Schedule tasks at intervals or after delays. Runs in a background thread.""" 

182 

183 def __init__(self): 

184 self._queue = PriorityTaskQueue() 

185 self._running = False 

186 self._thread: threading.Thread | None = None 

187 self._lock = threading.Lock() 

188 self._scheduled_count: int = 0 

189 self._executed_count: int = 0 

190 

191 def submit(self, task: Task) -> bool: 

192 """Submit a task for immediate execution.""" 

193 ok = self._queue.enqueue(task) 

194 if ok: 

195 self._scheduled_count += 1 

196 return ok 

197 

198 def schedule_after(self, func: Callable, delay: float, *args, **kwargs) -> Task: 

199 """Schedule a task to run after delay seconds.""" 

200 task = Task( 

201 task_id=str(uuid.uuid4()), 

202 func=func, 

203 args=args, 

204 kwargs=kwargs, 

205 deadline=time.monotonic() + delay, 

206 priority=int(delay * 1000), # sooner = lower priority 

207 ) 

208 self.submit(task) 

209 return task 

210 

211 def schedule_at_interval(self, func: Callable, interval: float, *args, **kwargs) -> None: 

212 """Repeatedly schedule a task at fixed intervals. Uses a daemon thread.""" 

213 

214 def _loop(): 

215 while self._running: 

216 t = Task( 

217 task_id=str(uuid.uuid4()), 

218 func=func, 

219 args=args, 

220 kwargs=kwargs, 

221 ) 

222 self.submit(t) 

223 self._executed_count += 1 

224 time.sleep(interval) 

225 

226 t = threading.Thread(target=_loop, daemon=True) 

227 t.start() 

228 

229 def start(self) -> None: 

230 self._running = True 

231 

232 def stop(self) -> None: 

233 self._running = False 

234 

235 def run_once(self) -> Task | None: 

236 """Pull and run one task. Returns the executed task or None.""" 

237 task = self._queue.dequeue() 

238 if task: 

239 task.run() 

240 self._executed_count += 1 

241 return task 

242 return None 

243 

244 def run_loop(self, max_tasks: int = 0) -> int: 

245 """Run tasks until queue empty or max_tasks reached. Returns count executed.""" 

246 count = 0 

247 self._running = True 

248 while self._running: 

249 task = self._queue.dequeue() 

250 if task is None: 

251 break 

252 task.run() 

253 count += 1 

254 if 0 < max_tasks <= count: 

255 break 

256 self._running = False 

257 self._executed_count += count 

258 return count 

259 

260 @property 

261 def pending(self) -> int: 

262 return self._queue.size 

263 

264 @property 

265 def stats(self) -> dict[str, Any]: 

266 return { 

267 **self._queue.stats, 

268 "scheduled_count": self._scheduled_count, 

269 "executed_count": self._executed_count, 

270 "pending": self.pending, 

271 } 

272 

273 

274# ============================================================================ 

275# WorkerPool 

276# ============================================================================ 

277 

278 

279class WorkerPool: 

280 """Simple thread pool pulling from a PriorityTaskQueue.""" 

281 

282 def __init__(self, num_workers: int = 4, max_queue_size: int = 1000): 

283 self._num_workers = num_workers 

284 self._queue = PriorityTaskQueue(max_size=max_queue_size) 

285 self._running = False 

286 self._workers: list[threading.Thread] = [] 

287 self._lock = threading.Lock() 

288 

289 def submit(self, func: Callable, *args, **kwargs) -> Task: 

290 task = Task( 

291 task_id=str(uuid.uuid4()), 

292 func=func, 

293 args=args, 

294 kwargs=kwargs, 

295 ) 

296 ok = self._queue.enqueue(task) 

297 if not ok: 

298 raise RuntimeError("WorkerPool queue is full") 

299 return task 

300 

301 def start(self) -> None: 

302 self._running = True 

303 for _ in range(self._num_workers): 

304 t = threading.Thread(target=self._worker_loop, daemon=True) 

305 t.start() 

306 self._workers.append(t) 

307 

308 def stop(self, wait: bool = True) -> None: 

309 self._running = False 

310 if wait: 

311 for t in self._workers: 

312 t.join(timeout=5.0) 

313 

314 def _worker_loop(self) -> None: 

315 while self._running: 

316 task = self._queue.dequeue() 

317 if task is None: 

318 time.sleep(0.01) 

319 continue 

320 task.run() 

321 

322 @property 

323 def pending(self) -> int: 

324 return self._queue.size 

325 

326 @property 

327 def stats(self) -> dict[str, Any]: 

328 return { 

329 **self._queue.stats, 

330 "workers": self._num_workers, 

331 "running": self._running, 

332 }