QuickLab-MCP 0.1.0__py3-none-any.whl

ml_mcp/server/http.py

@@ -0,0 +1,183 @@
+"""HTTP transport application for ModelLab FastMCP Server.
+
+Exposes Streamable HTTP endpoint (/mcp), health probes (/health/live, /health/ready),
+and enforces OAuth 2.1 / OIDC Bearer token authentication middleware.
+"""
+
+import logging
+
+from starlette.applications import Starlette
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+from starlette.types import ASGIApp, Receive, Scope, Send
+
+from ml_mcp.config import get_settings
+from ml_mcp.infrastructure.object_storage.s3 import get_storage_service
+from ml_mcp.infrastructure.postgres.session import get_db_manager
+from ml_mcp.infrastructure.redis.client import get_redis_manager
+from ml_mcp.server.app import create_server
+from ml_mcp.server.auth import get_token_validator
+from ml_mcp.server.context import set_current_principal
+
+logger = logging.getLogger(__name__)
+
+
+class BearerAuthMiddleware:
+    """Pure ASGI middleware enforcing Bearer token authentication on protected MCP endpoints."""
+
+    def __init__(self, app: ASGIApp) -> None:
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        path: str = scope.get("path", "")
+        # Health probes bypass token authentication
+        if path.startswith("/health"):
+            await self.app(scope, receive, send)
+            return
+
+        settings = get_settings()
+        if not settings.auth.enabled:
+            await self.app(scope, receive, send)
+            return
+
+        # Extract Authorization header
+        headers = dict(scope.get("headers", []))
+        auth_header_bytes = headers.get(b"authorization", b"")
+        auth_header = auth_header_bytes.decode("utf-8") if auth_header_bytes else ""
+
+        if not auth_header or not auth_header.startswith("Bearer "):
+            response = JSONResponse(
+                {
+                    "error": "Unauthorized",
+                    "code": "AUTHENTICATION_REQUIRED",
+                    "detail": "Missing Bearer token in Authorization header",
+                },
+                status_code=401,
+            )
+            await response(scope, receive, send)
+            return
+
+        token = auth_header.split(" ", 1)[1].strip()
+        try:
+            validator = get_token_validator()
+            principal = validator.validate_token(token)
+        except Exception as exc:
+            logger.warning("Authentication failure: %s", exc)
+            response = JSONResponse(
+                {
+                    "error": "Unauthorized",
+                    "code": "INVALID_TOKEN",
+                    "detail": str(exc),
+                },
+                status_code=401,
+            )
+            await response(scope, receive, send)
+            return
+
+        # Bind authenticated principal to context variable for downstream handlers
+        set_current_principal(principal)
+        try:
+            await self.app(scope, receive, send)
+        finally:
+            set_current_principal(None)
+
+
+async def health_live(request: Request) -> JSONResponse:
+    """Liveness probe: verifies process is alive and responsive."""
+    return JSONResponse({"status": "alive", "service": "modellab-mcp-api"})
+
+
+async def health_ready(request: Request) -> JSONResponse:
+    """Readiness probe: verifies operational connectivity to Postgres, Redis, and Object Store."""
+    checks: dict[str, str] = {}
+    all_healthy = True
+
+    # 1. Database Check
+    try:
+        from sqlalchemy import text
+
+        async with get_db_manager().session() as sess:
+            await sess.execute(text("SELECT 1"))
+        checks["database"] = "healthy"
+    except Exception as exc:
+        checks["database"] = f"unhealthy: {exc}"
+        all_healthy = False
+
+    # 2. Redis Check
+    try:
+        redis_mgr = get_redis_manager()
+        healthy = await redis_mgr.check_health()
+        checks["redis"] = "healthy" if healthy else "degraded"
+    except Exception as exc:
+        checks["redis"] = f"degraded: {exc}"
+        # Redis failure can fall back to in-memory, but flag status
+
+    # 3. Object Store Check
+    try:
+        storage = get_storage_service()
+        # Verify storage client initialized
+        if storage.client is not None:
+            checks["object_store"] = "healthy"
+        else:
+            checks["object_store"] = "uninitialized"
+            all_healthy = False
+    except Exception as exc:
+        checks["object_store"] = f"unhealthy: {exc}"
+        all_healthy = False
+
+    status_code = 200 if all_healthy else 503
+    return JSONResponse(
+        {
+            "status": "ready" if all_healthy else "degraded",
+            "checks": checks,
+        },
+        status_code=status_code,
+    )
+
+
+def create_http_app() -> Starlette:
+    """Create configured Starlette application wrapping FastMCP with auth and operational endpoints."""
+    settings = get_settings()
+    server = create_server()
+
+    # FastMCP streamable HTTP app mounted at configured path (default: /mcp)
+    app = server.http_app(
+        path=settings.mcp.streamable_http_path,
+        transport="streamable-http",
+    )
+
+    # Register health probes
+    app.add_route("/health/live", health_live, methods=["GET"])
+    app.add_route("/health/ready", health_ready, methods=["GET"])
+
+    # Register BearerAuthMiddleware
+    app.add_middleware(BearerAuthMiddleware)
+
+    return app
+
+
+def run_http() -> None:
+    """CLI entrypoint for running the ModelLab HTTP server using uvicorn."""
+    import uvicorn
+
+    from ml_mcp.infrastructure.telemetry.logging import configure_logging
+
+    configure_logging()
+    settings = get_settings()
+
+    uvicorn.run(
+        "ml_mcp.server.http:create_http_app",
+        factory=True,
+        host=settings.mcp.host,
+        port=settings.mcp.port,
+        log_config=None,  # Use structured stderr logger
+        access_log=False,
+    )
+
+
+if __name__ == "__main__":
+    run_http()

ml_mcp/server/middleware.py

@@ -0,0 +1,121 @@
+"""Security, authorization, and rate-limiting middleware pipeline."""
+
+import logging
+from typing import Any
+
+from ml_mcp.config import get_settings
+from ml_mcp.domain.errors import ResourceLimitExceededError
+from ml_mcp.domain.policies import (
+    TOOL_POLICIES,
+    Principal,
+    Role,
+)
+from ml_mcp.infrastructure.postgres.models import AuditEventOrm
+from ml_mcp.infrastructure.postgres.repositories.audit import AuditRepository
+from ml_mcp.infrastructure.postgres.session import get_db_manager
+from ml_mcp.infrastructure.redis.rate_limiter import SlidingWindowRateLimiter
+from ml_mcp.infrastructure.telemetry.logging import (
+    clear_request_context,
+    set_request_context,
+)
+from ml_mcp.infrastructure.telemetry.otel import get_instruments
+
+logger = logging.getLogger(__name__)
+
+
+# Default local principal for STDIO operations (trusted local process boundary)
+LOCAL_STDIO_PRINCIPAL = Principal(
+    principal_id="local-user",
+    tenant_id="default-tenant",
+    role=Role.ADMIN,
+)
+
+
+class SecurityPipeline:
+    """Orchestrates authentication, scope checks, rate limiting, and audit logging for tool executions."""
+
+    def __init__(self, rate_limiter: SlidingWindowRateLimiter | None = None) -> None:
+        self.rate_limiter = rate_limiter or SlidingWindowRateLimiter()
+
+    async def pre_tool_call(
+        self,
+        tool_name: str,
+        principal: Principal,
+        request_id: str | None = None,
+    ) -> None:
+        """Enforce rate limits, scopes, and set request context."""
+        policy = TOOL_POLICIES.get(tool_name)
+        if not policy:
+            return
+
+        # 1. Telemetry context
+        set_request_context(
+            request_id=request_id,
+            tenant_id=principal.tenant_id,
+            principal_id=principal.principal_id,
+        )
+
+        # 2. Scope enforcement
+        principal.enforce_permission(policy.required_scope)
+
+        # 3. Rate limiting enforcement
+        settings = get_settings()
+        limits_map = {
+            "read": settings.rate_limit.read_per_minute,
+            "metadata": settings.rate_limit.metadata_per_minute,
+            "experiment": settings.rate_limit.experiment_per_minute,
+            "analysis": settings.rate_limit.analysis_per_minute,
+        }
+        limit = limits_map.get(policy.rate_limit_category, 60)
+        res = await self.rate_limiter.check_rate_limit(
+            identifier=f"{principal.tenant_id}:{principal.principal_id}",
+            category=policy.rate_limit_category,
+            limit=limit,
+        )
+        if not res.allowed:
+            get_instruments().rate_limit_denials_total.add(1)
+            raise ResourceLimitExceededError(
+                f"Rate limit exceeded for tool category '{policy.rate_limit_category}'",
+                limit=limit,
+                current=limit + 1,
+            )
+
+    async def post_tool_call(
+        self,
+        tool_name: str,
+        principal: Principal,
+        result: Any,
+        success: bool = True,
+        error_code: str | None = None,
+    ) -> None:
+        """Emit audit log entry if required by policy and clear telemetry context."""
+        policy = TOOL_POLICIES.get(tool_name)
+        if policy and policy.requires_audit:
+            try:
+                async with get_db_manager().session() as sess:
+                    audit_repo = AuditRepository(sess)
+                    await audit_repo.record_event(
+                        AuditEventOrm(
+                            principal_id=principal.principal_id,
+                            tenant_id=principal.tenant_id,
+                            action=f"TOOL_{tool_name.upper()}",
+                            resource_type="tool",
+                            resource_id=tool_name,
+                            result="SUCCESS" if success else "ERROR",
+                            reason_code=error_code,
+                        )
+                    )
+            except Exception as exc:
+                logger.warning("Failed to record audit event for %s: %s", tool_name, exc)
+
+        clear_request_context()
+
+
+_pipeline: SecurityPipeline | None = None
+
+
+def get_security_pipeline() -> SecurityPipeline:
+    global _pipeline
+    if _pipeline is None:
+        _pipeline = SecurityPipeline()
+    return _pipeline

ml_mcp/server/stdio.py

@@ -0,0 +1,42 @@
+"""STDIO transport runner for ModelLab FastMCP Server.
+
+Adheres strictly to standard stream discipline:
+- stdout: Exclusively reserved for valid MCP JSON-RPC protocol frames.
+- stderr: All structured JSON logs, diagnostic output, and trace messages.
+"""
+
+import asyncio
+import logging
+import sys
+
+from ml_mcp.infrastructure.telemetry.logging import configure_logging
+from ml_mcp.server.app import create_server
+
+logger = logging.getLogger(__name__)
+
+
+async def run_stdio_server() -> None:
+    """Initialize server components and launch FastMCP STDIO loop."""
+    # Strict stream discipline: logs go to stderr ONLY
+    configure_logging()
+    logger.info("Initializing ModelLab FastMCP server on STDIO transport...")
+
+    server = create_server("ModelLab-STDIO")
+
+    # Run FastMCP stdio server
+    await server.run_stdio_async()
+
+
+def run_stdio() -> None:
+    """Synchronous entrypoint for STDIO launcher."""
+    try:
+        asyncio.run(run_stdio_server())
+    except (KeyboardInterrupt, SystemExit):
+        sys.stderr.write("ModelLab STDIO server terminated gracefully.\n")
+    except Exception as exc:
+        sys.stderr.write(f"Fatal error in ModelLab STDIO server: {exc}\n")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    run_stdio()

ml_mcp/workers/__init__.py

@@ -0,0 +1 @@
+"""Isolated ML worker package."""

ml_mcp/workers/daemon.py

@@ -0,0 +1,130 @@
+"""Worker daemon dequeuing and executing ML experiment tasks from Redis."""
+
+import asyncio
+import logging
+import signal
+import sys
+from typing import Any
+
+from ml_mcp.infrastructure.object_storage.s3 import S3StorageService, get_storage_service
+from ml_mcp.infrastructure.postgres.session import DatabaseManager, get_db_manager
+from ml_mcp.infrastructure.queue.task_queue import RedisTaskQueue
+from ml_mcp.infrastructure.telemetry.logging import configure_logging
+from ml_mcp.workers.runner import WorkerRunner
+
+logger = logging.getLogger(__name__)
+
+
+class WorkerDaemon:
+    """Asynchronous background worker daemon executing experiment tasks from Redis."""
+
+    def __init__(
+        self,
+        task_queue: RedisTaskQueue | None = None,
+        storage_service: S3StorageService | None = None,
+        db_manager: DatabaseManager | None = None,
+        drain_outbox: bool = False,
+    ) -> None:
+        self.queue = task_queue or RedisTaskQueue()
+        self.storage = storage_service or get_storage_service()
+        self.db_mgr = db_manager or get_db_manager()
+        self.drain_outbox = drain_outbox
+        self._stop_event = asyncio.Event()
+        self._running = False
+
+    def stop(self, *args: Any) -> None:
+        """Signal the daemon loop to terminate gracefully."""
+        logger.info("Shutdown signal received; initiating graceful termination of worker daemon...")
+        self._stop_event.set()
+
+    async def run(self) -> None:
+        """Run continuous dequeue-execute loop until stopped."""
+        self._running = True
+        logger.info("Worker daemon started, polling for experiment tasks...")
+
+        while not self._stop_event.is_set():
+            try:
+                # 1. Drain pending outbox events into Redis queue if enabled
+                if self.drain_outbox:
+                    try:
+                        async with self.db_mgr.session() as session:
+                            from ml_mcp.infrastructure.queue.outbox_processor import OutboxProcessor
+
+                            processor = OutboxProcessor(session, self.queue)
+                            await processor.process_batch(limit=20)
+                    except Exception as outbox_exc:
+                        logger.debug("Outbox drain cycle check: %s", outbox_exc)
+
+                # 2. Dequeue experiment task
+                task_message = await self.queue.dequeue(timeout_seconds=2)
+                if task_message is None:
+                    continue
+
+                logger.info(
+                    "Dequeued task %s for experiment %s (tenant %s)",
+                    task_message.task_id,
+                    task_message.experiment_id,
+                    task_message.tenant_id,
+                )
+
+                async with self.db_mgr.session() as session:
+                    runner = WorkerRunner(session=session, storage_service=self.storage)
+                    try:
+                        result = await runner.execute_experiment(
+                            experiment_id=task_message.experiment_id,
+                            tenant_id=task_message.tenant_id,
+                        )
+                        logger.info(
+                            "Task %s succeeded for experiment %s: %s",
+                            task_message.task_id,
+                            task_message.experiment_id,
+                            result.get("status"),
+                        )
+                    except Exception as exc:
+                        logger.error(
+                            "Task %s failed for experiment %s: %s",
+                            task_message.task_id,
+                            task_message.experiment_id,
+                            exc,
+                        )
+
+            except asyncio.CancelledError:
+                logger.info("Worker daemon cancelled")
+                break
+            except Exception as exc:
+                logger.error("Unexpected error in worker daemon polling loop: %s", exc)
+                # Brief sleep on unexpected error to prevent rapid spinning
+                await asyncio.sleep(1)
+
+        self._running = False
+        logger.info("Worker daemon stopped cleanly.")
+
+
+def _register_signals(daemon: WorkerDaemon) -> None:
+    """Attach graceful shutdown handlers for SIGINT and SIGTERM."""
+    loop = asyncio.get_running_loop()
+    signals = [signal.SIGINT]
+    if hasattr(signal, "SIGTERM"):
+        signals.append(signal.SIGTERM)
+
+    for sig in signals:
+        try:
+            loop.add_signal_handler(sig, daemon.stop)
+        except (NotImplementedError, AttributeError):
+            # Fallback for platforms without loop.add_signal_handler (e.g. Windows)
+            signal.signal(sig, lambda s, f: daemon.stop())
+
+
+async def main() -> None:
+    """Daemon CLI entrypoint."""
+    configure_logging()
+    daemon = WorkerDaemon()
+    _register_signals(daemon)
+    await daemon.run()
+
+
+if __name__ == "__main__":
+    try:
+        asyncio.run(main())
+    except (KeyboardInterrupt, SystemExit):
+        sys.exit(0)

ml_mcp/workers/metrics.py

@@ -0,0 +1,110 @@
+"""Comprehensive evaluation metrics computation engine for classification and regression."""
+
+import contextlib
+
+import numpy as np
+from sklearn.metrics import (
+    accuracy_score,
+    average_precision_score,
+    brier_score_loss,
+    confusion_matrix,
+    explained_variance_score,
+    f1_score,
+    log_loss,
+    mean_absolute_error,
+    mean_absolute_percentage_error,
+    mean_squared_error,
+    precision_score,
+    r2_score,
+    recall_score,
+    roc_auc_score,
+)
+
+from ml_mcp.domain.value_objects import TaskType
+
+
+class MetricsCalculator:
+    """Calculates all standard ML evaluation metrics without placeholders or approximations."""
+
+    @staticmethod
+    def compute_classification_metrics(
+        y_true: np.ndarray,
+        y_pred: np.ndarray,
+        y_proba: np.ndarray | None = None,
+        task_type: TaskType = TaskType.BINARY_CLASSIFICATION,
+    ) -> dict[str, float | list[list[int]]]:
+        """Compute full classification metric suite."""
+        metrics: dict[str, float | list[list[int]]] = {}
+
+        metrics["accuracy"] = float(accuracy_score(y_true, y_pred))
+        metrics["f1"] = float(
+            f1_score(
+                y_true,
+                y_pred,
+                average="weighted" if task_type == TaskType.MULTICLASS_CLASSIFICATION else "binary",
+                zero_division=0,
+            )
+        )
+        metrics["f1_macro"] = float(f1_score(y_true, y_pred, average="macro", zero_division=0))
+        metrics["precision"] = float(
+            precision_score(
+                y_true,
+                y_pred,
+                average="weighted" if task_type == TaskType.MULTICLASS_CLASSIFICATION else "binary",
+                zero_division=0,
+            )
+        )
+        metrics["recall"] = float(
+            recall_score(
+                y_true,
+                y_pred,
+                average="weighted" if task_type == TaskType.MULTICLASS_CLASSIFICATION else "binary",
+                zero_division=0,
+            )
+        )
+
+        # Confusion matrix
+        cm = confusion_matrix(y_true, y_pred)
+        metrics["confusion_matrix"] = cm.tolist()
+
+        # Probability-based metrics
+        if y_proba is not None:
+            with contextlib.suppress(Exception):
+                if task_type == TaskType.BINARY_CLASSIFICATION:
+                    # y_proba can be 1D or 2D
+                    proba_pos = (
+                        y_proba[:, 1] if y_proba.ndim == 2 and y_proba.shape[1] > 1 else y_proba
+                    )
+                    metrics["roc_auc"] = float(roc_auc_score(y_true, proba_pos))
+                    metrics["pr_auc"] = float(average_precision_score(y_true, proba_pos))
+                    metrics["brier_score"] = float(brier_score_loss(y_true, proba_pos))
+                    metrics["log_loss"] = float(log_loss(y_true, y_proba))
+                elif task_type == TaskType.MULTICLASS_CLASSIFICATION:
+                    metrics["roc_auc"] = float(roc_auc_score(y_true, y_proba, multi_class="ovr"))
+                    metrics["log_loss"] = float(log_loss(y_true, y_proba))
+
+        return metrics
+
+    @staticmethod
+    def compute_regression_metrics(
+        y_true: np.ndarray,
+        y_pred: np.ndarray,
+    ) -> dict[str, float]:
+        """Compute full regression metric suite."""
+        mse = float(mean_squared_error(y_true, y_pred))
+        rmse = float(np.sqrt(mse))
+        mae = float(mean_absolute_error(y_true, y_pred))
+        r2 = float(r2_score(y_true, y_pred))
+
+        metrics: dict[str, float] = {
+            "mse": mse,
+            "rmse": rmse,
+            "mae": mae,
+            "r2": r2,
+            "explained_variance": float(explained_variance_score(y_true, y_pred)),
+        }
+
+        with contextlib.suppress(Exception):
+            metrics["mape"] = float(mean_absolute_percentage_error(y_true, y_pred))
+
+        return metrics