File size: 7,668 Bytes
92c4ae6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 | """
Agent Task Registry
Manages asyncio tasks for agent execution with proper cancellation support.
Enables tracking and cancellation of running agents.
"""
import asyncio
import logging
from datetime import datetime
from typing import Dict, Optional, Set
from dataclasses import dataclass, field
logger = logging.getLogger(__name__)
@dataclass
class AgentTask:
"""Represents a running agent task"""
task_id: str
agent_id: str
agent_run_id: str
task: asyncio.Task
user_id: str
started_at: datetime = field(default_factory=datetime.now)
status: str = "running" # running, cancelled, completed, failed
def cancel(self) -> bool:
"""Cancel the underlying asyncio task"""
if not self.task.done():
self.task.cancel()
self.status = "cancelled"
return True
return False
class AgentTaskRegistry:
"""
Global registry for managing agent tasks.
Provides:
- Task registration for running agents
- Task cancellation by agent_id or task_id
- Task status tracking
- Cleanup of completed tasks
"""
_instance = None
def __new__(cls):
if cls._instance is None:
cls._instance = super().__new__(cls)
cls._instance._initialized = False
return cls._instance
def __init__(self):
if self._initialized:
return
self._tasks: Dict[str, AgentTask] = {} # task_id -> AgentTask
self._agent_tasks: Dict[str, Set[str]] = {} # agent_id -> set of task_ids
self._run_tasks: Dict[str, str] = {} # agent_run_id -> task_id
self._initialized = True
logger.info("AgentTaskRegistry initialized")
def register_task(
self,
task_id: str,
agent_id: str,
agent_run_id: str,
task: asyncio.Task,
user_id: str
) -> None:
"""Register a new agent task"""
agent_task = AgentTask(
task_id=task_id,
agent_id=agent_id,
agent_run_id=agent_run_id,
task=task,
user_id=user_id
)
self._tasks[task_id] = agent_task
# Track by agent_id
if agent_id not in self._agent_tasks:
self._agent_tasks[agent_id] = set()
self._agent_tasks[agent_id].add(task_id)
# Track by agent_run_id
self._run_tasks[agent_run_id] = task_id
logger.info(f"Registered task {task_id} for agent {agent_id}, run {agent_run_id}")
def unregister_task(self, task_id: str) -> None:
"""Unregister a completed task"""
if task_id not in self._tasks:
return
agent_task = self._tasks[task_id]
# Remove from agent_tasks
if agent_task.agent_id in self._agent_tasks:
self._agent_tasks[agent_task.agent_id].discard(task_id)
if not self._agent_tasks[agent_task.agent_id]:
del self._agent_tasks[agent_task.agent_id]
# Remove from run_tasks
if agent_task.agent_run_id in self._run_tasks:
del self._run_tasks[agent_task.agent_run_id]
# Remove from tasks
del self._tasks[task_id]
logger.info(f"Unregistered task {task_id}")
async def cancel_task(self, task_id: str) -> bool:
"""
Cancel a task by task_id and wait for cancellation to complete.
This method now properly waits for the task to handle the cancellation
signal before unregistering it, preventing race conditions in tests.
"""
if task_id not in self._tasks:
logger.warning(f"Task {task_id} not found in registry")
return False
agent_task = self._tasks[task_id]
success = agent_task.cancel()
if success:
logger.info(f"Cancelled task {task_id}")
# Wait for task to actually be cancelled (handles async propagation)
# This prevents race conditions where task isn't fully cancelled when unregistered
try:
await asyncio.wait_for(agent_task.task, timeout=5.0)
except (asyncio.CancelledError, asyncio.TimeoutError):
# CancelledError is expected when task handles cancellation
# TimeoutError means task didn't respond to cancellation within 5s
pass
self.unregister_task(task_id)
return success
async def cancel_agent_tasks(self, agent_id: str) -> int:
"""Cancel all running tasks for an agent"""
if agent_id not in self._agent_tasks:
logger.warning(f"No tasks found for agent {agent_id}")
return 0
task_ids = list(self._agent_tasks[agent_id])
cancelled_count = 0
for task_id in task_ids:
if await self.cancel_task(task_id):
cancelled_count += 1
logger.info(f"Cancelled {cancelled_count} tasks for agent {agent_id}")
return cancelled_count
async def cancel_agent_run(self, agent_run_id: str) -> bool:
"""Cancel a specific agent run"""
if agent_run_id not in self._run_tasks:
logger.warning(f"Agent run {agent_run_id} not found in registry")
return False
task_id = self._run_tasks[agent_run_id]
return await self.cancel_task(task_id)
def get_task(self, task_id: str) -> Optional[AgentTask]:
"""Get task by task_id"""
return self._tasks.get(task_id)
def get_agent_tasks(self, agent_id: str) -> list[AgentTask]:
"""Get all tasks for an agent"""
if agent_id not in self._agent_tasks:
return []
return [
self._tasks[task_id]
for task_id in self._agent_tasks[agent_id]
]
def is_agent_running(self, agent_id: str) -> bool:
"""Check if an agent has any running tasks"""
return agent_id in self._agent_tasks and len(self._agent_tasks[agent_id]) > 0
def get_task_id_by_run(self, agent_run_id: str) -> Optional[str]:
"""Get task_id by agent_run_id"""
return self._run_tasks.get(agent_run_id)
async def cleanup_completed_tasks(self) -> int:
"""Clean up completed/failed tasks"""
to_remove = []
for task_id, agent_task in self._tasks.items():
if agent_task.task.done():
to_remove.append(task_id)
for task_id in to_remove:
self.unregister_task(task_id)
if to_remove:
logger.info(f"Cleaned up {len(to_remove)} completed tasks")
return len(to_remove)
def get_all_running_agents(self) -> Dict[str, list[str]]:
"""Get all agents with running tasks"""
return {
agent_id: list(task_ids)
for agent_id, task_ids in self._agent_tasks.items()
}
def _reset(self) -> None:
"""
Reset the registry to initial state.
WARNING: This method is only for test use. It clears all registry state.
Do not call this in production code.
"""
self._tasks.clear()
self._agent_tasks.clear()
self._run_tasks.clear()
self._initialized = False
# Global registry instance
agent_task_registry = AgentTaskRegistry()
def register_agent_task(
agent_id: str,
agent_run_id: str,
task: asyncio.Task,
user_id: str
) -> str:
"""Helper function to register an agent task and return task_id"""
import uuid
task_id = str(uuid.uuid4())
agent_task_registry.register_task(
task_id=task_id,
agent_id=agent_id,
agent_run_id=agent_run_id,
task=task,
user_id=user_id
)
return task_id
|