Coverage for agentos/background/supervisor.py: 35%
284 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 12:29 +0800
1"""
2Agent Supervision Tree — v1.11.0
4Resource-bounded agent hierarchy with monitoring, quotas, and auto-recovery.
5Inspired by OS process supervision (systemd, supervisord).
7Features:
8- Hierarchical supervision tree: parent monitors children
9- Resource quotas per agent (time, cost, tokens, concurrency)
10- Heartbeat-based health monitoring
11- Auto-kill for runaway agents exceeding quotas
12- Graceful degradation: kill child, preserve parent
13- Aggregate progress across the tree
14- Supervision events for external monitoring
16Usage:
17 sup = AgentSupervisor()
18 child = await sup.spawn(
19 name="data_analyzer",
20 loop_factory=lambda: AgentLoop(...),
21 quotas=AgentQuota(max_duration=600, max_cost_usd=2.0),
22 )
23 result = await sup.await_child(child.id, timeout=600)
24"""
26from __future__ import annotations
28import asyncio
29import time
30import uuid
31from collections.abc import Callable, Coroutine
32from dataclasses import dataclass, field
33from enum import StrEnum
34from typing import Any
36# ── Enums ────────────────────────────────────────────────────────
39class SupervisionEventType(StrEnum):
40 """Types of supervision events."""
42 SPAWNED = "spawned"
43 STARTED = "started"
44 HEARTBEAT = "heartbeat"
45 PROGRESS = "progress"
46 QUOTA_WARNING = "quota_warning" # Nearing quota limit
47 QUOTA_EXCEEDED = "quota_exceeded" # Quota hit, killed
48 HEARTBEAT_LOST = "heartbeat_lost" # Child unresponsive
49 COMPLETED = "completed"
50 FAILED = "failed"
51 CANCELLED = "cancelled"
52 KILLED = "killed" # Killed by supervisor
55@dataclass
56class SupervisionEvent:
57 """Event emitted by the supervision tree."""
59 type: SupervisionEventType
60 child_id: str
61 child_name: str
62 timestamp: float = field(default_factory=time.time)
63 data: dict[str, Any] = field(default_factory=dict)
64 message: str = ""
66 def to_dict(self) -> dict:
67 return {
68 "type": self.type.value,
69 "child_id": self.child_id,
70 "child_name": self.child_name,
71 "timestamp": self.timestamp,
72 "data": self.data,
73 "message": self.message,
74 }
77# ── Data Models ──────────────────────────────────────────────────
80@dataclass
81class AgentQuota:
82 """Resource limits for a supervised agent."""
84 max_duration_seconds: float = 3600.0 # Wall-clock time budget
85 max_cost_usd: float = 10.0 # Cost budget
86 max_tokens: int = 1_000_000 # Token budget
87 max_iterations: int = 500 # Max loop iterations
88 heartbeat_interval: float = 10.0 # Seconds between heartbeats
89 heartbeat_timeout: float = 30.0 # Seconds before considered dead
90 max_retries: int = 0 # Auto-restart on failure (0=no restart)
91 retry_delay: float = 5.0 # Delay before restart
92 cooldown_period: float = 60.0 # Rate limit on restarts
95@dataclass
96class AgentQuotaUsage:
97 """Current resource consumption of a supervised agent."""
99 elapsed_seconds: float = 0.0
100 cost_usd: float = 0.0
101 tokens_used: int = 0
102 iterations: int = 0
103 heartbeats_received: int = 0
104 last_heartbeat: float = 0.0
105 restarts: int = 0
106 last_restart: float = 0.0
108 @property
109 def duration_percent(self) -> float:
110 return 0.0 # Set externally with quota context
112 @property
113 def cost_percent(self) -> float:
114 return 0.0
116 def to_dict(self) -> dict:
117 return {
118 "elapsed_seconds": self.elapsed_seconds,
119 "cost_usd": self.cost_usd,
120 "tokens_used": self.tokens_used,
121 "iterations": self.iterations,
122 "heartbeats_received": self.heartbeats_received,
123 "last_heartbeat": self.last_heartbeat,
124 "restarts": self.restarts,
125 }
128@dataclass
129class SupervisedAgent:
130 """An agent running under supervision."""
132 id: str = field(default_factory=lambda: uuid.uuid4().hex[:12])
133 name: str = ""
134 quotas: AgentQuota = field(default_factory=AgentQuota)
135 usage: AgentQuotaUsage = field(default_factory=AgentQuotaUsage)
136 status: str = "pending" # pending/running/paused/completed/failed/killed
137 started_at: float = 0.0
138 finished_at: float = 0.0
139 result: Any = None
140 error: str = ""
141 metadata: dict[str, Any] = field(default_factory=dict)
143 # Internal
144 _task: asyncio.Task | None = field(default=None, repr=False)
145 _heartbeat_task: asyncio.Task | None = field(default=None, repr=False)
146 _pause_event: asyncio.Event | None = field(default=None, repr=False)
147 _kill_event: asyncio.Event | None = field(default=None, repr=False)
149 @property
150 def is_alive(self) -> bool:
151 return self.status in ("running", "paused")
153 @property
154 def duration_seconds(self) -> float:
155 end = self.finished_at or time.time()
156 return end - self.started_at if self.started_at else 0.0
158 def to_dict(self) -> dict:
159 return {
160 "id": self.id,
161 "name": self.name,
162 "quotas": {
163 "max_duration_seconds": self.quotas.max_duration_seconds,
164 "max_cost_usd": self.quotas.max_cost_usd,
165 "max_tokens": self.quotas.max_tokens,
166 "max_iterations": self.quotas.max_iterations,
167 },
168 "usage": self.usage.to_dict(),
169 "status": self.status,
170 "started_at": self.started_at,
171 "finished_at": self.finished_at,
172 "error": self.error,
173 "metadata": self.metadata,
174 }
177@dataclass
178class SupervisorConfig:
179 """Global supervisor configuration."""
181 max_children: int = 20
182 monitor_interval: float = 1.0 # Seconds between health checks
183 event_history_size: int = 500 # Max events to retain
184 auto_kill_on_quota: bool = True
185 log_events: bool = True
188# ── Callback types ───────────────────────────────────────────────
190EventCallback = Callable[[SupervisionEvent], None]
193# ── Agent Supervisor ─────────────────────────────────────────────
196class AgentSupervisor:
197 """
198 Hierarchical supervision tree for long-running multi-agent tasks.
200 Monitors children for:
201 - Resource quota violations (time, cost, tokens)
202 - Heartbeat loss (crash/hang detection)
203 - Progress stalls
205 Actions:
206 - Auto-kill runaway agents
207 - Graceful restart (optional)
208 - Event emission for external monitoring
209 """
211 def __init__(
212 self,
213 config: SupervisorConfig | None = None,
214 on_event: EventCallback | None = None,
215 ):
216 self.config = config or SupervisorConfig()
217 self._on_event = on_event
218 self._children: dict[str, SupervisedAgent] = {}
219 self._events: list[SupervisionEvent] = []
220 self._monitor_task: asyncio.Task | None = None
221 self._lock = asyncio.Lock()
223 # ── Public API ───────────────────────────────────────────────
225 async def spawn(
226 self,
227 name: str,
228 task: str = "",
229 loop_factory: Callable[[], Any] | None = None,
230 agent_loop: Any = None,
231 quotas: AgentQuota | None = None,
232 on_heartbeat: Callable[[], Coroutine] | None = None,
233 metadata: dict[str, Any] | None = None,
234 ) -> str:
235 """
236 Spawn a new child agent under supervision.
238 Returns child_id for monitoring/control.
239 """
240 async with self._lock:
241 if len(self._children) >= self.config.max_children:
242 raise RuntimeError(f"Max children ({self.config.max_children}) reached")
244 child = SupervisedAgent(
245 name=name,
246 quotas=quotas or AgentQuota(),
247 metadata=metadata or {},
248 )
249 child._pause_event = asyncio.Event()
250 child._pause_event.set() # Not paused
251 child._kill_event = asyncio.Event()
253 self._children[child.id] = child
254 self._emit(SupervisionEvent(SupervisionEventType.SPAWNED, child.id, name))
256 # Start monitoring in background
257 if not self._monitor_task or self._monitor_task.done():
258 self._monitor_task = asyncio.create_task(self._monitor_loop())
260 # Start heartbeat task
261 if on_heartbeat:
262 child._heartbeat_task = asyncio.create_task(self._heartbeat_loop(child, on_heartbeat))
264 # Start execution
265 child._task = asyncio.create_task(self._run_child(child, task, loop_factory, agent_loop))
267 return child.id
269 async def get_child(self, child_id: str) -> SupervisedAgent | None:
270 """Get child agent by ID."""
271 return self._children.get(child_id)
273 async def list_children(self) -> list[SupervisedAgent]:
274 """List all children with their status."""
275 return list(self._children.values())
277 async def await_child(self, child_id: str, timeout: float | None = None) -> Any:
278 """Wait for a child to complete, return its result."""
279 child = self._children.get(child_id)
280 if not child:
281 raise KeyError(f"Child {child_id} not found")
282 if not child._task:
283 raise RuntimeError(f"Child {child_id} has no running task")
285 try:
286 return await asyncio.wait_for(child._task, timeout=timeout)
287 except TimeoutError:
288 # Kill the child on timeout
289 await self.kill_child(child_id, reason="await timeout")
290 raise
292 async def pause_child(self, child_id: str) -> bool:
293 """Pause a running child."""
294 child = self._children.get(child_id)
295 if not child or not child.is_alive or not child._pause_event:
296 return False
297 child._pause_event.clear()
298 child.status = "paused"
299 return True
301 async def resume_child(self, child_id: str) -> bool:
302 """Resume a paused child."""
303 child = self._children.get(child_id)
304 if not child or child.status != "paused" or not child._pause_event:
305 return False
306 child._pause_event.set()
307 child.status = "running"
308 self._emit(
309 SupervisionEvent(SupervisionEventType.PROGRESS, child.id, child.name, message="Resumed")
310 )
311 return True
313 async def kill_child(self, child_id: str, reason: str = "") -> bool:
314 """Force-kill a child agent."""
315 child = self._children.get(child_id)
316 if not child or not child.is_alive:
317 return False
319 if child._kill_event:
320 child._kill_event.set()
321 if child._task and not child._task.done():
322 child._task.cancel()
324 child.status = "killed"
325 child.finished_at = time.time()
326 child.error = reason
328 self._emit(
329 SupervisionEvent(
330 SupervisionEventType.KILLED,
331 child.id,
332 child.name,
333 message=reason,
334 )
335 )
336 return True
338 async def aggregate_progress(self) -> dict[str, Any]:
339 """Aggregate progress across all children."""
340 children = list(self._children.values())
341 total = len(children)
342 completed = sum(1 for c in children if c.status == "completed")
343 failed = sum(1 for c in children if c.status in ("failed", "killed"))
344 running = sum(1 for c in children if c.status == "running")
346 # Aggregate costs
347 total_cost = sum(c.usage.cost_usd for c in children)
348 total_tokens = sum(c.usage.tokens_used for c in children)
350 return {
351 "total_children": total,
352 "completed": completed,
353 "failed": failed,
354 "running": running,
355 "total_cost_usd": total_cost,
356 "total_tokens": total_tokens,
357 "percent_complete": (completed / total * 100) if total > 0 else 0,
358 "children": [c.to_dict() for c in children],
359 }
361 async def shutdown(self, timeout: float = 10.0):
362 """Graceful shutdown: pause new spawns, wait for children, kill stragglers."""
363 # Kill all running children
364 for child_id in list(self._children.keys()):
365 await self.kill_child(child_id, reason="supervisor shutdown")
367 if self._monitor_task and not self._monitor_task.done():
368 self._monitor_task.cancel()
370 # Wait for children to die
371 deadline = time.time() + timeout
372 for child in self._children.values():
373 if child._task and not child._task.done():
374 remaining = max(0, deadline - time.time())
375 try:
376 await asyncio.wait_for(child._task, timeout=remaining)
377 except (TimeoutError, asyncio.CancelledError):
378 pass
380 # ── Internal ─────────────────────────────────────────────────
382 async def _run_child(
383 self,
384 child: SupervisedAgent,
385 task: str,
386 loop_factory: Callable[[], Any] | None,
387 agent_loop: Any,
388 ):
389 """Execute a child agent with full supervision."""
390 child.status = "running"
391 child.started_at = time.time()
392 self._emit(SupervisionEvent(SupervisionEventType.STARTED, child.id, child.name))
394 try:
395 if loop_factory:
396 loop = loop_factory()
397 elif agent_loop:
398 loop = agent_loop
399 else:
400 raise ValueError("Must provide loop_factory or agent_loop")
402 # Wrap loop to check for pause/kill signals
403 original_on_iteration = getattr(loop, "on_iteration", None)
405 async def supervised_on_iteration(iteration: int, tool_results: list):
406 # Check kill signal
407 if child._kill_event and child._kill_event.is_set():
408 raise asyncio.CancelledError("Killed by supervisor")
410 # Check pause signal
411 if child._pause_event:
412 await child._pause_event.wait()
414 # Update usage
415 child.usage.iterations = iteration
416 child.usage.elapsed_seconds = time.time() - child.started_at
418 # Quota checks
419 if child.usage.elapsed_seconds > child.quotas.max_duration_seconds:
420 if self.config.auto_kill_on_quota:
421 child._kill_event.set()
422 raise TimeoutError("Duration quota exceeded")
423 else:
424 self._emit(
425 SupervisionEvent(
426 SupervisionEventType.QUOTA_WARNING,
427 child.id,
428 child.name,
429 message=f"Duration at {child.usage.elapsed_seconds:.0f}s / {child.quotas.max_duration_seconds}s",
430 )
431 )
433 if original_on_iteration:
434 original_on_iteration(iteration, tool_results)
436 loop.on_iteration = supervised_on_iteration
438 # Run with retry logic
439 for attempt in range(child.quotas.max_retries + 1):
440 try:
441 result = await loop.run(task, session_id=child.id)
442 child.result = result.output if hasattr(result, "output") else result
443 child.usage.cost_usd = getattr(result, "cost_usd", 0.0)
444 child.usage.tokens_used = sum(getattr(result, "tokens_used", {}).values())
445 child.status = "completed"
446 child.finished_at = time.time()
447 self._emit(
448 SupervisionEvent(SupervisionEventType.COMPLETED, child.id, child.name)
449 )
450 return child.result
451 except (TimeoutError, asyncio.CancelledError):
452 raise
453 except Exception as e:
454 if attempt < child.quotas.max_retries:
455 # Check cooldown
456 since_last = time.time() - child.usage.last_restart
457 if child.usage.restarts > 0 and since_last < child.quotas.cooldown_period:
458 await asyncio.sleep(child.quotas.cooldown_period - since_last)
459 child.usage.restarts += 1
460 child.usage.last_restart = time.time()
461 await asyncio.sleep(child.quotas.retry_delay)
462 continue
463 child.status = "failed"
464 child.finished_at = time.time()
465 child.error = str(e)
466 self._emit(
467 SupervisionEvent(
468 SupervisionEventType.FAILED,
469 child.id,
470 child.name,
471 message=str(e),
472 )
473 )
474 raise
476 except asyncio.CancelledError:
477 child.status = "killed"
478 child.finished_at = time.time()
479 except Exception as e:
480 child.status = "failed"
481 child.finished_at = time.time()
482 child.error = str(e)
483 self._emit(
484 SupervisionEvent(
485 SupervisionEventType.FAILED,
486 child.id,
487 child.name,
488 message=str(e),
489 )
490 )
492 async def _heartbeat_loop(
493 self,
494 child: SupervisedAgent,
495 on_heartbeat: Callable[[], Coroutine],
496 ):
497 """Send periodic heartbeats and update usage."""
498 interval = child.quotas.heartbeat_interval
499 while child.is_alive:
500 try:
501 await asyncio.sleep(interval)
502 if not child.is_alive:
503 break
504 await on_heartbeat()
505 child.usage.heartbeats_received += 1
506 child.usage.last_heartbeat = time.time()
507 self._emit(
508 SupervisionEvent(
509 SupervisionEventType.HEARTBEAT,
510 child.id,
511 child.name,
512 data={"heartbeats": child.usage.heartbeats_received},
513 )
514 )
515 except asyncio.CancelledError:
516 break
517 except Exception:
518 pass
520 async def _monitor_loop(self):
521 """Monitor all children for health and quota violations."""
522 while True:
523 try:
524 await asyncio.sleep(self.config.monitor_interval)
525 now = time.time()
527 for child in list(self._children.values()):
528 if not child.is_alive:
529 continue
531 # Heartbeat timeout check
532 if (
533 child.quotas.heartbeat_timeout > 0
534 and child.usage.last_heartbeat > 0
535 and now - child.usage.last_heartbeat > child.quotas.heartbeat_timeout
536 ):
537 self._emit(
538 SupervisionEvent(
539 SupervisionEventType.HEARTBEAT_LOST,
540 child.id,
541 child.name,
542 message=f"No heartbeat for {now - child.usage.last_heartbeat:.0f}s",
543 )
544 )
545 await self.kill_child(child.id, reason="heartbeat lost")
547 # Duration check
548 elapsed = now - child.started_at if child.started_at else 0
549 if elapsed > child.quotas.max_duration_seconds * 0.9:
550 self._emit(
551 SupervisionEvent(
552 SupervisionEventType.QUOTA_WARNING,
553 child.id,
554 child.name,
555 message=f"90% duration used: {elapsed:.0f}s / {child.quotas.max_duration_seconds}s",
556 )
557 )
559 except asyncio.CancelledError:
560 break
561 except Exception:
562 pass
564 def _emit(self, event: SupervisionEvent):
565 """Emit a supervision event."""
566 if self.config.log_events:
567 self._events.append(event)
568 if len(self._events) > self.config.event_history_size:
569 self._events = self._events[-self.config.event_history_size :]
571 if self._on_event:
572 try:
573 self._on_event(event)
574 except Exception:
575 pass