aegra-api 0.1.0__py3-none-any.whl

aegra_api/core/auth_handlers.py

@@ -0,0 +1,248 @@
+"""Authorization handler support for @auth.on.* decorators.
+
+This module provides integration with authorization handlers,
+allowing users to define fine-grained access control rules using @auth.on.*
+decorators in their auth.py files.
+"""
+
+from typing import Any
+
+from fastapi import HTTPException
+from langgraph_sdk import Auth
+from langgraph_sdk.auth.types import AuthContext as LangGraphAuthContext
+
+from aegra_api.core.auth_middleware import get_auth_instance
+from aegra_api.models.auth import User
+
+
+class AuthContextWrapper:
+    """Wrapper to convert Aegra User model to AuthContext.
+
+    AuthContext expects a BaseUser-compatible object. Our User model
+    implements the BaseUser protocol (identity, permissions, display_name, __getitem__),
+    so we can use it directly after ensuring compatibility.
+    """
+
+    def __init__(
+        self,
+        user: User,
+        resource: str,
+        action: str,
+    ):
+        """Initialize auth context wrapper.
+
+        Args:
+            user: Authenticated user from Aegra's User model
+            resource: Resource being accessed (e.g., "threads", "assistants")
+            action: Action being performed (e.g., "create", "read", "update")
+        """
+        self.user = user
+        self.resource = resource
+        self.action = action
+        self.permissions = user.permissions or []
+
+    def to_langgraph_context(self) -> LangGraphAuthContext:
+        """Convert to LangGraph AuthContext.
+
+        Our User model implements the BaseUser protocol (identity, permissions,
+        display_name, __getitem__, __contains__, __iter__), so it's compatible
+        with LangGraph's AuthContext.
+
+        Returns:
+            AuthContext instance compatible with @auth.on handlers
+        """
+        return LangGraphAuthContext(
+            user=self.user,  # Our User model implements BaseUser protocol
+            resource=self.resource,  # type: ignore
+            action=self.action,  # type: ignore
+            permissions=self.permissions,
+        )
+
+
+async def handle_event(
+    ctx: AuthContextWrapper | None,
+    value: dict[str, Any],
+) -> dict[str, Any] | None:
+    """Call the appropriate @auth.on.* handler for authorization.
+
+    This function resolves the most specific handler for the given resource
+    and action, calls it with the auth context and value, and interprets
+    the result.
+
+    **Default Behavior (Non-Interruptive):**
+    - If no auth is configured → allows by default (returns None)
+    - If no handlers are defined → allows by default (returns None)
+    - If handler returns None/True → allows (returns None)
+    - If handler returns dict → allows with filters applied (returns dict)
+    - **Only interrupts if handler returns False or raises exception**
+
+    This ensures that developers using raw Aegra without custom authorization
+    handlers will have a working system out-of-the-box. Handlers are purely
+    additive - they can inject metadata, apply filters, or deny access, but
+    they don't interrupt the flow unless explicitly configured to do so.
+
+    Handler resolution priority (most specific first):
+    1. Resource+action specific (e.g., "threads", "create")
+    2. Resource-specific (e.g., "threads", "*")
+    3. Action-specific (e.g., "*", "create")
+    4. Global handler ("*", "*")
+
+    Args:
+        ctx: Auth context wrapper with user, resource, and action
+        value: The data being authorized (request body, search filters, etc.)
+               This dict may be modified by the handler (e.g., injecting metadata)
+
+    Returns:
+        None: Request allowed, no filters to apply
+        dict: Filter dict to apply to queries (e.g., {"user_id": "123"})
+
+    Raises:
+        HTTPException(403): If handler returns False or raises AssertionError
+        HTTPException(500): If handler raises unexpected exception or returns invalid type
+    """
+    if ctx is None:
+        # No auth context means no authorization check needed
+        # This allows the request to proceed normally
+        return None
+
+    auth = get_auth_instance()
+    if auth is None:
+        # No auth configured, allow by default
+        # This ensures raw Aegra works out-of-the-box without interruption
+        return None
+
+    # Convert to AuthContext
+    auth_ctx = ctx.to_langgraph_context()
+
+    # Find the most specific handler
+    handler = _get_handler(auth, auth_ctx.resource, auth_ctx.action)
+    if handler is None:
+        # No handler for this resource/action, allow by default
+        # Developers can use Aegra without defining handlers - it won't break
+        return None
+
+    try:
+        # Call the handler with context and value
+        result = await handler(ctx=auth_ctx, value=value)
+    except Auth.exceptions.HTTPException as e:
+        # Handler raised HTTP exception, convert to FastAPI HTTPException
+        raise HTTPException(
+            status_code=e.status_code,
+            detail=e.detail,
+            headers=dict(e.headers) if hasattr(e, "headers") and e.headers else None,
+        ) from e
+    except AssertionError as e:
+        # Handler used assert for authorization check
+        raise HTTPException(status_code=403, detail=str(e)) from e
+    except Exception as e:
+        # Unexpected error in handler
+        raise HTTPException(status_code=500, detail=f"Authorization error: {str(e)}") from e
+
+    # Interpret handler result
+    if result in (None, True):
+        # Allow request, no filters
+        return None
+
+    if result is False:
+        # Deny request
+        raise HTTPException(status_code=403, detail="Forbidden")
+
+    if isinstance(result, dict):
+        # Return filter dict to apply to queries
+        return result
+
+    # Invalid return type
+    raise HTTPException(
+        status_code=500,
+        detail=f"Auth handler returned invalid type: {type(result)}. Expected dict, None, True, or False.",
+    )
+
+
+def _get_handler(
+    auth: Auth,
+    resource: str,
+    action: str,
+) -> Any | None:
+    """Find the most specific handler for resource+action.
+
+    Handler resolution follows this priority order (most specific first):
+    1. (resource, action) - e.g., ("threads", "create")
+    2. (resource, "*") - e.g., ("threads", "*")
+    3. ("*", action) - e.g., ("*", "create")
+    4. ("*", "*") - global handler
+
+    Args:
+        auth: Auth instance with registered handlers
+        resource: Resource name (e.g., "threads", "assistants")
+        action: Action name (e.g., "create", "read", "update")
+
+    Returns:
+        Handler function or None if no handler found
+    """
+    # Check cache first
+    key = (resource, action)
+    if key in auth._handler_cache:
+        return auth._handler_cache[key]
+
+    # Priority order (most specific first)
+    keys = [
+        (resource, action),  # Most specific: exact resource+action
+        (resource, "*"),  # Resource-specific: all actions on resource
+        ("*", action),  # Action-specific: all resources for action
+        ("*", "*"),  # Global: all resources and actions
+    ]
+
+    # Find first matching handler
+    for check_key in keys:
+        if check_key in auth._handlers and auth._handlers[check_key]:
+            # Get the last registered handler (most recent wins)
+            handler = auth._handlers[check_key][-1]
+            # Cache the result
+            auth._handler_cache[key] = handler
+            return handler
+
+    # Check global handlers (fallback for backward compatibility)
+    if auth._global_handlers:
+        handler = auth._global_handlers[-1]
+        auth._handler_cache[key] = handler
+        return handler
+
+    return None
+
+
+def build_auth_context(
+    user: User,
+    resource: str,
+    action: str,
+) -> AuthContextWrapper:
+    """Build AuthContextWrapper from user and operation info.
+
+    This is a convenience function to create an AuthContextWrapper with
+    the correct resource and action for a given operation.
+
+    Args:
+        user: Authenticated user from require_auth dependency
+        resource: Resource being accessed (e.g., "threads", "assistants")
+        action: Action being performed (e.g., "create", "read", "update")
+
+    Returns:
+        AuthContextWrapper instance ready to pass to handle_event()
+
+    Example:
+        ```python
+        @router.post("/threads")
+        async def create_thread(
+            request: ThreadCreate,
+            user: User = Depends(require_auth),
+        ):
+            ctx = build_auth_context(user, "threads", "create")
+            value = request.model_dump()
+            filters = await handle_event(ctx, value)
+            # Use filters or modified value...
+        ```
+    """
+    return AuthContextWrapper(
+        user=user,
+        resource=resource,
+        action=action,
+    )

aegra_api/core/auth_middleware.py

@@ -0,0 +1,331 @@
+"""
+Authentication middleware integration for Aegra.
+
+This module integrates authentication system with FastAPI
+using Starlette's AuthenticationMiddleware.
+"""
+
+import functools
+import importlib
+import importlib.util
+import sys
+from pathlib import Path
+from typing import Any
+
+import structlog
+from langgraph_sdk import Auth
+from langgraph_sdk.auth.types import MinimalUserDict
+from starlette.authentication import (
+    AuthCredentials,
+    AuthenticationBackend,
+    AuthenticationError,
+    BaseUser,
+)
+from starlette.requests import HTTPConnection
+from starlette.responses import JSONResponse
+
+from aegra_api.config import get_config_dir, load_auth_config
+from aegra_api.models.errors import AgentProtocolError
+from aegra_api.settings import settings
+
+logger = structlog.getLogger(__name__)
+
+
+class LangGraphUser(BaseUser):
+    """
+    User wrapper that implements Starlette's BaseUser interface
+    while preserving auth data.
+    """
+
+    def __init__(self, user_data: Auth.types.MinimalUserDict):
+        self._user_data = user_data
+
+    @property
+    def identity(self) -> str:
+        return self._user_data["identity"]
+
+    @property
+    def is_authenticated(self) -> bool:
+        return self._user_data.get("is_authenticated", True)
+
+    @property
+    def display_name(self) -> str:
+        return self._user_data.get("display_name", self.identity)
+
+    def __getattr__(self, name: str) -> Any:
+        """Allow access to any additional fields from auth data"""
+        if name in self._user_data:
+            return self._user_data[name]
+        raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'")
+
+    def to_dict(self) -> MinimalUserDict:
+        """Return the underlying user data dict"""
+        return self._user_data.copy()
+
+
+class LangGraphAuthBackend(AuthenticationBackend):
+    """
+    Authentication backend that uses the auth system.
+
+    This bridges @auth.authenticate handlers with
+    Starlette's AuthenticationMiddleware.
+    """
+
+    def __init__(self) -> None:
+        self.auth_instance = self._load_auth_instance()
+
+    def _load_auth_instance(self) -> Auth | None:
+        """Load the auth instance from config or fallback to hardcoded candidates.
+
+        Resolution order:
+        1. Load from aegra.json auth.path config
+        2. If no auth file found, returns None (noop handled directly in authenticate())
+
+        Returns:
+            Auth instance or None if not found (noop handled in authenticate() method)
+        """
+        # 1. Try loading from config
+        try:
+            auth_config = load_auth_config()
+            if auth_config and "path" in auth_config:
+                auth_path = auth_config["path"]
+                logger.info(f"Loading auth from config path: {auth_path}")
+                auth_instance = self._load_from_path(auth_path)
+                if auth_instance:
+                    return auth_instance
+                logger.warning(f"Failed to load auth from config path: {auth_path}")
+        except Exception as e:
+            logger.warning(f"Error loading auth config: {e}")
+
+        logger.debug("No auth instance found from config")
+        return None
+
+    def _load_from_path(self, path: str) -> Auth | None:
+        """Load auth instance from path in format './file.py:var' or 'module:var'.
+
+        Relative paths are resolved from the config file directory.
+
+        Args:
+            path: Import path in format './file.py:variable' or 'module.path:variable'
+
+        Returns:
+            Auth instance or None if loading fails
+        """
+        if ":" not in path:
+            logger.error(f"Invalid auth path format (missing ':'): {path}")
+            return None
+
+        module_path, var_name = path.rsplit(":", 1)
+
+        # Handle file path format: ./file.py or ./path/to/file.py or ../file.py
+        is_file_path = module_path.endswith(".py") or module_path.startswith("./") or module_path.startswith("../")
+        if is_file_path:
+            file_path = Path(module_path)
+
+            # Resolve relative paths from config directory
+            if not file_path.is_absolute():
+                config_dir = get_config_dir()
+                if config_dir:
+                    file_path = (config_dir / file_path).resolve()
+                else:
+                    # Fallback to CWD if no config found
+                    file_path = (Path.cwd() / file_path).resolve()
+
+            return self._load_from_file(file_path, var_name)
+
+        # Handle module format: module.path
+        return self._load_from_module(module_path, var_name)
+
+    def _load_from_file(self, file_path: Path, var_name: str) -> Auth | None:
+        """Load auth instance from a file path.
+
+        Args:
+            file_path: Path to the Python file
+            var_name: Name of the variable to load
+
+        Returns:
+            Auth instance or None if loading fails
+        """
+        try:
+            if not file_path.exists():
+                logger.warning(f"Auth file not found: {file_path}")
+                return None
+
+            if not file_path.is_file():
+                logger.warning(f"Auth path is not a file: {file_path} (is directory: {file_path.is_dir()})")
+                return None
+
+            # Create a unique module name based on the file path
+            module_name = f"auth_module_{file_path.stem}"
+
+            spec = importlib.util.spec_from_file_location(module_name, str(file_path))
+            if spec is None or spec.loader is None:
+                logger.error(f"Could not load auth module from {file_path}")
+                return None
+
+            auth_module = importlib.util.module_from_spec(spec)
+            sys.modules[module_name] = auth_module
+            spec.loader.exec_module(auth_module)
+
+            auth_instance = getattr(auth_module, var_name, None)
+            if not isinstance(auth_instance, Auth):
+                logger.error(f"Variable '{var_name}' in {file_path} is not an Auth instance")
+                return None
+
+            logger.info(f"Successfully loaded auth instance from {file_path}:{var_name}")
+            return auth_instance
+
+        except Exception as e:
+            logger.error(f"Error loading auth from {file_path}: {e}", exc_info=True)
+            return None
+
+    def _load_from_module(self, module_path: str, var_name: str) -> Auth | None:
+        """Load auth instance from an installed module.
+
+        Args:
+            module_path: Dotted module path (e.g., 'mypackage.auth')
+            var_name: Name of the variable to load
+
+        Returns:
+            Auth instance or None if loading fails
+        """
+        try:
+            module = importlib.import_module(module_path)
+            auth_instance = getattr(module, var_name, None)
+
+            if not isinstance(auth_instance, Auth):
+                logger.error(f"Variable '{var_name}' in module {module_path} is not an Auth instance")
+                return None
+
+            logger.info(f"Successfully loaded auth instance from {module_path}:{var_name}")
+            return auth_instance
+
+        except ImportError as e:
+            logger.error(f"Could not import module {module_path}: {e}")
+            return None
+        except Exception as e:
+            logger.error(f"Error loading auth from {module_path}: {e}", exc_info=True)
+            return None
+
+    async def authenticate(self, conn: HTTPConnection) -> tuple[AuthCredentials, BaseUser] | None:
+        """
+        Authenticate request using the configured auth system.
+
+        Args:
+            conn: HTTP connection containing request headers
+
+        Returns:
+            Tuple of (credentials, user) if authenticated, None otherwise
+
+        Raises:
+            AuthenticationError: If authentication fails
+        """
+        # Handle noop auth when no auth instance is configured
+        # Default to noop (anonymous) authentication when no auth file is found,
+        # regardless of AUTH_TYPE setting. This ensures the server works out-of-the-box.
+        if self.auth_instance is None:
+            logger.debug("No auth file configured, defaulting to noop (anonymous) authentication")
+            # Return anonymous user when no auth is configured
+            user_data: Auth.types.MinimalUserDict = {
+                "identity": "anonymous",
+                "display_name": "Anonymous User",
+                "is_authenticated": True,
+            }
+            credentials = AuthCredentials([])
+            user = LangGraphUser(user_data)
+            return credentials, user
+
+        if self.auth_instance._authenticate_handler is None:
+            logger.warning("No authenticate handler configured, skipping authentication")
+            return None
+
+        try:
+            # Convert headers to dict format expected by auth handlers
+            headers = {
+                key.decode() if isinstance(key, bytes) else key: value.decode() if isinstance(value, bytes) else value
+                for key, value in conn.headers.items()
+            }
+
+            # Call the authenticate handler
+            user_data = await self.auth_instance._authenticate_handler(headers)
+
+            if not user_data or not isinstance(user_data, dict):
+                raise AuthenticationError("Invalid user data returned from auth handler")
+
+            if "identity" not in user_data:
+                raise AuthenticationError("Auth handler must return 'identity' field")
+
+            # Extract permissions for credentials
+            permissions = user_data.get("permissions", [])
+            if isinstance(permissions, str):
+                permissions = [permissions]
+
+            # Create Starlette-compatible user and credentials
+            credentials = AuthCredentials(permissions)
+            user = LangGraphUser(user_data)
+
+            logger.debug(f"Successfully authenticated user: {user.identity}")
+            return credentials, user
+
+        except Auth.exceptions.HTTPException as e:
+            logger.warning(f"Authentication failed: {e.detail}")
+            raise AuthenticationError(e.detail) from e
+
+        except Exception as e:
+            logger.error(f"Unexpected error during authentication: {e}", exc_info=True)
+            raise AuthenticationError("Authentication system error") from e
+
+
+def get_auth_backend() -> AuthenticationBackend:
+    """
+    Get authentication backend based on AUTH_TYPE environment variable.
+
+    Returns:
+        AuthenticationBackend instance
+    """
+    auth_type = settings.app.AUTH_TYPE
+
+    if auth_type in ["noop", "custom"]:
+        logger.debug(f"Using auth backend with type: {auth_type}")
+        return LangGraphAuthBackend()
+    else:
+        logger.warning(f"Unknown AUTH_TYPE: {auth_type}, using noop")
+        return LangGraphAuthBackend()
+
+
+def on_auth_error(conn: HTTPConnection, exc: AuthenticationError) -> JSONResponse:
+    """
+    Handle authentication errors in Agent Protocol format.
+
+    Args:
+        conn: HTTP connection
+        exc: Authentication error
+
+    Returns:
+        JSON response with Agent Protocol error format
+    """
+    logger.warning(f"Authentication error for {conn.url}: {exc}")
+
+    return JSONResponse(
+        status_code=401,
+        content=AgentProtocolError(
+            error="unauthorized",
+            message=str(exc),
+            details={"authentication_required": True},
+        ).model_dump(),
+    )
+
+
+@functools.lru_cache(maxsize=1)
+def get_auth_instance() -> Auth | None:
+    """Get cached Auth instance for use by other modules.
+
+    Uses LRU cache to ensure only one Auth instance is loaded per process.
+    This allows other modules to access the same Auth instance used by
+    the middleware without re-loading it.
+
+    Returns:
+        Auth instance or None if not configured/found
+    """
+    backend = LangGraphAuthBackend()
+    return backend.auth_instance

aegra_api/core/database.py

@@ -0,0 +1,123 @@
+"""Database manager with LangGraph integration"""
+
+import structlog
+from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+from langgraph.store.postgres.aio import AsyncPostgresStore
+from psycopg.rows import dict_row
+from psycopg_pool import AsyncConnectionPool
+from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
+
+from aegra_api.config import load_store_config
+from aegra_api.settings import settings
+
+logger = structlog.get_logger(__name__)
+
+
+class DatabaseManager:
+    """Manages database connections and LangGraph persistence components"""
+
+    def __init__(self) -> None:
+        self.engine: AsyncEngine | None = None
+
+        # Shared pool for LangGraph components (Checkpointer + Store)
+        self.lg_pool: AsyncConnectionPool | None = None
+        self._checkpointer: AsyncPostgresSaver | None = None
+        self._store: AsyncPostgresStore | None = None
+        self._database_url = settings.db.database_url
+
+    async def initialize(self) -> None:
+        """Initialize database connections and LangGraph components"""
+        # Idempotency check: if already initialized, do nothing
+        if self.engine:
+            return
+
+        # 1. SQLAlchemy Engine (app metadata, uses asyncpg)
+        # We strictly limit this pool because the main load
+        # is handled by LangGraph components.
+        self.engine = create_async_engine(
+            self._database_url,
+            pool_size=settings.pool.SQLALCHEMY_POOL_SIZE,
+            max_overflow=settings.pool.SQLALCHEMY_MAX_OVERFLOW,
+            pool_pre_ping=True,
+            echo=settings.db.DB_ECHO_LOG,
+        )
+
+        lg_max = settings.pool.LANGGRAPH_MAX_POOL_SIZE
+        lg_kwargs = {
+            "autocommit": True,
+            "prepare_threshold": 0,  # Optimization for PgBouncer/Kubernetes compatibility
+            "row_factory": dict_row,  # LangGraph requires dictionary rows, not tuples
+        }
+
+        # Create a single shared pool.
+        # 'open=False' is important to avoid RuntimeWarning; we open it explicitly below.
+        self.lg_pool = AsyncConnectionPool(
+            conninfo=settings.db.database_url_sync,
+            min_size=settings.pool.LANGGRAPH_MIN_POOL_SIZE,
+            max_size=lg_max,
+            open=False,
+            kwargs=lg_kwargs,
+            check=AsyncConnectionPool.check_connection,
+        )
+
+        # Explicitly open the pool
+        await self.lg_pool.open()
+
+        # 2. Initialize LangGraph components using the shared pool
+        # Passing 'conn=self.lg_pool' prevents components from creating their own pools.
+
+        logger.info(f"Initializing LangGraph components with shared pool (max {lg_max} conns)...")
+
+        self._checkpointer = AsyncPostgresSaver(conn=self.lg_pool)
+        await self._checkpointer.setup()  # Ensure tables exist
+
+        # Load store configuration for semantic search (if configured)
+        store_config = load_store_config()
+        index_config = store_config.get("index") if store_config else None
+
+        self._store = AsyncPostgresStore(conn=self.lg_pool, index=index_config)
+        await self._store.setup()  # Ensure tables exist
+
+        if index_config:
+            embed_model = index_config.get("embed", "unknown")
+            logger.info(f"Semantic store enabled with embeddings: {embed_model}")
+
+        logger.info("✅ Database and LangGraph components initialized")
+
+    async def close(self) -> None:
+        """Close database connections"""
+        # Close SQLAlchemy engine
+        if self.engine:
+            await self.engine.dispose()
+            self.engine = None
+
+        # Close shared LangGraph pool
+        if self.lg_pool:
+            await self.lg_pool.close()
+            self.lg_pool = None
+            self._checkpointer = None
+            self._store = None
+
+        logger.info("✅ Database connections closed")
+
+    def get_checkpointer(self) -> AsyncPostgresSaver:
+        """Return the live AsyncPostgresSaver instance."""
+        if self._checkpointer is None:
+            raise RuntimeError("Database not initialized")
+        return self._checkpointer
+
+    def get_store(self) -> AsyncPostgresStore:
+        """Return the live AsyncPostgresStore instance."""
+        if self._store is None:
+            raise RuntimeError("Database not initialized")
+        return self._store
+
+    def get_engine(self) -> AsyncEngine:
+        """Get the SQLAlchemy engine for metadata tables"""
+        if not self.engine:
+            raise RuntimeError("Database not initialized")
+        return self.engine
+
+
+# Global database manager instance
+db_manager = DatabaseManager()