agent-killswitch 0.1.0__py3-none-any.whl

agent_killswitch/core/models.py

@@ -0,0 +1,173 @@
+"""Data models for agent-killswitch.
+
+Defines the core data structures used throughout the library.
+All models are dataclasses for zero-dependency operation (no Pydantic required).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any
+
+from .enums import CircuitState, HeartbeatStatus, KillLevel, KillScope
+
+
+@dataclass(frozen=True)
+class KillEvent:
+    """Record of a kill switch activation event.
+
+    Immutable record that captures all details of a kill switch
+    activation for audit trail purposes.
+
+    Attributes:
+        event_id: Unique identifier for this event.
+        level: The kill level activated (PAUSE/STOP/KILL).
+        scope: The scope of the activation.
+        scope_target: Target identifier (agent_id, crew_id, etc.).
+        reason: Human-readable reason for the activation.
+        activated_by: Identifier of who/what activated the switch.
+        timestamp: UTC timestamp of the activation.
+        latency_ms: Activation latency in milliseconds.
+        metadata: Additional metadata for audit logging.
+    """
+
+    event_id: str
+    level: KillLevel
+    scope: KillScope
+    scope_target: str | None
+    reason: str
+    activated_by: str
+    timestamp: datetime
+    latency_ms: float
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass(frozen=True)
+class KillStatus:
+    """Current status of a kill switch check.
+
+    Returned by KillSwitch.check() to provide full status information.
+
+    Attributes:
+        is_killed: Whether any kill switch is active.
+        level: The active kill level, or None if not killed.
+        scope: The scope that matched, or None.
+        scope_target: The target that matched, or None.
+        reason: The reason for the kill, or None.
+        activated_at: When the kill was activated, or None.
+    """
+
+    is_killed: bool
+    level: KillLevel | None = None
+    scope: KillScope | None = None
+    scope_target: str | None = None
+    reason: str | None = None
+    activated_at: datetime | None = None
+
+    @staticmethod
+    def safe() -> KillStatus:
+        """Return a safe (not killed) status."""
+        return KillStatus(is_killed=False)
+
+    @staticmethod
+    def killed(
+        level: KillLevel,
+        scope: KillScope,
+        scope_target: str | None = None,
+        reason: str = "",
+        activated_at: datetime | None = None,
+    ) -> KillStatus:
+        """Return a killed status with details."""
+        return KillStatus(
+            is_killed=True,
+            level=level,
+            scope=scope,
+            scope_target=scope_target,
+            reason=reason,
+            activated_at=activated_at,
+        )
+
+
+@dataclass
+class AgentRegistration:
+    """Registration record for an agent in the kill switch system.
+
+    Attributes:
+        agent_id: Unique identifier for the agent.
+        parent_id: ID of the parent agent (for cascading), or None if root.
+        registered_at: UTC timestamp of registration.
+        metadata: Additional metadata about the agent.
+    """
+
+    agent_id: str
+    parent_id: str | None = None
+    registered_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class HeartbeatRecord:
+    """Record of an agent's heartbeat state.
+
+    Attributes:
+        agent_id: The agent being monitored.
+        last_heartbeat: UTC timestamp of last heartbeat.
+        status: Current health status.
+        miss_count: Number of consecutive missed heartbeats.
+        interval_seconds: Expected heartbeat interval.
+    """
+
+    agent_id: str
+    last_heartbeat: datetime | None = None
+    status: HeartbeatStatus = HeartbeatStatus.UNKNOWN
+    miss_count: int = 0
+    interval_seconds: float = 5.0
+
+
+@dataclass
+class CircuitBreakerState:
+    """State of a circuit breaker.
+
+    Attributes:
+        name: Name/identifier of the circuit breaker.
+        state: Current circuit state.
+        failure_count: Consecutive failure count.
+        success_count: Consecutive success count in HALF_OPEN state.
+        last_failure_time: When the last failure occurred.
+        last_state_change: When the state last changed.
+        failure_threshold: Number of failures to trip the breaker.
+        recovery_timeout: Seconds to wait before trying HALF_OPEN.
+        half_open_max_calls: Max calls allowed in HALF_OPEN state.
+    """
+
+    name: str
+    state: CircuitState = CircuitState.CLOSED
+    failure_count: int = 0
+    success_count: int = 0
+    last_failure_time: float | None = None
+    last_state_change: float = field(default_factory=lambda: 0.0)
+    failure_threshold: int = 5
+    recovery_timeout: float = 30.0
+    half_open_max_calls: int = 1
+
+
+@dataclass
+class BudgetStatus:
+    """Budget tracking status for an agent.
+
+    Attributes:
+        agent_id: The agent being tracked.
+        budget_limit: Maximum allowed cost.
+        current_cost: Current accumulated cost.
+        utilization: Budget utilization as a fraction (0.0 to 1.0+).
+        alert_thresholds: Thresholds that trigger alerts.
+        alerts_fired: Thresholds that have already fired alerts.
+    """
+
+    agent_id: str
+    budget_limit: float = 0.0
+    current_cost: float = 0.0
+    utilization: float = 0.0
+    alert_thresholds: list[float] = field(default_factory=lambda: [0.50, 0.80, 0.95])
+    alerts_fired: set[float] = field(default_factory=set)

agent_killswitch/decorators.py

@@ -0,0 +1,302 @@
+"""Decorators for easy kill switch integration.
+
+Provides decorator-based APIs for common patterns:
+- @killswitch_protected: Check kill switch before/after function execution
+- @with_heartbeat: Send heartbeats during long-running operations
+- @circuit_breaker: Wrap functions with circuit breaker protection
+
+All decorators support both sync and async functions.
+"""
+
+from __future__ import annotations
+
+import functools
+import inspect
+import logging
+import threading
+from collections.abc import Callable
+from typing import Any, TypeVar, Union
+
+from .core.circuit_breaker import CircuitBreaker
+from .core.enums import KillScope
+from .core.killswitch import KillSwitch
+from .core.models import KillStatus
+
+logger = logging.getLogger(__name__)
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+class KillSwitchTriggered(Exception):
+    """Raised when a kill switch check finds an active kill.
+
+    Attributes:
+        status: The KillStatus that triggered the exception.
+    """
+
+    def __init__(self, status: KillStatus) -> None:
+        self.status = status
+        level_str = status.level.value if status.level else "unknown"
+        super().__init__(
+            f"Kill switch triggered: level={level_str}, "
+            f"scope={status.scope}, target={status.scope_target}, "
+            f"reason={status.reason}"
+        )
+
+
+def killswitch_protected(
+    killswitch: KillSwitch | None = None,
+    scope: Union[str, KillScope] = KillScope.AGENT,
+    agent_id: str | None = None,
+    crew_id: str | None = None,
+    on_killed: Callable[[KillStatus], None] | None = None,
+    raise_on_kill: bool = True,
+) -> Callable[[F], F]:
+    """Decorator that checks the kill switch before and after execution.
+
+    If the kill switch is active, either raises KillSwitchTriggered
+    or calls the on_killed callback (depending on raise_on_kill).
+
+    Example::
+
+        ks = KillSwitch()
+
+
+        @killswitch_protected(killswitch=ks, agent_id="my-agent")
+        async def my_agent_task():
+            # ... agent work ...
+            pass
+
+    Args:
+        killswitch: The KillSwitch instance to check. If None, the
+            function runs without protection (useful for testing).
+        scope: The scope to check. Default AGENT.
+        agent_id: The agent ID for scope checks.
+        crew_id: The crew ID for crew scope checks.
+        on_killed: Callback when kill is detected.
+        raise_on_kill: If True, raise KillSwitchTriggered. Default True.
+
+    Returns:
+        Decorator function.
+    """
+    if isinstance(scope, str):
+        scope = KillScope(scope)
+
+    def decorator(func: F) -> F:
+        if inspect.iscoroutinefunction(func):
+
+            @functools.wraps(func)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                if killswitch is not None:
+                    status = killswitch.check(
+                        scope_target=agent_id,
+                        crew_id=crew_id,
+                    )
+                    if status.is_killed:
+                        if on_killed is not None:
+                            on_killed(status)
+                        if raise_on_kill:
+                            raise KillSwitchTriggered(status)
+                        return None
+
+                result = await func(*args, **kwargs)
+
+                if killswitch is not None:
+                    status = killswitch.check(
+                        scope_target=agent_id,
+                        crew_id=crew_id,
+                    )
+                    if status.is_killed:
+                        if on_killed is not None:
+                            on_killed(status)
+                        if raise_on_kill:
+                            raise KillSwitchTriggered(status)
+
+                return result
+
+            return async_wrapper  # type: ignore[return-value]
+        else:
+
+            @functools.wraps(func)
+            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+                if killswitch is not None:
+                    status = killswitch.check(
+                        scope_target=agent_id,
+                        crew_id=crew_id,
+                    )
+                    if status.is_killed:
+                        if on_killed is not None:
+                            on_killed(status)
+                        if raise_on_kill:
+                            raise KillSwitchTriggered(status)
+                        return None
+
+                result = func(*args, **kwargs)
+
+                if killswitch is not None:
+                    status = killswitch.check(
+                        scope_target=agent_id,
+                        crew_id=crew_id,
+                    )
+                    if status.is_killed:
+                        if on_killed is not None:
+                            on_killed(status)
+                        if raise_on_kill:
+                            raise KillSwitchTriggered(status)
+
+                return result
+
+            return sync_wrapper  # type: ignore[return-value]
+
+    return decorator
+
+
+def with_heartbeat(
+    monitor: Any = None,
+    agent_id: str | None = None,
+    interval: float = 5.0,
+) -> Callable[[F], F]:
+    """Decorator that sends heartbeats during function execution.
+
+    Starts a background thread that sends periodic heartbeats to the
+    HeartbeatMonitor while the decorated function is running.
+
+    Example::
+
+        from agent_killswitch import HeartbeatMonitor
+
+        monitor = HeartbeatMonitor()
+
+
+        @with_heartbeat(monitor=monitor, agent_id="my-agent", interval=5.0)
+        async def long_running_agent():
+            # heartbeats are sent automatically
+            await do_work()
+
+    Args:
+        monitor: HeartbeatMonitor instance. If None, no heartbeats sent.
+        agent_id: The agent ID for heartbeats.
+        interval: Heartbeat interval in seconds. Default 5.0.
+
+    Returns:
+        Decorator function.
+    """
+
+    def decorator(func: F) -> F:
+        if inspect.iscoroutinefunction(func):
+
+            @functools.wraps(func)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                if monitor is None or agent_id is None:
+                    return await func(*args, **kwargs)
+
+                monitor.start(agent_id, interval_seconds=interval)
+                stop_event = threading.Event()
+
+                def heartbeat_loop() -> None:
+                    while not stop_event.is_set():
+                        monitor.heartbeat(agent_id)
+                        stop_event.wait(interval)
+
+                thread = threading.Thread(
+                    target=heartbeat_loop,
+                    daemon=True,
+                    name=f"heartbeat-{agent_id}",
+                )
+                thread.start()
+
+                try:
+                    return await func(*args, **kwargs)
+                finally:
+                    stop_event.set()
+                    thread.join(timeout=interval + 1.0)
+                    monitor.stop(agent_id)
+
+            return async_wrapper  # type: ignore[return-value]
+        else:
+
+            @functools.wraps(func)
+            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+                if monitor is None or agent_id is None:
+                    return func(*args, **kwargs)
+
+                monitor.start(agent_id, interval_seconds=interval)
+                stop_event = threading.Event()
+
+                def heartbeat_loop() -> None:
+                    while not stop_event.is_set():
+                        monitor.heartbeat(agent_id)
+                        stop_event.wait(interval)
+
+                thread = threading.Thread(
+                    target=heartbeat_loop,
+                    daemon=True,
+                    name=f"heartbeat-{agent_id}",
+                )
+                thread.start()
+
+                try:
+                    return func(*args, **kwargs)
+                finally:
+                    stop_event.set()
+                    thread.join(timeout=interval + 1.0)
+                    monitor.stop(agent_id)
+
+            return sync_wrapper  # type: ignore[return-value]
+
+    return decorator
+
+
+def circuit_breaker(
+    name: str = "default",
+    failure_threshold: int = 5,
+    recovery_timeout: float = 30.0,
+    excluded_exceptions: tuple[type[BaseException], ...] | None = None,
+) -> Callable[[F], F]:
+    """Decorator that wraps a function with circuit breaker protection.
+
+    Creates a CircuitBreaker instance per decorated function and
+    routes all calls through it.
+
+    Example::
+
+        @circuit_breaker(name="openai", failure_threshold=3, recovery_timeout=60)
+        async def call_openai(prompt: str) -> str:
+            return await openai.chat(prompt)
+
+    Args:
+        name: Circuit breaker name.
+        failure_threshold: Failures before opening.
+        recovery_timeout: Seconds before trying half-open.
+        excluded_exceptions: Exceptions that don't count as failures.
+
+    Returns:
+        Decorator function.
+    """
+    breaker = CircuitBreaker(
+        name=name,
+        failure_threshold=failure_threshold,
+        recovery_timeout=recovery_timeout,
+        excluded_exceptions=excluded_exceptions,
+    )
+
+    def decorator(func: F) -> F:
+        if inspect.iscoroutinefunction(func):
+
+            @functools.wraps(func)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                return await breaker.async_call(func, *args, **kwargs)
+
+            # Expose the breaker for inspection
+            async_wrapper._circuit_breaker = breaker  # type: ignore[attr-defined]
+            return async_wrapper  # type: ignore[return-value]
+        else:
+
+            @functools.wraps(func)
+            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+                return breaker.call(func, *args, **kwargs)
+
+            sync_wrapper._circuit_breaker = breaker  # type: ignore[attr-defined]
+            return sync_wrapper  # type: ignore[return-value]
+
+    return decorator

agent_killswitch/integrations/__init__.py

@@ -0,0 +1,17 @@
+"""Framework integrations for agent-killswitch.
+
+Each integration is a thin adapter that connects the core kill switch
+to a specific AI agent framework. Integrations are lazy-loaded to
+avoid requiring the framework as a dependency.
+
+Available integrations:
+- langchain: LangChain callback handler
+- langgraph: LangGraph node and conditional edge
+- crewai: CrewAI pre-action hook
+- openai_agents: OpenAI Agents SDK guardrail
+- asyncio_tasks: Generic asyncio.Task cancellation
+"""
+
+from __future__ import annotations
+
+__all__: list[str] = []

agent_killswitch/integrations/asyncio_tasks.py

@@ -0,0 +1,211 @@
+"""asyncio integration for agent-killswitch.
+
+Provides utilities for managing asyncio.Task lifecycle with kill switch
+integration. Supports task group cancellation and periodic kill switch
+checks during async operations.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+from collections.abc import Callable, Coroutine
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+class KillSwitchTaskGroup:
+    """Managed group of asyncio tasks with kill switch integration.
+
+    Wraps asyncio tasks and periodically checks the kill switch.
+    If the kill switch activates, all tasks in the group are cancelled.
+
+    Example::
+
+        from agent_killswitch import KillSwitch
+        from agent_killswitch.integrations.asyncio_tasks import KillSwitchTaskGroup
+
+        ks = KillSwitch()
+
+
+        async def worker(name: str):
+            while True:
+                await asyncio.sleep(1)
+                print(f"{name} working...")
+
+
+        async def main():
+            group = KillSwitchTaskGroup(
+                killswitch=ks,
+                agent_id="async-agents",
+                check_interval=1.0,
+            )
+            group.add_task(worker("a"))
+            group.add_task(worker("b"))
+
+            await group.run()  # Blocks until kill switch or all done
+
+    Args:
+        killswitch: The KillSwitch instance.
+        agent_id: Agent ID for scope checks.
+        crew_id: Optional crew ID.
+        check_interval: How often to check kill switch (seconds). Default 1.0.
+        on_cancelled: Callback when tasks are cancelled by kill switch.
+    """
+
+    def __init__(
+        self,
+        killswitch: Any,
+        agent_id: str | None = None,
+        crew_id: str | None = None,
+        check_interval: float = 1.0,
+        on_cancelled: Callable[[], None] | None = None,
+    ) -> None:
+        self._ks = killswitch
+        self._agent_id = agent_id
+        self._crew_id = crew_id
+        self._check_interval = check_interval
+        self._on_cancelled = on_cancelled
+        self._coros: list[Coroutine[Any, Any, Any]] = []
+        self._tasks: list[asyncio.Task[Any]] = []
+        self._killed = False
+
+    def add_task(self, coro: Coroutine[Any, Any, Any]) -> None:
+        """Add a coroutine to the task group.
+
+        Args:
+            coro: The coroutine to add. Must not have been awaited yet.
+        """
+        self._coros.append(coro)
+
+    async def run(self) -> list[Any]:
+        """Run all tasks with kill switch monitoring.
+
+        Starts all coroutines as tasks and monitors the kill switch.
+        Returns when all tasks complete or kill switch activates.
+
+        Returns:
+            List of results from completed tasks. Results may be
+            asyncio.CancelledError for tasks cancelled by kill switch.
+        """
+        asyncio.get_event_loop()
+
+        # Start all tasks
+        for coro in self._coros:
+            task = asyncio.ensure_future(coro)
+            self._tasks.append(task)
+        self._coros.clear()
+
+        # Start monitor
+        monitor = asyncio.ensure_future(self._monitor_loop())
+
+        # Wait for all tasks to complete (or be cancelled)
+        results: list[Any] = []
+        try:
+            done, _pending = await asyncio.wait(
+                self._tasks,
+                return_when=asyncio.ALL_COMPLETED,
+            )
+            for task in done:
+                try:
+                    results.append(task.result())
+                except (asyncio.CancelledError, Exception) as exc:
+                    results.append(exc)
+        except asyncio.CancelledError:
+            results = [asyncio.CancelledError() for _ in self._tasks]
+        finally:
+            monitor.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await monitor
+
+        return results
+
+    async def _monitor_loop(self) -> None:
+        """Periodically check kill switch and cancel tasks if active."""
+        while True:
+            await asyncio.sleep(self._check_interval)
+
+            status = self._ks.check(
+                scope_target=self._agent_id,
+                crew_id=self._crew_id,
+            )
+
+            if status.is_killed and not self._killed:
+                self._killed = True
+                level_str = status.level.value if status.level else "unknown"
+                logger.warning(
+                    "asyncio_task_group_kill",
+                    extra={
+                        "agent_id": self._agent_id,
+                        "level": level_str,
+                        "task_count": len(self._tasks),
+                    },
+                )
+
+                # Cancel all running tasks
+                for task in self._tasks:
+                    if not task.done():
+                        task.cancel()
+
+                if self._on_cancelled is not None:
+                    try:
+                        self._on_cancelled()
+                    except Exception as exc:
+                        logger.error(
+                            "on_cancelled callback error",
+                            extra={"error": str(exc)},
+                        )
+
+                break
+
+    async def cancel_all(self) -> None:
+        """Manually cancel all tasks in the group."""
+        for task in self._tasks:
+            if not task.done():
+                task.cancel()
+
+    @property
+    def is_killed(self) -> bool:
+        """Whether the task group was killed by the kill switch."""
+        return self._killed
+
+
+async def run_with_killswitch(
+    coro: Coroutine[Any, Any, Any],
+    killswitch: Any,
+    agent_id: str | None = None,
+    crew_id: str | None = None,
+    check_interval: float = 1.0,
+) -> Any:
+    """Run a single coroutine with kill switch monitoring.
+
+    Convenience function for running a single async operation with
+    periodic kill switch checks.
+
+    Args:
+        coro: The coroutine to run.
+        killswitch: The KillSwitch instance.
+        agent_id: Agent ID for scope checks.
+        crew_id: Optional crew ID.
+        check_interval: Kill switch check interval in seconds.
+
+    Returns:
+        The result of the coroutine.
+
+    Raises:
+        asyncio.CancelledError: If killed by kill switch.
+    """
+    group = KillSwitchTaskGroup(
+        killswitch=killswitch,
+        agent_id=agent_id,
+        crew_id=crew_id,
+        check_interval=check_interval,
+    )
+    group.add_task(coro)
+    results = await group.run()
+
+    if results and isinstance(results[0], BaseException):
+        raise results[0]
+    return results[0] if results else None