QuickLab-MCP 0.1.0__py3-none-any.whl

ml_mcp/infrastructure/queue/outbox_processor.py

@@ -0,0 +1,67 @@
+"""Transactional Outbox processor dispatching experiment events to the task queue."""
+
+import logging
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from ml_mcp.infrastructure.postgres.repositories.outbox import OutboxRepository
+from ml_mcp.infrastructure.queue.task_queue import RedisTaskQueue
+
+logger = logging.getLogger(__name__)
+
+
+class OutboxProcessor:
+    """Polls transactional outbox and guarantees reliable message delivery to Redis task queue."""
+
+    def __init__(
+        self,
+        session: AsyncSession,
+        task_queue: RedisTaskQueue | None = None,
+        max_retries: int = 5,
+    ) -> None:
+        self.session = session
+        self.repo = OutboxRepository(session)
+        self.queue = task_queue or RedisTaskQueue()
+        self.max_retries = max_retries
+
+    async def process_batch(self, limit: int = 50) -> int:
+        """Process scheduled pending outbox events and dispatch them to the worker queue."""
+        events = await self.repo.fetch_pending(limit=limit)
+        delivered_count = 0
+
+        for event in events:
+            try:
+                payload = event.payload_json
+                experiment_id = payload.get("experiment_id", event.aggregate_id)
+                tenant_id = payload.get("tenant_id", "system")
+
+                # Push to Redis queue
+                await self.queue.enqueue(
+                    experiment_id=experiment_id,
+                    tenant_id=tenant_id,
+                    payload=payload,
+                )
+
+                # Mark delivered in outbox table
+                await self.repo.mark_delivered(event.id)
+                delivered_count += 1
+                logger.info(
+                    "Successfully dispatched outbox event %s for experiment %s",
+                    event.id,
+                    experiment_id,
+                )
+
+            except Exception as exc:
+                new_retry = event.retry_count + 1
+                logger.error(
+                    "Failed to dispatch outbox event %s (attempt %d/%d): %s",
+                    event.id,
+                    new_retry,
+                    self.max_retries,
+                    exc,
+                )
+                await self.repo.mark_failed(
+                    event.id, retry_count=new_retry, max_retries=self.max_retries
+                )
+
+        return delivered_count

ml_mcp/infrastructure/queue/task_queue.py

@@ -0,0 +1,98 @@
+"""Durable task queue implementation backed by Redis list/streams with local queue fallback."""
+
+import asyncio
+import json
+import logging
+import time
+from dataclasses import asdict, dataclass
+from typing import Any
+
+from ml_mcp.config import Settings, get_settings
+from ml_mcp.infrastructure.postgres.base import generate_uuid7
+from ml_mcp.infrastructure.redis.client import RedisManager, get_redis_manager
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TaskMessage:
+    """Message payload transmitted from Outbox processor to isolated ML worker."""
+
+    task_id: str
+    experiment_id: str
+    tenant_id: str
+    payload: dict[str, Any]
+    created_at: float
+
+    def to_json(self) -> str:
+        return json.dumps(asdict(self))
+
+    @classmethod
+    def from_json(cls, data_str: str) -> "TaskMessage":
+        data = json.loads(data_str)
+        return cls(**data)
+
+
+class RedisTaskQueue:
+    """Reliable Redis-backed queue for orchestrating ML worker tasks."""
+
+    def __init__(
+        self,
+        redis_manager: RedisManager | None = None,
+        settings: Settings | None = None,
+    ) -> None:
+        self.settings = settings or get_settings()
+        self.redis_mgr = redis_manager or get_redis_manager()
+        self.queue_name = self.settings.redis.task_queue_name
+        # In-memory queue fallback for isolated testing
+        self._local_queue: asyncio.Queue[TaskMessage] = asyncio.Queue()
+
+    async def enqueue(
+        self,
+        experiment_id: str,
+        tenant_id: str,
+        payload: dict[str, Any],
+    ) -> str:
+        """Enqueue an experiment task for execution. Returns task_id."""
+        task_id = generate_uuid7()
+        message = TaskMessage(
+            task_id=task_id,
+            experiment_id=experiment_id,
+            tenant_id=tenant_id,
+            payload=payload,
+            created_at=time.time(),
+        )
+
+        try:
+            client = self.redis_mgr.client
+            await client.rpush(self.queue_name, message.to_json())
+            logger.info("Enqueued task %s for experiment %s in Redis", task_id, experiment_id)
+        except Exception as exc:
+            logger.warning("Redis unavailable, enqueuing to in-memory queue fallback: %s", exc)
+            await self._local_queue.put(message)
+
+        return task_id
+
+    async def dequeue(self, timeout_seconds: int = 5) -> TaskMessage | None:
+        """Pop the next available task from the queue with timeout."""
+        try:
+            client = self.redis_mgr.client
+            result = await client.blpop(self.queue_name, timeout=timeout_seconds)
+            if result:
+                _, data_str = result
+                return TaskMessage.from_json(data_str)
+            return None
+        except Exception:
+            # Fallback to local queue
+            try:
+                return await asyncio.wait_for(self._local_queue.get(), timeout=timeout_seconds)
+            except TimeoutError:
+                return None
+
+    async def get_queue_depth(self) -> int:
+        """Return total count of waiting tasks in the queue."""
+        try:
+            client = self.redis_mgr.client
+            return await client.llen(self.queue_name)
+        except Exception:
+            return self._local_queue.qsize()

ml_mcp/infrastructure/redis/__init__.py

@@ -0,0 +1,11 @@
+"""Redis client and rate limiting infrastructure."""
+
+from ml_mcp.infrastructure.redis.client import RedisManager, get_redis_manager
+from ml_mcp.infrastructure.redis.rate_limiter import RateLimitResult, SlidingWindowRateLimiter
+
+__all__ = [
+    "RedisManager",
+    "get_redis_manager",
+    "SlidingWindowRateLimiter",
+    "RateLimitResult",
+]

ml_mcp/infrastructure/redis/client.py

@@ -0,0 +1,64 @@
+"""Async Redis connection management."""
+
+import logging
+
+from redis.asyncio import ConnectionPool, Redis
+
+from ml_mcp.config import Settings, get_settings
+
+logger = logging.getLogger(__name__)
+
+
+class RedisManager:
+    """Manages Redis connection pool and async client."""
+
+    def __init__(self, settings: Settings | None = None) -> None:
+        self.settings = settings or get_settings()
+        self._pool: ConnectionPool | None = None
+        self._client: Redis | None = None
+
+    def initialize(self, custom_url: str | None = None) -> None:
+        redis_url = custom_url or self.settings.redis.url
+        self._pool = ConnectionPool.from_url(
+            redis_url,
+            max_connections=self.settings.redis.max_connections,
+            decode_responses=True,
+            socket_connect_timeout=0.5,
+            socket_timeout=1.0,
+        )
+        self._client = Redis(connection_pool=self._pool)
+        logger.info("Redis client initialized")
+
+    @property
+    def client(self) -> Redis:
+        if self._client is None:
+            self.initialize()
+        assert self._client is not None
+        return self._client
+
+    async def check_health(self) -> bool:
+        """Ping Redis server to assert availability."""
+        try:
+            return await self.client.ping() is True
+        except Exception as exc:
+            logger.error("Redis healthcheck failed: %s", exc)
+            return False
+
+    async def close(self) -> None:
+        if self._client is not None:
+            await self._client.aclose()
+        if self._pool is not None:
+            await self._pool.disconnect()
+        self._client = None
+        self._pool = None
+        logger.info("Redis connection closed")
+
+
+_redis_manager: RedisManager | None = None
+
+
+def get_redis_manager() -> RedisManager:
+    global _redis_manager
+    if _redis_manager is None:
+        _redis_manager = RedisManager()
+    return _redis_manager

ml_mcp/infrastructure/redis/rate_limiter.py

@@ -0,0 +1,129 @@
+"""Sliding window distributed rate limiter using Redis sorted sets and atomic Lua script."""
+
+import time
+from collections import defaultdict
+from dataclasses import dataclass
+
+from redis.asyncio import Redis
+
+from ml_mcp.infrastructure.redis.client import RedisManager, get_redis_manager
+
+LUA_RATE_LIMIT_SCRIPT = """
+local key = KEYS[1]
+local now = tonumber(ARGV[1])
+local window_seconds = tonumber(ARGV[2])
+local limit = tonumber(ARGV[3])
+local member = ARGV[4]
+
+local window_start = now - window_seconds
+
+-- 1. Remove entries older than window_start
+redis.call('ZREMRANGEBYSCORE', key, '-inf', window_start)
+
+-- 2. Count existing entries
+local current_count = redis.call('ZCARD', key)
+
+if current_count < limit then
+    -- Allowed: add entry and update TTL
+    redis.call('ZADD', key, now, member)
+    redis.call('EXPIRE', key, math.ceil(window_seconds * 2))
+    local remaining = limit - (current_count + 1)
+    return {1, limit, remaining, tostring(now + window_seconds), "0.0"}
+else
+    -- Denied: find oldest entry to determine reset time
+    local oldest = redis.call('ZRANGE', key, 0, 0, 'WITHSCORES')
+    local reset_time = now + window_seconds
+    if oldest and #oldest >= 2 then
+        reset_time = tonumber(oldest[2]) + window_seconds
+    end
+    local retry_after = reset_time - now
+    if retry_after < 0 then
+        retry_after = 0
+    end
+    return {0, limit, 0, tostring(reset_time), tostring(retry_after)}
+end
+"""
+
+
+@dataclass
+class RateLimitResult:
+    allowed: bool
+    limit: int
+    remaining: int
+    reset_epoch_seconds: float
+    retry_after_seconds: float
+
+
+class SlidingWindowRateLimiter:
+    """Sliding-window rate limiter ensuring atomic Lua script execution across all tool tiers."""
+
+    def __init__(self, redis_manager: RedisManager | None = None) -> None:
+        self.redis_mgr = redis_manager or get_redis_manager()
+        # In-memory fallback for testing or when Redis is unavailable
+        self._local_history: dict[str, list[float]] = defaultdict(list)
+
+    async def check_rate_limit(
+        self,
+        identifier: str,
+        category: str = "default",
+        limit: int = 60,
+        window_seconds: int = 60,
+    ) -> RateLimitResult:
+        """Check and record a rate limit hit atomically using Redis Lua script.
+
+        Returns:
+            RateLimitResult with allowed status, remaining tokens, and reset epoch.
+        """
+        key = f"ratelimit:{category}:{identifier}"
+        now = time.time()
+        member = f"{now}:{time.time_ns()}"
+        window_start = now - window_seconds
+
+        try:
+            client: Redis = self.redis_mgr.client
+            result = await client.eval(
+                LUA_RATE_LIMIT_SCRIPT,
+                1,
+                key,
+                str(now),
+                str(window_seconds),
+                str(limit),
+                member,
+            )
+            allowed = bool(result[0])
+            limit_val = int(result[1])
+            remaining_val = int(result[2])
+            reset_epoch = float(result[3])
+            retry_after = float(result[4])
+            return RateLimitResult(
+                allowed=allowed,
+                limit=limit_val,
+                remaining=remaining_val,
+                reset_epoch_seconds=reset_epoch,
+                retry_after_seconds=retry_after,
+            )
+
+        except Exception:
+            # In-memory fallback if Redis is offline/unconfigured in testing
+            history = self._local_history[key]
+            valid = [ts for ts in history if ts > window_start]
+            self._local_history[key] = valid
+
+            if len(valid) >= limit:
+                reset_time = valid[0] + window_seconds if valid else now + window_seconds
+                return RateLimitResult(
+                    allowed=False,
+                    limit=limit,
+                    remaining=0,
+                    reset_epoch_seconds=reset_time,
+                    retry_after_seconds=max(0.0, reset_time - now),
+                )
+
+            valid.append(now)
+            return RateLimitResult(
+                allowed=True,
+                limit=limit,
+                remaining=max(0, limit - len(valid)),
+                reset_epoch_seconds=now + window_seconds,
+                retry_after_seconds=0.0,
+            )

ml_mcp/infrastructure/telemetry/__init__.py

@@ -0,0 +1,23 @@
+"""Telemetry and observability module."""
+
+from ml_mcp.infrastructure.telemetry.logging import (
+    clear_request_context,
+    configure_logging,
+    get_logger,
+    set_request_context,
+)
+from ml_mcp.infrastructure.telemetry.otel import (
+    get_meter,
+    get_tracer,
+    init_telemetry,
+)
+
+__all__ = [
+    "configure_logging",
+    "get_logger",
+    "set_request_context",
+    "clear_request_context",
+    "init_telemetry",
+    "get_tracer",
+    "get_meter",
+]

ml_mcp/infrastructure/telemetry/logging.py

@@ -0,0 +1,155 @@
+"""Structured JSON logging strictly directed to sys.stderr to preserve STDIO transport."""
+
+import contextvars
+import datetime
+import json
+import logging
+import sys
+from typing import Any
+
+# Context variables for correlation
+ctx_request_id: contextvars.ContextVar[str | None] = contextvars.ContextVar(
+    "ctx_request_id", default=None
+)
+ctx_trace_id: contextvars.ContextVar[str | None] = contextvars.ContextVar(
+    "ctx_trace_id", default=None
+)
+ctx_tenant_id: contextvars.ContextVar[str | None] = contextvars.ContextVar(
+    "ctx_tenant_id", default=None
+)
+ctx_principal_id: contextvars.ContextVar[str | None] = contextvars.ContextVar(
+    "ctx_principal_id", default=None
+)
+
+
+def set_request_context(
+    request_id: str | None = None,
+    trace_id: str | None = None,
+    tenant_id: str | None = None,
+    principal_id: str | None = None,
+) -> None:
+    """Set request-scoped telemetry and correlation IDs."""
+    if request_id is not None:
+        ctx_request_id.set(request_id)
+    if trace_id is not None:
+        ctx_trace_id.set(trace_id)
+    if tenant_id is not None:
+        ctx_tenant_id.set(tenant_id)
+    if principal_id is not None:
+        ctx_principal_id.set(principal_id)
+
+
+def clear_request_context() -> None:
+    """Reset all correlation IDs in current context."""
+    ctx_request_id.set(None)
+    ctx_trace_id.set(None)
+    ctx_tenant_id.set(None)
+    ctx_principal_id.set(None)
+
+
+class JSONFormatter(logging.Formatter):
+    """Custom JSON formatter producing structured log entries for centralized log ingestion."""
+
+    _SENSITIVE_KEYS = {
+        "password",
+        "secret",
+        "token",
+        "authorization",
+        "key",
+        "api_key",
+        "access_token",
+        "secret_key",
+    }
+
+    def _sanitize(self, val: Any) -> Any:
+        if isinstance(val, dict):
+            return {
+                k: (
+                    "***REDACTED***"
+                    if any(s in k.lower() for s in self._SENSITIVE_KEYS)
+                    else self._sanitize(v)
+                )
+                for k, v in val.items()
+            }
+        if isinstance(val, (list, tuple)):
+            return [self._sanitize(item) for item in val]
+        return val
+
+    def format(self, record: logging.LogRecord) -> str:
+        payload: dict[str, Any] = {
+            "timestamp": datetime.datetime.now(datetime.UTC).isoformat(),
+            "level": record.levelname,
+            "logger": record.name,
+            "message": record.getMessage(),
+            "request_id": ctx_request_id.get(),
+            "trace_id": ctx_trace_id.get(),
+            "tenant_id": ctx_tenant_id.get(),
+            "principal_id": ctx_principal_id.get(),
+        }
+
+        if record.exc_info:
+            payload["exception"] = self.formatException(record.exc_info)
+
+        # Include custom extra fields attached to the log record
+        standard_attrs = {
+            "name",
+            "msg",
+            "args",
+            "levelname",
+            "levelno",
+            "pathname",
+            "filename",
+            "module",
+            "exc_info",
+            "exc_text",
+            "stack_info",
+            "lineno",
+            "funcName",
+            "created",
+            "msecs",
+            "relativeCreated",
+            "thread",
+            "threadName",
+            "processName",
+            "process",
+            "message",
+        }
+        extras = {
+            k: (
+                "***REDACTED***"
+                if any(s in k.lower() for s in self._SENSITIVE_KEYS)
+                else self._sanitize(v)
+            )
+            for k, v in record.__dict__.items()
+            if k not in standard_attrs and not k.startswith("_")
+        }
+        if extras:
+            payload["extra"] = extras
+
+        return json.dumps(payload, default=str)
+
+
+def configure_logging(level: str = "INFO") -> None:
+    """Configure root logger to use JSON formatting directed exclusively to stderr."""
+    root_logger = logging.getLogger()
+    root_logger.setLevel(level.upper())
+
+    # Remove any existing handlers
+    for handler in list(root_logger.handlers):
+        root_logger.removeHandler(handler)
+
+    # Output strictly to sys.stderr
+    stderr_handler = logging.StreamHandler(sys.stderr)
+    stderr_handler.setFormatter(JSONFormatter())
+    root_logger.addHandler(stderr_handler)
+
+    # Ensure critical third-party loggers do not write raw stdout
+    for noisy in ("uvicorn", "uvicorn.access", "uvicorn.error", "fastmcp"):
+        noisy_logger = logging.getLogger(noisy)
+        noisy_logger.handlers = []
+        noisy_logger.propagate = True
+
+
+def get_logger(name: str) -> logging.Logger:
+    """Return a logger instance by name."""
+    return logging.getLogger(name)

ml_mcp/infrastructure/telemetry/otel.py

@@ -0,0 +1,141 @@
+"""OpenTelemetry instrumentation setup for traces and metrics."""
+
+import logging
+
+from opentelemetry import metrics, trace
+from opentelemetry.sdk.metrics import MeterProvider
+from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+
+logger = logging.getLogger(__name__)
+
+
+# Global metric instruments container
+class MetricInstruments:
+    def __init__(self, meter: metrics.Meter) -> None:
+        self.meter = meter
+        self.mcp_requests_total = meter.create_counter(
+            name="mcp_requests_total",
+            description="Total number of MCP requests handled",
+            unit="1",
+        )
+        self.mcp_request_duration_seconds = meter.create_histogram(
+            name="mcp_request_duration_seconds",
+            description="Latency distribution of MCP requests",
+            unit="s",
+        )
+        self.mcp_errors_total = meter.create_counter(
+            name="mcp_errors_total",
+            description="Total number of MCP error responses returned",
+            unit="1",
+        )
+        self.auth_failures_total = meter.create_counter(
+            name="auth_failures_total",
+            description="Total authentication failures",
+            unit="1",
+        )
+        self.authorization_denials_total = meter.create_counter(
+            name="authorization_denials_total",
+            description="Total authorization rejections",
+            unit="1",
+        )
+        self.rate_limit_denials_total = meter.create_counter(
+            name="rate_limit_denials_total",
+            description="Total rate limit throttle events",
+            unit="1",
+        )
+        self.experiments_created_total = meter.create_counter(
+            name="experiments_created_total",
+            description="Total experiments created",
+            unit="1",
+        )
+        self.experiments_succeeded_total = meter.create_counter(
+            name="experiments_succeeded_total",
+            description="Total experiments successfully completed",
+            unit="1",
+        )
+        self.experiments_failed_total = meter.create_counter(
+            name="experiments_failed_total",
+            description="Total experiments that failed during execution",
+            unit="1",
+        )
+        self.experiment_runtime_seconds = meter.create_histogram(
+            name="experiment_runtime_seconds",
+            description="Execution duration of ML experiment runs",
+            unit="s",
+        )
+        self.artifact_bytes_written = meter.create_counter(
+            name="artifact_bytes_written",
+            description="Total bytes written to artifact store",
+            unit="By",
+        )
+        self.analysis_runtime_seconds = meter.create_histogram(
+            name="analysis_runtime_seconds",
+            description="Runtime of structured analysis evaluation",
+            unit="s",
+        )
+
+
+_instruments: MetricInstruments | None = None
+
+
+def init_telemetry(
+    service_name: str = "ml-mcp-service",
+    enabled: bool = False,
+    otlp_endpoint: str = "http://localhost:4317",
+) -> MetricInstruments:
+    """Initialize OpenTelemetry SDK with resource attributes and OTLP exporters if enabled."""
+    global _instruments
+    resource = Resource.create({"service.name": service_name})
+
+    if enabled:
+        try:
+            from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import OTLPMetricExporter
+            from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
+
+            # Traces
+            tracer_provider = TracerProvider(resource=resource)
+            span_exporter = OTLPSpanExporter(endpoint=otlp_endpoint, insecure=True)
+            tracer_provider.add_span_processor(BatchSpanProcessor(span_exporter))
+            trace.set_tracer_provider(tracer_provider)
+
+            # Metrics
+            metric_exporter = OTLPMetricExporter(endpoint=otlp_endpoint, insecure=True)
+            reader = PeriodicExportingMetricReader(metric_exporter, export_interval_millis=15000)
+            meter_provider = MeterProvider(resource=resource, metric_readers=[reader])
+            metrics.set_meter_provider(meter_provider)
+            logger.info("OpenTelemetry initialized with OTLP endpoint: %s", otlp_endpoint)
+        except Exception as exc:
+            logger.warning(
+                "Failed to initialize OTLP exporter, falling back to default providers: %s", exc
+            )
+            trace.set_tracer_provider(TracerProvider(resource=resource))
+            metrics.set_meter_provider(MeterProvider(resource=resource))
+    else:
+        trace.set_tracer_provider(TracerProvider(resource=resource))
+        metrics.set_meter_provider(MeterProvider(resource=resource))
+
+    meter = metrics.get_meter("ml-mcp", "0.1.0")
+    _instruments = MetricInstruments(meter)
+    return _instruments
+
+
+def get_tracer(name: str = "ml-mcp") -> trace.Tracer:
+    """Return an OpenTelemetry Tracer instance."""
+    return trace.get_tracer(name, "0.1.0")
+
+
+def get_meter(name: str = "ml-mcp") -> metrics.Meter:
+    """Return an OpenTelemetry Meter instance."""
+    return metrics.get_meter(name, "0.1.0")
+
+
+def get_instruments() -> MetricInstruments:
+    """Return instantiated metric counters and histograms."""
+    global _instruments
+    if _instruments is None:
+        meter = metrics.get_meter("ml-mcp", "0.1.0")
+        _instruments = MetricInstruments(meter)
+    return _instruments

ml_mcp/mcp/__init__.py

@@ -0,0 +1 @@
+"""MCP protocol integration layer."""