agno 2.3.19__py3-none-any.whl → 2.3.20__py3-none-any.whl

agno/run/__init__.py

@@ -1,6 +1,4 @@
 from agno.run.base import RunContext, RunStatus
+from agno.run.cancel import get_cancellation_manager, set_cancellation_manager
 
-__all__ = [
-    "RunContext",
-    "RunStatus",
-]
+__all__ = ["RunContext", "RunStatus", "get_cancellation_manager", "set_cancellation_manager"]

agno/run/cancel.py

@@ -1,64 +1,37 @@
 """Run cancellation management."""
 
-import threading
 from typing import Dict
 
-from agno.exceptions import RunCancelledException
+from agno.run.cancellation_management.base import BaseRunCancellationManager
+from agno.run.cancellation_management.in_memory_cancellation_manager import InMemoryRunCancellationManager
 from agno.utils.log import logger
 
+# Global cancellation manager instance
+_cancellation_manager: BaseRunCancellationManager = InMemoryRunCancellationManager()
 
-class RunCancellationManager:
-    """Manages cancellation state for agent runs."""
-
-    def __init__(self):
-        self._cancelled_runs: Dict[str, bool] = {}
-        self._lock = threading.Lock()
-
-    def register_run(self, run_id: str) -> None:
-        """Register a new run as not cancelled."""
-        with self._lock:
-            self._cancelled_runs[run_id] = False
-
-    def cancel_run(self, run_id: str) -> bool:
-        """Cancel a run by marking it as cancelled.
-
-        Returns:
-            bool: True if run was found and cancelled, False if run not found.
-        """
-        with self._lock:
-            if run_id in self._cancelled_runs:
-                self._cancelled_runs[run_id] = True
-                logger.info(f"Run {run_id} marked for cancellation")
-                return True
-            else:
-                logger.warning(f"Attempted to cancel unknown run {run_id}")
-                return False
-
-    def is_cancelled(self, run_id: str) -> bool:
-        """Check if a run is cancelled."""
-        with self._lock:
-            return self._cancelled_runs.get(run_id, False)
-
-    def cleanup_run(self, run_id: str) -> None:
-        """Remove a run from tracking (called when run completes)."""
-        with self._lock:
-            if run_id in self._cancelled_runs:
-                del self._cancelled_runs[run_id]
-
-    def raise_if_cancelled(self, run_id: str) -> None:
-        """Check if a run should be cancelled and raise exception if so."""
-        if self.is_cancelled(run_id):
-            logger.info(f"Cancelling run {run_id}")
-            raise RunCancelledException(f"Run {run_id} was cancelled")
-
-    def get_active_runs(self) -> Dict[str, bool]:
-        """Get all currently tracked runs and their cancellation status."""
-        with self._lock:
-            return self._cancelled_runs.copy()
 
+def set_cancellation_manager(manager: BaseRunCancellationManager) -> None:
+    """Set a custom cancellation manager.
 
-# Global cancellation manager instance
-_cancellation_manager = RunCancellationManager()
+    Args:
+        manager: A BaseRunCancellationManager instance or subclass.
+
+    Example:
+        ```python
+        class MyCustomManager(BaseRunCancellationManager):
+            ....
+
+        set_cancellation_manager(MyCustomManager())
+        ```
+    """
+    global _cancellation_manager
+    _cancellation_manager = manager
+    logger.info(f"Cancellation manager set to {type(manager).__name__}")
+
+
+def get_cancellation_manager() -> BaseRunCancellationManager:
+    """Get the current cancellation manager instance."""
+    return _cancellation_manager
 
 
 def register_run(run_id: str) -> None:
@@ -66,16 +39,56 @@ def register_run(run_id: str) -> None:
     _cancellation_manager.register_run(run_id)
 
 
+async def aregister_run(run_id: str) -> None:
+    """Register a new run for cancellation tracking (async version)."""
+    await _cancellation_manager.aregister_run(run_id)
+
+
 def cancel_run(run_id: str) -> bool:
     """Cancel a run."""
     return _cancellation_manager.cancel_run(run_id)
 
 
+async def acancel_run(run_id: str) -> bool:
+    """Cancel a run (async version)."""
+    return await _cancellation_manager.acancel_run(run_id)
+
+
+def is_cancelled(run_id: str) -> bool:
+    """Check if a run is cancelled."""
+    return _cancellation_manager.is_cancelled(run_id)
+
+
+async def ais_cancelled(run_id: str) -> bool:
+    """Check if a run is cancelled (async version)."""
+    return await _cancellation_manager.ais_cancelled(run_id)
+
+
 def cleanup_run(run_id: str) -> None:
     """Clean up cancellation tracking for a completed run."""
     _cancellation_manager.cleanup_run(run_id)
 
 
+async def acleanup_run(run_id: str) -> None:
+    """Clean up cancellation tracking for a completed run (async version)."""
+    await _cancellation_manager.acleanup_run(run_id)
+
+
 def raise_if_cancelled(run_id: str) -> None:
     """Check if a run should be cancelled and raise exception if so."""
     _cancellation_manager.raise_if_cancelled(run_id)
+
+
+async def araise_if_cancelled(run_id: str) -> None:
+    """Check if a run should be cancelled and raise exception if so (async version)."""
+    await _cancellation_manager.araise_if_cancelled(run_id)
+
+
+def get_active_runs() -> Dict[str, bool]:
+    """Get all currently tracked runs and their cancellation status."""
+    return _cancellation_manager.get_active_runs()
+
+
+async def aget_active_runs() -> Dict[str, bool]:
+    """Get all currently tracked runs and their cancellation status (async version)."""
+    return await _cancellation_manager.aget_active_runs()

agno/run/cancellation_management/__init__.py

@@ -0,0 +1,9 @@
+from agno.run.cancellation_management.base import BaseRunCancellationManager
+from agno.run.cancellation_management.in_memory_cancellation_manager import InMemoryRunCancellationManager
+from agno.run.cancellation_management.redis_cancellation_manager import RedisRunCancellationManager
+
+__all__ = [
+    "BaseRunCancellationManager",
+    "InMemoryRunCancellationManager",
+    "RedisRunCancellationManager",
+]

agno/run/cancellation_management/base.py

@@ -0,0 +1,78 @@
+from abc import ABC, abstractmethod
+from typing import Dict
+
+
+class BaseRunCancellationManager(ABC):
+    """Manages cancellation state for agent runs.
+
+    This class can be extended to implement custom cancellation logic.
+    Use set_cancellation_manager() to replace the global instance with your own.
+    """
+
+    @abstractmethod
+    def register_run(self, run_id: str) -> None:
+        """Register a new run as not cancelled."""
+        pass
+
+    @abstractmethod
+    async def aregister_run(self, run_id: str) -> None:
+        """Register a new run as not cancelled (async version)."""
+        pass
+
+    @abstractmethod
+    def cancel_run(self, run_id: str) -> bool:
+        """Cancel a run by marking it as cancelled.
+
+        Returns:
+            bool: True if run was found and cancelled, False if run not found.
+        """
+        pass
+
+    @abstractmethod
+    async def acancel_run(self, run_id: str) -> bool:
+        """Cancel a run by marking it as cancelled (async version).
+
+        Returns:
+            bool: True if run was found and cancelled, False if run not found.
+        """
+        pass
+
+    @abstractmethod
+    def is_cancelled(self, run_id: str) -> bool:
+        """Check if a run is cancelled."""
+        pass
+
+    @abstractmethod
+    async def ais_cancelled(self, run_id: str) -> bool:
+        """Check if a run is cancelled (async version)."""
+        pass
+
+    @abstractmethod
+    def cleanup_run(self, run_id: str) -> None:
+        """Remove a run from tracking (called when run completes)."""
+        pass
+
+    @abstractmethod
+    async def acleanup_run(self, run_id: str) -> None:
+        """Remove a run from tracking (called when run completes) (async version)."""
+        pass
+
+    @abstractmethod
+    def raise_if_cancelled(self, run_id: str) -> None:
+        """Check if a run should be cancelled and raise exception if so."""
+        pass
+
+    @abstractmethod
+    async def araise_if_cancelled(self, run_id: str) -> None:
+        """Check if a run should be cancelled and raise exception if so (async version)."""
+        pass
+
+    @abstractmethod
+    def get_active_runs(self) -> Dict[str, bool]:
+        """Get all currently tracked runs and their cancellation status."""
+        pass
+
+    @abstractmethod
+    async def aget_active_runs(self) -> Dict[str, bool]:
+        """Get all currently tracked runs and their cancellation status (async version)."""
+        pass

agno/run/cancellation_management/in_memory_cancellation_manager.py

@@ -0,0 +1,100 @@
+"""Run cancellation management."""
+
+import asyncio
+import threading
+from typing import Dict
+
+from agno.exceptions import RunCancelledException
+from agno.run.cancellation_management.base import BaseRunCancellationManager
+from agno.utils.log import logger
+
+
+class InMemoryRunCancellationManager(BaseRunCancellationManager):
+    def __init__(self):
+        self._cancelled_runs: Dict[str, bool] = {}
+        self._lock = threading.Lock()
+        self._async_lock = asyncio.Lock()
+
+    def register_run(self, run_id: str) -> None:
+        """Register a new run as not cancelled."""
+        with self._lock:
+            self._cancelled_runs[run_id] = False
+
+    async def aregister_run(self, run_id: str) -> None:
+        """Register a new run as not cancelled (async version)."""
+        async with self._async_lock:
+            self._cancelled_runs[run_id] = False
+
+    def cancel_run(self, run_id: str) -> bool:
+        """Cancel a run by marking it as cancelled.
+
+        Returns:
+            bool: True if run was found and cancelled, False if run not found.
+        """
+        with self._lock:
+            if run_id in self._cancelled_runs:
+                self._cancelled_runs[run_id] = True
+                logger.info(f"Run {run_id} marked for cancellation")
+                return True
+            else:
+                logger.warning(f"Attempted to cancel unknown run {run_id}")
+                return False
+
+    async def acancel_run(self, run_id: str) -> bool:
+        """Cancel a run by marking it as cancelled (async version).
+
+        Returns:
+            bool: True if run was found and cancelled, False if run not found.
+        """
+        async with self._async_lock:
+            if run_id in self._cancelled_runs:
+                self._cancelled_runs[run_id] = True
+                logger.info(f"Run {run_id} marked for cancellation")
+                return True
+            else:
+                logger.warning(f"Attempted to cancel unknown run {run_id}")
+                return False
+
+    def is_cancelled(self, run_id: str) -> bool:
+        """Check if a run is cancelled."""
+        with self._lock:
+            return self._cancelled_runs.get(run_id, False)
+
+    async def ais_cancelled(self, run_id: str) -> bool:
+        """Check if a run is cancelled (async version)."""
+        async with self._async_lock:
+            return self._cancelled_runs.get(run_id, False)
+
+    def cleanup_run(self, run_id: str) -> None:
+        """Remove a run from tracking (called when run completes)."""
+        with self._lock:
+            if run_id in self._cancelled_runs:
+                del self._cancelled_runs[run_id]
+
+    async def acleanup_run(self, run_id: str) -> None:
+        """Remove a run from tracking (called when run completes) (async version)."""
+        async with self._async_lock:
+            if run_id in self._cancelled_runs:
+                del self._cancelled_runs[run_id]
+
+    def raise_if_cancelled(self, run_id: str) -> None:
+        """Check if a run should be cancelled and raise exception if so."""
+        if self.is_cancelled(run_id):
+            logger.info(f"Cancelling run {run_id}")
+            raise RunCancelledException(f"Run {run_id} was cancelled")
+
+    async def araise_if_cancelled(self, run_id: str) -> None:
+        """Check if a run should be cancelled and raise exception if so (async version)."""
+        if await self.ais_cancelled(run_id):
+            logger.info(f"Cancelling run {run_id}")
+            raise RunCancelledException(f"Run {run_id} was cancelled")
+
+    def get_active_runs(self) -> Dict[str, bool]:
+        """Get all currently tracked runs and their cancellation status."""
+        with self._lock:
+            return self._cancelled_runs.copy()
+
+    async def aget_active_runs(self) -> Dict[str, bool]:
+        """Get all currently tracked runs and their cancellation status (async version)."""
+        async with self._async_lock:
+            return self._cancelled_runs.copy()

agno/run/cancellation_management/redis_cancellation_manager.py

@@ -0,0 +1,236 @@
+"""Redis-based run cancellation management."""
+
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+
+from agno.exceptions import RunCancelledException
+from agno.run.cancellation_management.base import BaseRunCancellationManager
+from agno.utils.log import logger
+
+# Defer import error until class instantiation
+_redis_available = True
+_redis_import_error: Optional[str] = None
+
+try:
+    from redis import Redis, RedisCluster
+    from redis.asyncio import Redis as AsyncRedis
+    from redis.asyncio import RedisCluster as AsyncRedisCluster
+except ImportError:
+    _redis_available = False
+    _redis_import_error = "`redis` not installed. Please install it using `pip install redis`"
+    # Type hints for when redis is not installed
+    if TYPE_CHECKING:
+        from redis import Redis, RedisCluster
+        from redis.asyncio import Redis as AsyncRedis
+        from redis.asyncio import RedisCluster as AsyncRedisCluster
+    else:
+        Redis = Any
+        RedisCluster = Any
+        AsyncRedis = Any
+        AsyncRedisCluster = Any
+
+
+class RedisRunCancellationManager(BaseRunCancellationManager):
+    """Redis-based cancellation manager for distributed run cancellation.
+    This manager stores run cancellation state in Redis, enabling cancellation
+    across multiple processes or services.
+
+    To use: call the set_cancellation_manager function to set the cancellation manager.
+    Args:
+        redis_client: Sync Redis client for sync methods. Can be Redis or RedisCluster.
+        async_redis_client: Async Redis client for async methods. Can be AsyncRedis or AsyncRedisCluster.
+        key_prefix: Prefix for Redis keys. Defaults to "agno:run:cancellation:".
+        ttl_seconds: TTL for keys in seconds. Defaults to 86400 (1 day).
+            Keys auto-expire to prevent orphaned keys if runs aren't cleaned up.
+            Set to None to disable expiration.
+    """
+
+    DEFAULT_TTL_SECONDS = 60 * 60 * 24  # 1 day
+
+    def __init__(
+        self,
+        redis_client: Optional[Union[Redis, RedisCluster]] = None,
+        async_redis_client: Optional[Union[AsyncRedis, AsyncRedisCluster]] = None,
+        key_prefix: str = "agno:run:cancellation:",
+        ttl_seconds: Optional[int] = DEFAULT_TTL_SECONDS,
+    ):
+        if not _redis_available:
+            raise ImportError(_redis_import_error)
+
+        super().__init__()
+        self.redis_client = redis_client
+        self.async_redis_client = async_redis_client
+        self.key_prefix = key_prefix
+        self.ttl_seconds = ttl_seconds
+
+        if redis_client is None and async_redis_client is None:
+            raise ValueError("At least one of redis_client or async_redis_client must be provided")
+
+    def _get_key(self, run_id: str) -> str:
+        """Get the Redis key for a run ID."""
+        return f"{self.key_prefix}{run_id}"
+
+    def _ensure_sync_client(self) -> Union[Redis, RedisCluster]:
+        """Ensure sync client is available."""
+        if self.redis_client is None:
+            raise RuntimeError("Sync Redis client not provided. Use async methods or provide a sync client.")
+        return self.redis_client
+
+    def _ensure_async_client(self) -> Union[AsyncRedis, AsyncRedisCluster]:
+        """Ensure async client is available."""
+        if self.async_redis_client is None:
+            raise RuntimeError("Async Redis client not provided. Use sync methods or provide an async client.")
+        return self.async_redis_client
+
+    def register_run(self, run_id: str) -> None:
+        """Register a new run as not cancelled."""
+        client = self._ensure_sync_client()
+        key = self._get_key(run_id)
+        client.set(key, "0", ex=self.ttl_seconds)
+
+    async def aregister_run(self, run_id: str) -> None:
+        """Register a new run as not cancelled (async version)."""
+        client = self._ensure_async_client()
+        key = self._get_key(run_id)
+        await client.set(key, "0", ex=self.ttl_seconds)
+
+    def cancel_run(self, run_id: str) -> bool:
+        """Cancel a run by marking it as cancelled.
+
+        Returns:
+            bool: True if run was found and cancelled, False if run not found.
+        """
+        client = self._ensure_sync_client()
+        key = self._get_key(run_id)
+
+        # Atomically set to "1" only if key exists (XX flag)
+        result = client.set(key, "1", ex=self.ttl_seconds, xx=True)
+
+        if result:
+            logger.info(f"Run {run_id} marked for cancellation")
+            return True
+        else:
+            logger.warning(f"Attempted to cancel unknown run {run_id}")
+            return False
+
+    async def acancel_run(self, run_id: str) -> bool:
+        """Cancel a run by marking it as cancelled (async version).
+
+        Returns:
+            bool: True if run was found and cancelled, False if run not found.
+        """
+        client = self._ensure_async_client()
+        key = self._get_key(run_id)
+
+        # Atomically set to "1" only if key exists (XX flag)
+        result = await client.set(key, "1", ex=self.ttl_seconds, xx=True)
+
+        if result:
+            logger.info(f"Run {run_id} marked for cancellation")
+            return True
+        else:
+            logger.warning(f"Attempted to cancel unknown run {run_id}")
+            return False
+
+    def is_cancelled(self, run_id: str) -> bool:
+        """Check if a run is cancelled."""
+        client = self._ensure_sync_client()
+        key = self._get_key(run_id)
+        value = client.get(key)
+        if value is None:
+            return False
+        # Redis returns bytes, handle both bytes and str
+        if isinstance(value, bytes):
+            return value == b"1"
+        return value == "1"
+
+    async def ais_cancelled(self, run_id: str) -> bool:
+        """Check if a run is cancelled (async version)."""
+        client = self._ensure_async_client()
+        key = self._get_key(run_id)
+        value = await client.get(key)
+        if value is None:
+            return False
+        # Redis returns bytes, handle both bytes and str
+        if isinstance(value, bytes):
+            return value == b"1"
+        return value == "1"
+
+    def cleanup_run(self, run_id: str) -> None:
+        """Remove a run from tracking (called when run completes)."""
+        client = self._ensure_sync_client()
+        key = self._get_key(run_id)
+        client.delete(key)
+
+    async def acleanup_run(self, run_id: str) -> None:
+        """Remove a run from tracking (called when run completes) (async version)."""
+        client = self._ensure_async_client()
+        key = self._get_key(run_id)
+        await client.delete(key)
+
+    def raise_if_cancelled(self, run_id: str) -> None:
+        """Check if a run should be cancelled and raise exception if so."""
+        if self.is_cancelled(run_id):
+            logger.info(f"Cancelling run {run_id}")
+            raise RunCancelledException(f"Run {run_id} was cancelled")
+
+    async def araise_if_cancelled(self, run_id: str) -> None:
+        """Check if a run should be cancelled and raise exception if so (async version)."""
+        if await self.ais_cancelled(run_id):
+            logger.info(f"Cancelling run {run_id}")
+            raise RunCancelledException(f"Run {run_id} was cancelled")
+
+    def get_active_runs(self) -> Dict[str, bool]:
+        """Get all currently tracked runs and their cancellation status.
+
+        Note: Uses scan_iter which works correctly with both standalone Redis
+        and Redis Cluster (scans all nodes in cluster mode).
+        """
+        client = self._ensure_sync_client()
+        result: Dict[str, bool] = {}
+
+        # scan_iter handles cluster mode correctly (scans all nodes)
+        pattern = f"{self.key_prefix}*"
+        for key in client.scan_iter(match=pattern, count=100):
+            # Extract run_id from key
+            if isinstance(key, bytes):
+                key = key.decode("utf-8")
+            run_id = key[len(self.key_prefix) :]
+
+            # Get value
+            value = client.get(key)
+            if value is not None:
+                if isinstance(value, bytes):
+                    is_cancelled = value == b"1"
+                else:
+                    is_cancelled = value == "1"
+                result[run_id] = is_cancelled
+
+        return result
+
+    async def aget_active_runs(self) -> Dict[str, bool]:
+        """Get all currently tracked runs and their cancellation status (async version).
+
+        Note: Uses scan_iter which works correctly with both standalone Redis
+        and Redis Cluster (scans all nodes in cluster mode).
+        """
+        client = self._ensure_async_client()
+        result: Dict[str, bool] = {}
+
+        # scan_iter handles cluster mode correctly (scans all nodes)
+        pattern = f"{self.key_prefix}*"
+        async for key in client.scan_iter(match=pattern, count=100):
+            # Extract run_id from key
+            if isinstance(key, bytes):
+                key = key.decode("utf-8")
+            run_id = key[len(self.key_prefix) :]
+
+            # Get value
+            value = await client.get(key)
+            if value is not None:
+                if isinstance(value, bytes):
+                    is_cancelled = value == b"1"
+                else:
+                    is_cancelled = value == "1"
+                result[run_id] = is_cancelled
+
+        return result