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

aeonic/__init__.py

@@ -0,0 +1,5 @@
+__version__ = "0.1.0"
+
+from aeonic.core.init import init
+
+__all__ = ["init", "__version__"]

aeonic/adapters/__init__.py

@@ -0,0 +1 @@
+"""Framework adapters for FastAPI and Django."""

aeonic/adapters/django.py

@@ -0,0 +1,275 @@
+from functools import wraps
+from time import time
+from typing import Any, Callable, Dict, Optional, cast
+
+from django.http import HttpRequest, HttpResponse
+
+from aeonic.core.buffer import SampleBuffer
+from aeonic.core.init import get_config
+from aeonic.core.sender import _is_2xx, emit_error_sample, enqueue
+from aeonic.core.types import AgentContext, Sample
+from aeonic.utils.safe_serialize import safe_serialize
+
+# In-memory buffers (per process)
+_BUFFERS: Dict[str, SampleBuffer] = {}
+_MAX_AGENT_BUFFERS = 100
+
+# Agent metadata storage for introspection
+# Maps function ID to agent context
+_AGENT_METADATA_MAP: Dict[int, AgentContext] = {}
+
+
+class agent_middleware:
+    """
+    Django middleware to observe responses and capture samples
+    ONLY for agent-marked routes.
+    """
+
+    def __init__(self, get_response: Callable[[HttpRequest], HttpResponse]):
+        self.get_response = get_response
+
+    def __call__(self, request: HttpRequest) -> HttpResponse:
+        response = self.get_response(request)
+
+        aeonic = getattr(request, "aeonic", None)
+        if not aeonic or not aeonic.get("is_agent"):
+            return response  # Ignore non-agent routes
+
+        # ⚠️ Skip sample collection for assessment calls
+        # Assessment calls are synthetic test scenarios and should not be tracked as production samples
+        # Django converts headers: X-Aeonic-Assessment -> HTTP_X_AEONIC_ASSESSMENT
+        assessment_header = request.META.get("HTTP_X_AEONIC_ASSESSMENT")
+        is_assessment = assessment_header is not None and str(assessment_header).lower() == "true"
+        if is_assessment:
+            # Log for debugging (optional)
+            import os
+
+            if os.environ.get("AEONIC_DEBUG"):
+                job_id = request.META.get("HTTP_X_AEONIC_JOB_ID")
+                job_id_str = str(job_id) if job_id else ""
+                print(
+                    f"[Aeonic] Skipping sample collection for assessment"
+                    f"{f' (job: {job_id_str})' if job_id_str else ''}"
+                )
+            return response  # Skip sample collection
+
+        agent: AgentContext = aeonic["agent"]
+        path: str = request.path
+        method: str = request.method
+
+        # ⚠️ Skip sample collection for blocked/quarantined agents
+        from aeonic.core.agent_status_cache import should_collect_samples
+
+        agent_key: str = agent.get("name") or f"{method}:{path}"
+        if not should_collect_samples(agent_key):
+            import os
+
+            if os.environ.get("AEONIC_DEBUG"):
+                print(f"[Aeonic] Skipping sample collection for blocked agent: {agent_key}")
+            return response
+
+        req_payload = safe_serialize(aeonic.get("req_payload"))
+
+        res_payload = None
+        try:
+            content_type = response.get("Content-Type", "")
+            if "json" in content_type.lower():
+                import json
+
+                try:
+                    content = response.content.decode("utf-8")
+                    if content:
+                        res_json = json.loads(content)
+                        res_payload = safe_serialize(res_json)
+                except Exception:
+                    pass
+        except Exception:
+            pass
+
+        status_code: int = getattr(response, "status_code", 200)
+        sample: Sample = {
+            "req": req_payload,
+            "res": res_payload,
+            "timestamp": int(time()),
+            "status_code": status_code,
+        }
+
+        config = get_config()
+        max_samples: int = config["max_samples"]
+
+        if _is_2xx(status_code):
+            buffer = _BUFFERS.get(agent_key)
+            if buffer is None:
+                if len(_BUFFERS) >= _MAX_AGENT_BUFFERS:
+                    first_key = next(iter(_BUFFERS))
+                    _BUFFERS.pop(first_key, None)
+
+                buffer = SampleBuffer(
+                    agent_key=agent_key,
+                    agent_context=agent,
+                    route={"path": path, "method": method},
+                    max_samples=max_samples,
+                )
+                _BUFFERS[agent_key] = buffer
+
+            ready: bool = buffer.add(sample)
+            if ready:
+                enqueue(
+                    {
+                        "api_key": config["api_key"],
+                        "service_name": config["service_name"],
+                        "agent": buffer.agent_context,
+                        "route": buffer.route,
+                        "samples": buffer.flush(),
+                    }
+                )
+        else:
+            emit_error_sample(
+                {
+                    "api_key": config["api_key"],
+                    "service_name": config["service_name"],
+                    "agent": agent,
+                    "route": {"path": path, "method": method},
+                    "samples": [sample],
+                }
+            )
+
+        return response
+
+
+def with_agent(agent: AgentContext) -> Callable[..., HttpResponse]:
+    """
+    Decorator to mark a Django view as an agent route.
+
+    This function:
+    - Stores agent metadata for later introspection
+    - Registers agent at declaration time (if SDK is initialized)
+    - Returns a decorator that attaches agent context to requests
+
+    Args:
+        agent: Agent context with name, type, model, etc.
+
+    Returns:
+        Decorator function for Django views
+    """
+
+    # Register agent at declaration time
+    _register_agent_declaration(agent, framework="django")
+
+    def decorator(view_func: Callable[..., HttpResponse]) -> Callable[..., HttpResponse]:
+        @wraps(view_func)
+        def wrapped_view(request: HttpRequest, *args: Any, **kwargs: Any) -> HttpResponse:
+            try:
+                req_payload = None
+
+                # For POST/PUT/PATCH, try to parse JSON body
+                if request.method in ["POST", "PUT", "PATCH"]:
+                    content_type = request.META.get("CONTENT_TYPE", "")
+                    if "json" in content_type.lower() and request.body:
+                        import json
+
+                        try:
+                            body_data = json.loads(request.body.decode("utf-8"))
+                            req_payload = safe_serialize(body_data)
+                        except Exception:
+                            pass
+                    elif request.body:
+                        # Non-JSON body, serialize as string
+                        req_payload = safe_serialize(request.body.decode("utf-8"))
+
+                # For GET requests, capture query parameters
+                elif request.method == "GET":
+                    if request.GET:
+                        req_payload = safe_serialize(dict(request.GET))
+                    else:
+                        # No query params, set empty dict instead of None
+                        req_payload = {}
+
+            except Exception:
+                req_payload = None
+
+            request.aeonic = {
+                "is_agent": True,
+                "agent": {**agent, "source": "manual"},
+                "req_payload": req_payload,
+            }
+
+            return view_func(request, *args, **kwargs)
+
+        # Store metadata for introspection (key by wrapped view ID)
+        agent_metadata: AgentContext = cast(AgentContext, {**agent, "source": "manual"})
+        _AGENT_METADATA_MAP[id(wrapped_view)] = agent_metadata
+
+        # Also store as function attribute for easier retrieval
+        wrapped_view.__aeonic_metadata__ = agent_metadata  # type: ignore[attr-defined]
+
+        return wrapped_view
+
+    return decorator
+
+
+def get_agent_metadata(func: Callable) -> Optional[AgentContext]:
+    """
+    Retrieve agent metadata from a function.
+
+    Used during route introspection to match views with agent metadata.
+
+    Args:
+        func: Function to check (view function)
+
+    Returns:
+        Agent context if found, None otherwise
+    """
+    # Try direct lookup first
+    metadata = _AGENT_METADATA_MAP.get(id(func))
+    if metadata:
+        return metadata
+
+    # Try checking if function has metadata attribute (for wrapped functions)
+    if hasattr(func, "__aeonic_metadata__"):
+        return getattr(func, "__aeonic_metadata__")  # type: ignore[no-any-return]
+
+    return None
+
+
+def _register_agent_declaration(agent: AgentContext, framework: str) -> None:
+    """
+    Register agent in the global registry at declaration time.
+
+    This function:
+    - Registers the agent with None route_path/http_method (will be updated by introspection)
+    - Does NOT emit inventory (introspection will do that after finding routes)
+    - Never throws
+
+    Args:
+        agent: Agent context
+        framework: Framework identifier
+    """
+    try:
+        from aeonic.core.agent_registry import register_agent
+        from aeonic.core.init import get_config
+
+        # Get config (may not be initialized yet)
+        try:
+            config = get_config()
+            service_name = config.get("service_name")
+        except RuntimeError:
+            service_name = None
+
+        # Register agent (with None values - introspection will update these)
+        register_agent(
+            name=agent.get("name", ""),
+            agent_type=agent.get("type"),
+            model=agent.get("model"),
+            route_path=None,  # Not available at declaration time - introspection will update
+            http_method=None,  # Not available at declaration time - introspection will update
+            framework=framework,
+            service_name=service_name,
+        )
+
+        # DO NOT emit inventory here - wait for introspection to complete
+        # Introspection will emit inventory after updating route_path and http_method
+
+    except Exception:
+        # Never throw - silent failure
+        pass

aeonic/adapters/fastapi.py

@@ -0,0 +1,337 @@
+from time import time
+from typing import Awaitable, Callable, Dict, Optional, cast
+
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+from starlette.responses import Response
+
+from aeonic.core.buffer import SampleBuffer
+from aeonic.core.init import get_config
+from aeonic.core.sender import _is_2xx, emit_error_sample, enqueue
+from aeonic.core.types import AgentContext, Sample
+from aeonic.utils.safe_serialize import safe_serialize
+
+# In-memory buffers (per process)
+_BUFFERS: Dict[str, SampleBuffer] = {}
+_MAX_AGENT_BUFFERS = 100
+
+# Agent metadata storage for introspection
+# Maps function ID to agent context
+_AGENT_METADATA_MAP: Dict[int, AgentContext] = {}
+# Also store by agent name as fallback (for cases where function IDs don't match)
+_AGENT_METADATA_BY_NAME: Dict[str, AgentContext] = {}
+# Store all metadata for fallback matching
+_ALL_METADATA: list = []  # List of (func_id, metadata) tuples
+
+
+class agent_middleware(BaseHTTPMiddleware):
+    """
+    Global middleware.
+    Observes responses and captures samples ONLY for agent routes.
+    """
+
+    async def dispatch(
+        self,
+        request: Request,
+        call_next: Callable[[Request], Awaitable[Response]],
+    ) -> Response:
+        # Request payload (best-effort) - must be captured BEFORE processing request
+        # Note: consuming request.stream() or .json() here would break the app if not careful.
+        # FastAPI request.json() caches the result, so it might be safe if the app also uses .json().
+        # However, checking `request.state` (populated by dependency) is safer.
+
+        # We process the request first
+        response: Response = await call_next(request)
+
+        aeonic = getattr(request.state, "aeonic", None)
+        if not aeonic or not aeonic.get("is_agent"):
+            return response  # Ignore non-agent routes
+
+        # ⚠️ Skip sample collection for assessment calls
+        # Assessment calls are synthetic test scenarios and should not be tracked as production samples
+        assessment_header = request.headers.get("x-aeonic-assessment")
+        is_assessment = assessment_header is not None and str(assessment_header).lower() == "true"
+        if is_assessment:
+            # Log for debugging (optional)
+            import os
+
+            if os.environ.get("AEONIC_DEBUG"):
+                job_id = request.headers.get("x-aeonic-job-id")
+                job_id_str = str(job_id) if job_id else ""
+                print(
+                    f"[Aeonic] Skipping sample collection for assessment"
+                    f"{f' (job: {job_id_str})' if job_id_str else ''}"
+                )
+            return response  # Skip sample collection
+
+        agent: AgentContext = aeonic["agent"]
+        path: str = request.url.path
+        method: str = request.method
+
+        # ⚠️ Skip sample collection for blocked/quarantined agents
+        from aeonic.core.agent_status_cache import should_collect_samples
+
+        agent_key: str = agent.get("name") or f"{method}:{path}"
+        if not should_collect_samples(agent_key):
+            if os.environ.get("AEONIC_DEBUG"):
+                print(f"[Aeonic] Skipping sample collection for blocked agent: {agent_key}")
+            return response
+
+        # Request payload (attempt to get from state first, which is populated by dependency)
+        req_payload = safe_serialize(aeonic.get("req_payload"))
+
+        # Response payload capture
+        res_payload = None
+        response_body = b""
+
+        # We need to capture the response body.
+        # Strategy: Try .body first (works for non-streaming responses), then body_iterator
+        try:
+            # First, try to get the body directly (for pre-rendered responses like JSONResponse)
+            if hasattr(response, "body") and response.body:
+                response_body = response.body
+            elif hasattr(response, "body_iterator"):
+                # Fallback: iterate through body_iterator (for streaming responses)
+                # CAUTION: This consumes the iterator. We must reconstruct the response.
+                chunks = []
+                async for chunk in response.body_iterator:
+                    chunks.append(chunk)
+                response_body = b"".join(chunks)
+
+                # Reconstruct the response with fresh content
+                response = Response(
+                    content=response_body,
+                    status_code=response.status_code,
+                    headers=dict(response.headers),
+                    media_type=response.media_type,
+                )
+
+            # Parse JSON if applicable
+            # Check Content-Type header instead of media_type (which may be None)
+            content_type = response.headers.get("content-type", "")
+
+            if "json" in content_type.lower() and response_body:
+                import json
+
+                try:
+                    body_json = json.loads(response_body)
+                    res_payload = safe_serialize(body_json)
+                except Exception:
+                    pass
+
+        except Exception:
+            # Fallback if anything goes wrong with body capture
+            pass
+
+        status_code: int = getattr(response, "status_code", 200)
+        sample: Sample = {
+            "req": req_payload,
+            "res": res_payload,
+            "timestamp": int(time()),
+            "status_code": status_code,
+        }
+
+        config = get_config()
+        max_samples: int = config["max_samples"]
+
+        if _is_2xx(status_code):
+            buffer = _BUFFERS.get(agent_key)
+            if buffer is None:
+                if len(_BUFFERS) >= _MAX_AGENT_BUFFERS:
+                    first_key = next(iter(_BUFFERS))
+                    _BUFFERS.pop(first_key, None)
+
+                buffer = SampleBuffer(
+                    agent_key=agent_key,
+                    agent_context=agent,
+                    route={"path": path, "method": method},
+                    max_samples=max_samples,
+                )
+                _BUFFERS[agent_key] = buffer
+
+            ready: bool = buffer.add(sample)
+            if ready:
+                enqueue(
+                    {
+                        "api_key": config["api_key"],
+                        "service_name": config["service_name"],
+                        "agent": buffer.agent_context,
+                        "route": buffer.route,
+                        "samples": buffer.flush(),
+                    }
+                )
+        else:
+            emit_error_sample(
+                {
+                    "api_key": config["api_key"],
+                    "service_name": config["service_name"],
+                    "agent": agent,
+                    "route": {"path": path, "method": method},
+                    "samples": [sample],
+                }
+            )
+
+        return response
+
+
+def with_agent(agent: AgentContext) -> Callable[[Request], Awaitable[None]]:
+    """
+    Dependency to mark a route as an agent route.
+
+    This function:
+    - Stores agent metadata for later introspection
+    - Registers agent at declaration time (if SDK is initialized)
+    - Returns a FastAPI dependency that attaches agent context to requests
+
+    Args:
+        agent: Agent context with name, type, model, etc.
+
+    Returns:
+        FastAPI dependency function
+    """
+
+    # Register agent at declaration time
+    _register_agent_declaration(agent, framework="fastapi")
+
+    async def dependency(request: Request) -> None:
+        try:
+            body = await request.json()
+        except Exception:
+            body = None
+
+        request.state.aeonic = {
+            "is_agent": True,
+            "agent": {**agent, "source": "manual"},
+            "req_payload": body,
+        }
+
+    # Store metadata for introspection (key by dependency function ID)
+    agent_metadata: AgentContext = cast(AgentContext, {**agent, "source": "manual"})
+    func_id = id(dependency)
+    _AGENT_METADATA_MAP[func_id] = agent_metadata
+
+    # Also store as function attribute for easier retrieval (handles wrapped functions)
+    dependency.__aeonic_metadata__ = agent_metadata  # type: ignore[attr-defined]
+
+    # Also store by agent name as fallback
+    agent_name = agent.get("name")
+    if agent_name:
+        _AGENT_METADATA_BY_NAME[agent_name] = agent_metadata
+
+    # Store in list for fallback matching
+    _ALL_METADATA.append((func_id, agent_metadata))
+
+    return dependency
+
+
+def get_agent_metadata(func: Callable) -> Optional[AgentContext]:
+    """
+    Retrieve agent metadata from a function.
+
+    Used during route introspection to match routes with agent metadata.
+
+    Args:
+        func: Function to check (route handler or dependency)
+
+    Returns:
+        Agent context if found, None otherwise
+    """
+    if func is None:
+        return None
+
+    # Try direct lookup first
+    func_id = id(func)
+    metadata = _AGENT_METADATA_MAP.get(func_id)
+    if metadata:
+        return metadata
+
+    # Try checking if function has metadata attribute (for wrapped functions)
+    if hasattr(func, "__aeonic_metadata__"):
+        return getattr(func, "__aeonic_metadata__")  # type: ignore[no-any-return]
+
+    # Try checking closure variables (for nested functions)
+    if hasattr(func, "__closure__") and func.__closure__:
+        # Check if any closure variable contains our metadata
+        for cell in func.__closure__:
+            try:
+                value = cell.cell_contents
+                if isinstance(value, dict) and value.get("source") == "manual":
+                    return cast(AgentContext, value)
+            except (ValueError, AttributeError):
+                continue
+
+    # For FastAPI dependencies wrapped in Depends, try to get the underlying callable
+    # FastAPI's DependencyInfo has a 'call' attribute
+    if hasattr(func, "call"):
+        return get_agent_metadata(func.call)
+
+    # Try checking __wrapped__ for decorated functions
+    if hasattr(func, "__wrapped__"):
+        wrapped_metadata = get_agent_metadata(func.__wrapped__)
+        if wrapped_metadata:
+            return wrapped_metadata
+
+    # Last resort: try to match by inspecting function closure/code
+    # FastAPI might wrap the dependency, so function IDs won't match
+    # But the closure might contain our agent context
+    if hasattr(func, "__closure__") and func.__closure__:
+        for cell in func.__closure__:
+            try:
+                value = cell.cell_contents
+                # Check if this is our dependency function
+                if callable(value) and hasattr(value, "__aeonic_metadata__"):
+                    return getattr(value, "__aeonic_metadata__")  # type: ignore[no-any-return]
+                # Check if closure contains agent context directly
+                if isinstance(value, dict) and value.get("source") == "manual":
+                    return cast(AgentContext, value)
+            except (ValueError, AttributeError):
+                continue
+
+    # Final fallback: check all stored metadata to see if any function matches
+    # This handles cases where FastAPI creates a wrapper
+    # Note: We can't easily match by function ID here, so we skip this fallback
+
+    return None
+
+
+def _register_agent_declaration(agent: AgentContext, framework: str) -> None:
+    """
+    Register agent in the global registry at declaration time.
+
+    This function:
+    - Registers the agent with None route_path/http_method (will be updated by introspection)
+    - Does NOT emit inventory (introspection will do that after finding routes)
+    - Never throws
+
+    Args:
+        agent: Agent context
+        framework: Framework identifier
+    """
+    try:
+        from aeonic.core.agent_registry import register_agent
+        from aeonic.core.init import get_config
+
+        # Get config (may not be initialized yet)
+        try:
+            config = get_config()
+            service_name = config.get("service_name")
+        except RuntimeError:
+            service_name = None
+
+        # Register agent (with None values - introspection will update these)
+        register_agent(
+            name=agent.get("name", ""),
+            agent_type=agent.get("type"),
+            model=agent.get("model"),
+            route_path=None,  # Not available at declaration time - introspection will update
+            http_method=None,  # Not available at declaration time - introspection will update
+            framework=framework,
+            service_name=service_name,
+        )
+
+        # DO NOT emit inventory here - wait for introspection to complete
+        # Introspection will emit inventory after updating route_path and http_method
+
+    except Exception:
+        # Never throw - silent failure
+        pass

aeonic/core/__init__.py

@@ -0,0 +1,4 @@
+from aeonic.core.buffer import SampleBuffer
+from aeonic.core.init import get_config, init
+
+__all__ = ["init", "get_config", "SampleBuffer"]