Spaces:
Sleeping
Sleeping
| """ | |
| 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__) | |
| 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 | |