kubiya-control-plane-api 0.9.15__py3-none-any.whl

control_plane_api/app/middleware/rate_limiting.py

@@ -0,0 +1,384 @@
+"""
+Rate limiting middleware using token bucket algorithm.
+
+Provides configurable rate limiting per client IP or user.
+"""
+
+from fastapi import Request, Response, HTTPException, status
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.types import ASGIApp
+from typing import Dict, Optional, Tuple, Any
+from datetime import datetime, timedelta
+import time
+import asyncio
+import structlog
+import hashlib
+from control_plane_api.app.exceptions import RateLimitError
+
+logger = structlog.get_logger()
+
+
+class TokenBucket:
+    """
+    Token bucket implementation for rate limiting.
+    
+    Each bucket starts with a capacity of tokens.
+    Tokens are consumed when requests are made.
+    Tokens are refilled at a constant rate.
+    """
+    
+    def __init__(
+        self,
+        capacity: int,
+        refill_rate: float,
+        refill_period: float = 60.0,
+    ):
+        """
+        Initialize token bucket.
+        
+        Args:
+            capacity: Maximum number of tokens in bucket
+            refill_rate: Number of tokens to add per period
+            refill_period: Period in seconds for refilling tokens
+        """
+        self.capacity = capacity
+        self.refill_rate = refill_rate
+        self.refill_period = refill_period
+        self.tokens = capacity
+        self.last_refill = time.time()
+        self.lock = asyncio.Lock()
+    
+    async def consume(self, tokens: int = 1) -> Tuple[bool, Dict[str, Any]]:
+        """
+        Try to consume tokens from the bucket.
+        
+        Args:
+            tokens: Number of tokens to consume
+            
+        Returns:
+            Tuple of (success, info_dict)
+        """
+        async with self.lock:
+            now = time.time()
+            
+            # Refill tokens based on time elapsed
+            time_elapsed = now - self.last_refill
+            tokens_to_add = (time_elapsed / self.refill_period) * self.refill_rate
+            
+            if tokens_to_add > 0:
+                self.tokens = min(self.capacity, self.tokens + tokens_to_add)
+                self.last_refill = now
+            
+            # Check if we have enough tokens
+            if self.tokens >= tokens:
+                self.tokens -= tokens
+                success = True
+                retry_after = None
+            else:
+                success = False
+                # Calculate when enough tokens will be available
+                tokens_needed = tokens - self.tokens
+                time_to_wait = (tokens_needed / self.refill_rate) * self.refill_period
+                retry_after = int(time_to_wait) + 1
+            
+            info = {
+                "limit": self.capacity,
+                "remaining": int(self.tokens),
+                "reset": int(now + self.refill_period),
+                "retry_after": retry_after,
+            }
+            
+            return success, info
+
+
+class RateLimiter:
+    """
+    Rate limiter managing multiple token buckets for different clients.
+    """
+    
+    def __init__(
+        self,
+        requests_per_minute: int = 60,
+        burst_size: Optional[int] = None,
+        cleanup_interval: int = 300,
+    ):
+        """
+        Initialize rate limiter.
+        
+        Args:
+            requests_per_minute: Number of requests allowed per minute
+            burst_size: Maximum burst size (defaults to requests_per_minute // 4)
+            cleanup_interval: Interval in seconds to clean up old buckets
+        """
+        self.requests_per_minute = requests_per_minute
+        self.burst_size = burst_size or max(10, requests_per_minute // 4)
+        self.buckets: Dict[str, TokenBucket] = {}
+        self.last_cleanup = time.time()
+        self.cleanup_interval = cleanup_interval
+        self.lock = asyncio.Lock()
+    
+    async def check_rate_limit(self, identifier: str) -> Tuple[bool, Dict[str, Any]]:
+        """
+        Check if request is within rate limit.
+        
+        Args:
+            identifier: Client identifier (IP, user ID, etc.)
+            
+        Returns:
+            Tuple of (allowed, headers_dict)
+        """
+        # Clean up old buckets periodically
+        await self._cleanup_buckets()
+        
+        # Get or create bucket for this identifier
+        bucket = await self._get_or_create_bucket(identifier)
+        
+        # Try to consume a token
+        allowed, info = await bucket.consume()
+        
+        # Build rate limit headers
+        headers = {
+            "X-RateLimit-Limit": str(info["limit"]),
+            "X-RateLimit-Remaining": str(info["remaining"]),
+            "X-RateLimit-Reset": str(info["reset"]),
+        }
+        
+        if not allowed and info.get("retry_after"):
+            headers["Retry-After"] = str(info["retry_after"])
+        
+        return allowed, headers
+    
+    async def _get_or_create_bucket(self, identifier: str) -> TokenBucket:
+        """Get existing bucket or create new one."""
+        async with self.lock:
+            if identifier not in self.buckets:
+                self.buckets[identifier] = TokenBucket(
+                    capacity=self.burst_size,
+                    refill_rate=self.requests_per_minute,
+                    refill_period=60.0,
+                )
+            return self.buckets[identifier]
+    
+    async def _cleanup_buckets(self):
+        """Remove old unused buckets to prevent memory leak."""
+        now = time.time()
+        if now - self.last_cleanup < self.cleanup_interval:
+            return
+        
+        async with self.lock:
+            # Remove buckets that haven't been used recently
+            cutoff_time = now - self.cleanup_interval
+            to_remove = []
+            
+            for identifier, bucket in self.buckets.items():
+                if bucket.last_refill < cutoff_time:
+                    to_remove.append(identifier)
+            
+            for identifier in to_remove:
+                del self.buckets[identifier]
+            
+            if to_remove:
+                logger.info(
+                    "rate_limiter_cleanup",
+                    removed_count=len(to_remove),
+                    remaining_count=len(self.buckets),
+                )
+            
+            self.last_cleanup = now
+
+
+class RateLimitMiddleware(BaseHTTPMiddleware):
+    """
+    Rate limiting middleware for FastAPI.
+    
+    Limits requests per client based on IP address or authenticated user.
+    """
+    
+    def __init__(
+        self,
+        app: ASGIApp,
+        requests_per_minute: int = 60,
+        burst_size: Optional[int] = None,
+        exclude_paths: Optional[list] = None,
+        identifier_callback: Optional[callable] = None,
+    ):
+        """
+        Initialize rate limit middleware.
+        
+        Args:
+            app: FastAPI application
+            requests_per_minute: Default rate limit
+            burst_size: Maximum burst size
+            exclude_paths: Paths to exclude from rate limiting
+            identifier_callback: Custom function to get client identifier
+        """
+        super().__init__(app)
+        self.rate_limiter = RateLimiter(requests_per_minute, burst_size)
+        self.exclude_paths = exclude_paths or ["/health", "/metrics", "/docs", "/openapi.json"]
+        self.identifier_callback = identifier_callback
+    
+    async def dispatch(self, request: Request, call_next):
+        """Process request with rate limiting."""
+        
+        # Check if path is excluded
+        if self._is_excluded(request.url.path):
+            return await call_next(request)
+        
+        # Get client identifier
+        identifier = await self._get_identifier(request)
+        
+        # Check rate limit
+        allowed, headers = await self.rate_limiter.check_rate_limit(identifier)
+        
+        if not allowed:
+            # Log rate limit exceeded
+            logger.warning(
+                "rate_limit_exceeded",
+                identifier=self._hash_identifier(identifier),
+                path=request.url.path,
+                method=request.method,
+            )
+            
+            # Raise rate limit error
+            raise RateLimitError(
+                limit=self.rate_limiter.requests_per_minute,
+                window="minute",
+                retry_after=int(headers.get("Retry-After", 60)),
+            )
+        
+        # Process request
+        response = await call_next(request)
+        
+        # Add rate limit headers to response
+        for header, value in headers.items():
+            response.headers[header] = value
+        
+        return response
+    
+    def _is_excluded(self, path: str) -> bool:
+        """Check if path is excluded from rate limiting."""
+        for excluded in self.exclude_paths:
+            if path.startswith(excluded):
+                return True
+        return False
+    
+    async def _get_identifier(self, request: Request) -> str:
+        """
+        Get client identifier for rate limiting.
+        
+        Priority:
+        1. Custom identifier callback
+        2. Authenticated user ID
+        3. Client IP address
+        """
+        # Use custom identifier callback if provided
+        if self.identifier_callback:
+            identifier = await self.identifier_callback(request)
+            if identifier:
+                return f"custom:{identifier}"
+        
+        # Check for authenticated user
+        if hasattr(request.state, "user") and request.state.user:
+            user_id = getattr(request.state.user, "id", None)
+            if user_id:
+                return f"user:{user_id}"
+        
+        # Fall back to IP address
+        client_host = request.client.host if request.client else "unknown"
+        
+        # Check for proxy headers
+        forwarded_for = request.headers.get("X-Forwarded-For")
+        if forwarded_for:
+            # Take the first IP in the chain
+            client_host = forwarded_for.split(",")[0].strip()
+        
+        real_ip = request.headers.get("X-Real-IP")
+        if real_ip:
+            client_host = real_ip
+        
+        return f"ip:{client_host}"
+    
+    def _hash_identifier(self, identifier: str) -> str:
+        """Hash identifier for logging (privacy)."""
+        return hashlib.sha256(identifier.encode()).hexdigest()[:16]
+
+
+# Per-endpoint rate limiting decorator
+class EndpointRateLimiter:
+    """
+    Decorator for per-endpoint rate limiting.
+    
+    Usage:
+        rate_limiter = EndpointRateLimiter()
+        
+        @app.get("/expensive-operation")
+        @rate_limiter.limit(requests_per_minute=10)
+        async def expensive_operation():
+            ...
+    """
+    
+    def __init__(self):
+        self.limiters: Dict[str, RateLimiter] = {}
+    
+    def limit(
+        self,
+        requests_per_minute: int,
+        burst_size: Optional[int] = None,
+        identifier_callback: Optional[callable] = None,
+    ):
+        """
+        Create rate limit decorator for endpoint.
+        
+        Args:
+            requests_per_minute: Rate limit for this endpoint
+            burst_size: Burst size for this endpoint
+            identifier_callback: Custom identifier function
+        """
+        def decorator(func):
+            # Create rate limiter for this endpoint
+            endpoint_id = f"{func.__module__}.{func.__name__}"
+            self.limiters[endpoint_id] = RateLimiter(
+                requests_per_minute=requests_per_minute,
+                burst_size=burst_size,
+            )
+            
+            async def wrapper(request: Request, *args, **kwargs):
+                # Get client identifier
+                if identifier_callback:
+                    identifier = await identifier_callback(request)
+                else:
+                    identifier = self._default_identifier(request)
+                
+                # Check rate limit
+                limiter = self.limiters[endpoint_id]
+                allowed, headers = await limiter.check_rate_limit(identifier)
+                
+                if not allowed:
+                    raise RateLimitError(
+                        limit=requests_per_minute,
+                        window="minute",
+                        retry_after=int(headers.get("Retry-After", 60)),
+                    )
+                
+                # Add headers to response (if we have access to it)
+                response = await func(request, *args, **kwargs)
+                if isinstance(response, Response):
+                    for header, value in headers.items():
+                        response.headers[header] = value
+                
+                return response
+            
+            return wrapper
+        return decorator
+    
+    def _default_identifier(self, request: Request) -> str:
+        """Default identifier extraction."""
+        if hasattr(request.state, "user") and request.state.user:
+            return f"user:{request.state.user.id}"
+        
+        client_host = request.client.host if request.client else "unknown"
+        return f"ip:{client_host}"
+
+
+# Global rate limiter instance for decorator usage
+endpoint_rate_limiter = EndpointRateLimiter()

control_plane_api/app/middleware/request_id.py

@@ -0,0 +1,202 @@
+"""
+Request ID middleware for request tracking and correlation.
+
+Adds a unique request ID to each request for tracking through logs and services.
+"""
+
+from fastapi import Request, Response
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.types import ASGIApp
+import structlog
+import logging
+import uuid
+import contextvars
+from typing import Optional
+
+logger = structlog.get_logger()
+
+# Context variable to store request ID
+request_id_context: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
+    "request_id", default=None
+)
+
+
+class RequestIDMiddleware(BaseHTTPMiddleware):
+    """
+    Middleware to add request ID to all requests.
+    
+    - Checks for existing X-Request-ID header
+    - Generates new ID if not present
+    - Adds ID to response headers
+    - Makes ID available in context for logging
+    """
+    
+    REQUEST_ID_HEADER = "X-Request-ID"
+    
+    def __init__(self, app: ASGIApp, header_name: str = None):
+        super().__init__(app)
+        if header_name:
+            self.REQUEST_ID_HEADER = header_name
+    
+    async def dispatch(self, request: Request, call_next):
+        """Process the request and add request ID."""
+        
+        # Get or generate request ID
+        request_id = (
+            request.headers.get(self.REQUEST_ID_HEADER) or
+            request.headers.get(self.REQUEST_ID_HEADER.lower()) or
+            self._generate_request_id()
+        )
+        
+        # Validate request ID format (basic security check)
+        if not self._is_valid_request_id(request_id):
+            request_id = self._generate_request_id()
+        
+        # Store in request state for easy access
+        request.state.request_id = request_id
+        
+        # Set context variable for logging
+        token = request_id_context.set(request_id)
+        
+        try:
+            # Log request start
+            logger.info(
+                "request_started",
+                request_id=request_id,
+                method=request.method,
+                path=str(request.url.path),
+                query=str(request.url.query) if request.url.query else None,
+                client=request.client.host if request.client else None,
+            )
+            
+            # Process request
+            response = await call_next(request)
+            
+            # Add request ID to response headers
+            response.headers[self.REQUEST_ID_HEADER] = request_id
+            
+            # Log request completion
+            logger.info(
+                "request_completed",
+                request_id=request_id,
+                status_code=response.status_code,
+            )
+            
+            return response
+            
+        except Exception as e:
+            # Log request failure
+            logger.error(
+                "request_failed",
+                request_id=request_id,
+                error=str(e),
+                error_type=type(e).__name__,
+            )
+            raise
+            
+        finally:
+            # Reset context variable
+            request_id_context.reset(token)
+    
+    def _generate_request_id(self) -> str:
+        """Generate a new request ID."""
+        return str(uuid.uuid4())
+    
+    def _is_valid_request_id(self, request_id: str) -> bool:
+        """
+        Validate request ID format.
+        
+        Accepts UUIDs and alphanumeric strings up to 128 characters.
+        """
+        if not request_id or len(request_id) > 128:
+            return False
+        
+        # Allow UUIDs, alphanumeric, hyphens, and underscores
+        import re
+        return bool(re.match(r'^[a-zA-Z0-9\-_]+$', request_id))
+
+
+def get_request_id() -> Optional[str]:
+    """
+    Get the current request ID from context.
+    
+    This can be used anywhere in the application to get the current request ID.
+    
+    Returns:
+        Request ID if in request context, None otherwise
+    """
+    return request_id_context.get()
+
+
+def set_request_id(request_id: str) -> None:
+    """
+    Set the request ID in context.
+    
+    This is useful for background tasks or other contexts where you want
+    to maintain the request ID.
+    
+    Args:
+        request_id: Request ID to set
+    """
+    request_id_context.set(request_id)
+
+
+class RequestIDLogProcessor:
+    """
+    Structlog processor to add request ID to all log entries.
+    
+    Use this in your structlog configuration:
+    
+    structlog.configure(
+        processors=[
+            RequestIDLogProcessor(),
+            structlog.processors.add_log_level,
+            structlog.processors.TimeStamper(fmt="iso"),
+            structlog.processors.JSONRenderer(),
+        ],
+        ...
+    )
+    """
+    
+    def __call__(self, logger, name, event_dict):
+        """Add request ID to log entry if available."""
+        request_id = get_request_id()
+        if request_id:
+            event_dict["request_id"] = request_id
+        return event_dict
+
+
+def setup_request_id_logging():
+    """
+    Configure structlog to include request ID in all logs.
+    
+    Call this in your app initialization.
+    """
+    import structlog
+    
+    structlog.configure(
+        processors=[
+            RequestIDLogProcessor(),  # Add request ID
+            structlog.contextvars.merge_contextvars,
+            structlog.processors.add_log_level,
+            structlog.processors.TimeStamper(fmt="iso"),
+            structlog.processors.JSONRenderer(),
+        ],
+        wrapper_class=structlog.make_filtering_bound_logger(
+            logging.INFO
+        ),
+        logger_factory=structlog.PrintLoggerFactory(),
+    )
+
+
+# For FastAPI dependency injection
+async def get_request_id_from_request(request: Request) -> str:
+    """
+    FastAPI dependency to get request ID.
+    
+    Usage:
+        @app.get("/example")
+        async def example(request_id: str = Depends(get_request_id_from_request)):
+            return {"request_id": request_id}
+    """
+    return getattr(request.state, "request_id", None) or str(uuid.uuid4())

control_plane_api/app/models/__init__.py

@@ -0,0 +1,40 @@
+# Database Models
+from control_plane_api.app.models.project import Project, ProjectStatus
+from control_plane_api.app.models.agent import Agent, AgentStatus
+from control_plane_api.app.models.team import Team, TeamStatus
+from control_plane_api.app.models.workflow import Workflow, WorkflowStatus
+from control_plane_api.app.models.workspace import Workspace
+from control_plane_api.app.models.user_profile import UserProfile
+from control_plane_api.app.models.session import Session
+from control_plane_api.app.models.execution import Execution, ExecutionStatus, ExecutionType, ExecutionTriggerSource
+from control_plane_api.app.models.presence import UserPresence
+from control_plane_api.app.models.environment import Environment, EnvironmentStatus
+from control_plane_api.app.models.associations import AgentEnvironment, TeamEnvironment, ExecutionParticipant, ParticipantRole
+from control_plane_api.app.models.job import Job, JobExecution, JobStatus, JobTriggerType, ExecutorType, PlanningMode
+from control_plane_api.app.models.llm_model import LLMModel
+from control_plane_api.app.models.orchestration import Namespace, TemporalNamespace
+from control_plane_api.app.models.project_management import ProjectAgent
+from control_plane_api.app.models.skill import SkillAssociation, Skill
+from control_plane_api.app.models.worker import WorkerQueue
+from control_plane_api.app.models.trace import Trace, Span, TraceStatus, SpanKind, SpanStatusCode
+
+__all__ = [
+    "Project", "ProjectStatus",
+    "Agent", "AgentStatus",
+    "Team", "TeamStatus",
+    "Workflow", "WorkflowStatus",
+    "Workspace",
+    "UserProfile",
+    "Session",
+    "Execution", "ExecutionStatus", "ExecutionType", "ExecutionTriggerSource",
+    "UserPresence",
+    "Environment", "EnvironmentStatus",
+    "AgentEnvironment", "TeamEnvironment",
+    "ExecutionParticipant", "ParticipantRole",
+    "Job", "JobExecution", "JobStatus", "JobTriggerType", "ExecutorType", "PlanningMode",
+    "LLMModel",
+    "Namespace", "TemporalNamespace", "ProjectAgent",
+    "SkillAssociation", "Skill",
+    "WorkerQueue",
+    "Trace", "Span", "TraceStatus", "SpanKind", "SpanStatusCode"
+]