aeonic-agentguard-sdk-python 0.1.0__py3-none-any.whl

aeonic/core/agent_registry.py

@@ -0,0 +1,209 @@
+"""
+Agent Registry - Global in-memory registry for declared agents.
+
+This module tracks all agents declared via SDK APIs (with_agent() / @with_agent)
+at declaration time, NOT via static code scanning.
+"""
+
+import threading
+from time import time
+from typing import Dict, List, Optional
+
+from aeonic.core.types import AgentDescriptor
+
+# Global registry storage
+_REGISTRY: Dict[str, AgentDescriptor] = {}
+_REGISTRY_LOCK = threading.Lock()
+_INVENTORY_EMITTED = False
+_INVENTORY_LOCK = threading.Lock()
+
+
+def register_agent(
+    name: str,
+    agent_type: Optional[str],
+    model: Optional[List[str]],
+    route_path: Optional[str],
+    http_method: Optional[str],
+    framework: str,
+    service_name: Optional[str],
+) -> None:
+    """
+    Register an agent in the global registry.
+
+    This function is idempotent and never throws exceptions.
+    Called at agent declaration time (when with_agent() is invoked or @with_agent is applied).
+
+    If an agent with the same name already exists but has None route_path/http_method,
+    this will update that registration with the actual values.
+
+    Args:
+        name: Agent name (unique identifier)
+        agent_type: Type of agent (e.g., "conversational", "task")
+        model: List of model identifiers
+        route_path: HTTP route path (if applicable)
+        http_method: HTTP method (if applicable)
+        framework: Framework identifier ("fastapi" or "django")
+        service_name: Service name from SDK config
+    """
+    try:
+        agent_key = _generate_agent_key(name, route_path, http_method, framework)
+
+        with _REGISTRY_LOCK:
+            # Idempotent: skip if already registered with same key
+            if agent_key in _REGISTRY:
+                return
+
+            # If we have actual route_path/http_method, check if there's an existing
+            # registration with the same name but None values (from early registration)
+            if route_path and http_method:
+                # Look for existing registration with same name but None path/method
+                existing_key = _generate_agent_key(name, None, None, framework)
+                if existing_key in _REGISTRY:
+                    # Update the existing registration
+                    existing = _REGISTRY[existing_key]
+                    # Preserve original declared_at
+                    updated_descriptor: AgentDescriptor = {
+                        "name": name,
+                        "type": agent_type or existing.get("type"),
+                        "model": model or existing.get("model"),
+                        "route_path": route_path,
+                        "http_method": http_method,
+                        "framework": framework,
+                        "service_name": service_name or existing.get("service_name"),
+                        "declared_at": existing.get("declared_at", time()),
+                    }
+                    # Remove old key and add new one
+                    del _REGISTRY[existing_key]
+                    _REGISTRY[agent_key] = updated_descriptor
+                    return
+
+                # Also check all existing agents for same name (in case key format differs)
+                for key, existing_agent in list(_REGISTRY.items()):
+                    if (
+                        existing_agent.get("name") == name
+                        and existing_agent.get("framework") == framework
+                        and existing_agent.get("route_path") is None
+                        and existing_agent.get("http_method") is None
+                    ):
+                        # Found a match - update it
+                        matched_descriptor: AgentDescriptor = {
+                            "name": name,
+                            "type": agent_type or existing_agent.get("type"),
+                            "model": model or existing_agent.get("model"),
+                            "route_path": route_path,
+                            "http_method": http_method,
+                            "framework": framework,
+                            "service_name": service_name or existing_agent.get("service_name"),
+                            "declared_at": existing_agent.get("declared_at", time()),
+                        }
+                        # Remove old and add new
+                        del _REGISTRY[key]
+                        _REGISTRY[agent_key] = matched_descriptor
+                        return
+
+            # Create new registration
+            new_descriptor: AgentDescriptor = {
+                "name": name,
+                "type": agent_type,
+                "model": model,
+                "route_path": route_path,
+                "http_method": http_method,
+                "framework": framework,
+                "service_name": service_name,
+                "declared_at": time(),
+            }
+
+            _REGISTRY[agent_key] = new_descriptor
+
+    except Exception:
+        # Never throw - silent failure
+        pass
+
+
+def get_all_agents() -> List[AgentDescriptor]:
+    """
+    Retrieve all registered agents.
+
+    Returns:
+        List of agent descriptors
+    """
+    try:
+        with _REGISTRY_LOCK:
+            return list(_REGISTRY.values())
+    except Exception:
+        return []
+
+
+def should_emit_inventory() -> bool:
+    """
+    Check if inventory should be emitted.
+
+    This checks if inventory has been emitted, but does NOT set the flag.
+    The flag is only set when the API call succeeds (via mark_inventory_emitted()).
+
+    Returns:
+        True if inventory should be emitted (hasn't been emitted yet), False otherwise
+    """
+    global _INVENTORY_EMITTED
+
+    try:
+        with _INVENTORY_LOCK:
+            return not _INVENTORY_EMITTED
+    except Exception:
+        return False
+
+
+def mark_inventory_emitted() -> None:
+    """
+    Mark inventory as successfully emitted.
+
+    This should only be called after a successful API response.
+    """
+    global _INVENTORY_EMITTED
+
+    try:
+        with _INVENTORY_LOCK:
+            _INVENTORY_EMITTED = True
+    except Exception:
+        pass
+
+
+def reset_inventory_emission() -> None:
+    """
+    Reset the inventory emission flag.
+
+    This is primarily used for testing purposes.
+    """
+    global _INVENTORY_EMITTED
+
+    try:
+        with _INVENTORY_LOCK:
+            _INVENTORY_EMITTED = False
+    except Exception:
+        pass
+
+
+def _generate_agent_key(
+    name: str,
+    route_path: Optional[str],
+    http_method: Optional[str],
+    framework: str,
+) -> str:
+    """
+    Generate a unique key for agent registration.
+
+    Args:
+        name: Agent name
+        route_path: HTTP route path
+        http_method: HTTP method
+        framework: Framework identifier
+
+    Returns:
+        Unique agent key
+    """
+    parts = [framework, name]
+    if route_path:
+        parts.append(route_path)
+    if http_method:
+        parts.append(http_method)
+    return ":".join(parts)

aeonic/core/agent_status_cache.py

@@ -0,0 +1,28 @@
+"""Agent status cache for blocking/quarantining agents."""
+
+import threading
+from typing import Dict, Optional
+
+_AGENT_STATUS_CACHE: Dict[str, str] = {}
+_CACHE_LOCK = threading.Lock()
+
+
+def get_agent_status(agent_key: str) -> Optional[str]:
+    """Get agent status from cache."""
+    with _CACHE_LOCK:
+        return _AGENT_STATUS_CACHE.get(agent_key)
+
+
+def set_agent_status(agent_key: str, status: str) -> None:
+    """Set agent status in cache."""
+    with _CACHE_LOCK:
+        _AGENT_STATUS_CACHE[agent_key] = status
+
+
+def should_collect_samples(agent_key: str) -> bool:
+    """Check if samples should be collected for agent."""
+    status = get_agent_status(agent_key)
+    # If no status in cache, assume healthy (collect samples)
+    if status is None:
+        return True
+    return status not in ("blocked", "quarantined")

aeonic/core/agentguard/__init__.py

@@ -0,0 +1,8 @@
+"""AgentGuard registration and management module."""
+
+from aeonic.core.agentguard.registration import (
+    get_agent_guard_credentials,
+    register_agent_guard,
+)
+
+__all__ = ["register_agent_guard", "get_agent_guard_credentials"]

aeonic/core/agentguard/registration.py

@@ -0,0 +1,93 @@
+"""AgentGuard registration with T-LAIOR backend."""
+
+import json
+import logging
+import os
+import threading
+import urllib.request
+from typing import Any, Dict, Optional
+
+from aeonic.core.agentguard.sse_client import start_sse_connection, stop_sse_connection
+
+logger = logging.getLogger("aeonic")
+
+_AGENT_GUARD_CREDENTIALS: Optional[Dict[str, Any]] = None
+_CREDENTIALS_LOCK = threading.Lock()
+
+
+def _get_endpoint() -> str:
+    """Get the Aeonic endpoint from environment variable."""
+    return os.environ.get("AEONIC_ENDPOINT", "http://45.79.111.106:3291")
+
+
+def register_agent_guard() -> Dict[str, Any]:
+    """
+    Register SDK as AgentGuard instance with T-LAIOR.
+
+    Returns:
+        Dictionary containing agentguard_id, auth_token, and intervals.
+
+    Raises:
+        Exception: If registration fails.
+    """
+    global _AGENT_GUARD_CREDENTIALS
+
+    from aeonic.core.init import get_config
+
+    config = get_config()
+    endpoint = _get_endpoint()
+
+    payload = {
+        "api_key": config["api_key"],
+        "agentguard_version": "0.1.0",
+        "capabilities": ["red-team", "leakage", "safety"],
+        "service_name": config.get("service_name"),
+    }
+
+    try:
+        url = f"{endpoint}/agentguard/register"
+        data = json.dumps(payload).encode("utf-8")
+        req = urllib.request.Request(
+            url,
+            data=data,
+            headers={"Content-Type": "application/json"},
+            method="POST",
+        )
+
+        with urllib.request.urlopen(req, timeout=10) as response:
+            response_data = json.loads(response.read())
+
+            with _CREDENTIALS_LOCK:
+                _AGENT_GUARD_CREDENTIALS = {
+                    "agentguard_id": response_data["agentguard_id"],
+                    "auth_token": response_data["auth_token"],
+                    "heartbeat_interval": response_data.get("heartbeat_interval", 30),
+                    "polling_interval": response_data.get("polling_interval", 5),
+                }
+
+            # Start SSE connection for real-time events
+            start_sse_connection()
+
+            logger.info("[Aeonic] AgentGuard registered successfully")
+            return _AGENT_GUARD_CREDENTIALS
+
+    except Exception as error:
+        logger.error(f"[Aeonic] Failed to register AgentGuard: {error}")
+        raise
+
+
+def get_agent_guard_credentials() -> Optional[Dict[str, Any]]:
+    """Get current AgentGuard credentials."""
+    with _CREDENTIALS_LOCK:
+        return _AGENT_GUARD_CREDENTIALS.copy() if _AGENT_GUARD_CREDENTIALS else None
+
+
+def stop_agent_guard() -> None:
+    """Stop AgentGuard and SSE connection."""
+    global _AGENT_GUARD_CREDENTIALS
+
+    # Stop SSE connection
+    stop_sse_connection()
+
+    with _CREDENTIALS_LOCK:
+        _AGENT_GUARD_CREDENTIALS = None

aeonic/core/agentguard/sse_client.py

@@ -0,0 +1,210 @@
+"""SSE client for real-time event streaming."""
+
+import json
+import logging
+import os
+import threading
+from typing import Any, Dict, Optional
+
+import requests
+
+from aeonic.core.agent_status_cache import set_agent_status
+
+logger = logging.getLogger("aeonic")
+
+_SSE_CONNECTION: Optional[threading.Thread] = None
+_STOP_SSE = threading.Event()
+_RECONNECT_ATTEMPTS = 0
+_MAX_RECONNECT_ATTEMPTS = 10
+_INITIAL_RECONNECT_DELAY = 1.0  # 1 second
+
+
+def _get_endpoint() -> str:
+    """Get the Aeonic endpoint from environment variable."""
+    return os.environ.get("AEONIC_ENDPOINT", "http://45.79.111.106:3291")
+
+
+def start_sse_connection() -> None:
+    """Start SSE connection for real-time events."""
+    global _SSE_CONNECTION, _RECONNECT_ATTEMPTS
+
+    from aeonic.core.agentguard.registration import get_agent_guard_credentials
+
+    credentials = get_agent_guard_credentials()
+    if not credentials:
+        logger.warning("[Aeonic] Cannot start SSE: AgentGuard not registered")
+        return
+
+    if _SSE_CONNECTION and _SSE_CONNECTION.is_alive():
+        # Already connected
+        return
+
+    _STOP_SSE.clear()
+    _RECONNECT_ATTEMPTS = 0
+
+    def sse_worker() -> None:
+        """SSE connection worker thread."""
+        global _RECONNECT_ATTEMPTS
+
+        endpoint = _get_endpoint()
+        url = f"{endpoint}/api/sdk/events/stream"
+
+        while not _STOP_SSE.is_set():
+            try:
+                headers = {
+                    "Authorization": f"Bearer {credentials['auth_token']}",
+                    "Accept": "text/event-stream",
+                    "Cache-Control": "no-cache",
+                }
+
+                response = requests.get(
+                    url,
+                    headers=headers,
+                    stream=True,
+                    timeout=None,  # No timeout for SSE
+                )
+
+                response.raise_for_status()
+
+                logger.info("[Aeonic] SSE connection established")
+                _RECONNECT_ATTEMPTS = 0
+
+                # Process SSE stream
+                for line in response.iter_lines(decode_unicode=True):
+                    if _STOP_SSE.is_set():
+                        break
+
+                    if not line:
+                        continue
+
+                    if line.startswith("data: "):
+                        try:
+                            # Remove "data: " prefix and parse JSON
+                            event_data = json.loads(line[6:])
+                            _handle_sse_event(event_data)
+                        except json.JSONDecodeError as e:
+                            logger.error(f"[Aeonic] Failed to parse SSE event: {e}")
+
+            except requests.exceptions.RequestException as e:
+                if _STOP_SSE.is_set():
+                    break
+
+                logger.warning(f"[Aeonic] SSE connection error: {e}")
+
+                # Reconnect with exponential backoff
+                if _RECONNECT_ATTEMPTS < _MAX_RECONNECT_ATTEMPTS:
+                    delay = _INITIAL_RECONNECT_DELAY * (2**_RECONNECT_ATTEMPTS)
+                    _RECONNECT_ATTEMPTS += 1
+                    logger.info(
+                        f"[Aeonic] Reconnecting SSE in {delay}s (attempt {_RECONNECT_ATTEMPTS})"
+                    )
+                    if _STOP_SSE.wait(delay):
+                        break
+                else:
+                    logger.error("[Aeonic] Max SSE reconnection attempts reached")
+                    break
+            except Exception as e:
+                logger.error(f"[Aeonic] Unexpected SSE error: {e}")
+                if _STOP_SSE.wait(5):  # Wait a bit before retrying unexpected errors
+                    break
+
+    _SSE_CONNECTION = threading.Thread(target=sse_worker, name="AeonicSSEThread", daemon=True)
+    _SSE_CONNECTION.start()
+
+
+def _handle_sse_event(data: Dict[str, Any]) -> None:
+    """Handle SSE event."""
+    event_type = data.get("type")
+
+    if event_type == "connected":
+        logger.info(f"[Aeonic] SSE connected: {data.get('message')}")
+
+    elif event_type == "heartbeat":
+        # Heartbeat received - connection is alive
+        logger.debug("[Aeonic] SSE heartbeat received")
+
+    elif event_type == "agent.status_changed":
+        _handle_agent_status_change(data)
+
+    elif event_type == "assessment.job_available":
+        _handle_assessment_job(data)
+
+    else:
+        logger.debug(f"[Aeonic] Unknown SSE event type: {event_type}")
+
+
+def _handle_agent_status_change(data: Dict[str, Any]) -> None:
+    """Handle agent status change event."""
+    agent_id = data.get("agent_id")
+    agent_key = data.get("agent_key")
+    agent_name = data.get("agent_name")
+    new_status = data.get("new_status")
+    previous_status = data.get("previous_status")
+
+    logger.info(
+        f"[Aeonic] Agent status changed: {agent_name or agent_id} "
+        f"({agent_key or agent_id}) {previous_status} -> {new_status}"
+    )
+
+    if not any([agent_key, agent_id, agent_name]):
+        logger.warning("[Aeonic] Received agent.status_changed event without any identifiers")
+        return
+
+    # Update agent status cache for all identifiers provided
+    identifiers: list[str] = [x for x in (agent_key, agent_id, agent_name) if x]
+    for identifier in identifiers:
+        if new_status:
+            set_agent_status(identifier, new_status)
+
+    # Logging based on status
+    display_name = agent_name or agent_key or agent_id
+    if new_status in ("blocked", "quarantined"):
+        logger.warning(
+            f"[Aeonic] Agent {display_name} is {new_status} - stopping sample collection"
+        )
+    elif new_status == "healthy" and previous_status in ("blocked", "quarantined"):
+        # Agent was unblocked - resume sample collection
+        logger.info(f"[Aeonic] Agent {display_name} is now healthy - resuming sample collection")
+
+
+def _handle_assessment_job(data: Dict[str, Any]) -> None:
+    """Handle assessment job available event."""
+    from aeonic.core.assessment.executor import execute_assessment
+
+    job = {
+        "job_id": data.get("job_id"),
+        "agent_id": data.get("agent_id"),
+        "scenario_id": data.get("scenario_id"),
+        "scenario_steps": data.get("scenario_steps"),
+        "playbook": data.get("playbook"),
+        "agent_endpoint": data.get("agent_endpoint"),
+    }
+
+    logger.info(f"[Aeonic] Assessment job received via SSE: {job.get('job_id')}")
+
+    # Execute assessment asynchronously (non-blocking) in a new thread
+    threading.Thread(
+        target=lambda: _run_assessment_safely(job, execute_assessment),
+        name=f"AeonicAssessment-{job.get('job_id')}",
+        daemon=True,
+    ).start()
+
+
+def _run_assessment_safely(job: Dict[str, Any], executor_func: Any) -> None:
+    """Run assessment and log errors."""
+    try:
+        executor_func(job)
+    except Exception as e:
+        logger.error(f"[Aeonic] Assessment {job.get('job_id', 'unknown')} failed: {e}")
+
+
+def stop_sse_connection() -> None:
+    """Stop SSE connection."""
+    global _SSE_CONNECTION
+
+    _STOP_SSE.set()
+
+    if _SSE_CONNECTION and _SSE_CONNECTION.is_alive():
+        # Try to join briefly
+        _SSE_CONNECTION.join(timeout=2)
+        _SSE_CONNECTION = None

aeonic/core/assessment/__init__.py

@@ -0,0 +1,5 @@
+"""Assessment execution module."""
+
+from aeonic.core.assessment.executor import execute_assessment
+
+__all__ = ["execute_assessment"]