spatial-memory-mcp 1.9.1__py3-none-any.whl

spatial_memory/core/logging.py

@@ -0,0 +1,194 @@
+"""Secure structured logging for Spatial Memory MCP Server.
+
+This module provides secure logging with sensitive data masking and
+optional request context tracking for observability.
+
+Features:
+    - Sensitive data masking (API keys, passwords)
+    - JSON structured logging format
+    - Request context integration ([req=xxx][agent=yyy] prefixes)
+    - Configurable log levels and formats
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import re
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    pass
+
+# Patterns to mask in logs
+SENSITIVE_PATTERNS = [
+    (re.compile(r'api[_-]?key["\']?\s*[:=]\s*["\']?[\w-]+', re.I), 'api_key=***MASKED***'),
+    (re.compile(r'sk-[a-zA-Z0-9]{20,}'), '***OPENAI_KEY***'),
+    (re.compile(r'password["\']?\s*[:=]\s*["\']?[^\s"\']+', re.I), 'password=***MASKED***'),
+]
+
+
+def _get_trace_context() -> tuple[str | None, str | None]:
+    """Get request context without importing at module level.
+
+    Returns:
+        Tuple of (request_id, agent_id) or (None, None) if no context.
+    """
+    try:
+        # Import here to avoid circular imports
+        from spatial_memory.core.tracing import get_current_context
+
+        ctx = get_current_context()
+        if ctx:
+            return ctx.request_id, ctx.agent_id
+    except ImportError:
+        pass
+    return None, None
+
+
+class SecureFormatter(logging.Formatter):
+    """Formatter that masks sensitive data and includes trace context."""
+
+    def __init__(
+        self,
+        fmt: str | None = None,
+        datefmt: str | None = None,
+        include_trace_context: bool = True,
+    ) -> None:
+        """Initialize the secure formatter.
+
+        Args:
+            fmt: Format string for log messages.
+            datefmt: Date format string.
+            include_trace_context: Whether to include [req=xxx][agent=yyy] prefix.
+        """
+        super().__init__(fmt=fmt, datefmt=datefmt)
+        self.include_trace_context = include_trace_context
+
+    def format(self, record: logging.LogRecord) -> str:
+        """Format log record and mask sensitive data.
+
+        Args:
+            record: The log record to format.
+
+        Returns:
+            Formatted log message with sensitive data masked and trace context.
+        """
+        message = super().format(record)
+
+        # Add trace context prefix if available
+        if self.include_trace_context:
+            request_id, agent_id = _get_trace_context()
+            if request_id:
+                prefix_parts = [f"[req={request_id}]"]
+                if agent_id:
+                    prefix_parts.append(f"[agent={agent_id}]")
+                prefix = "".join(prefix_parts) + " "
+                # Insert after timestamp and logger name
+                # Format: "2024-01-15 10:30:00 - logger - LEVEL - message"
+                # We want: "2024-01-15 10:30:00 - logger - LEVEL - [req=xxx] message"
+                parts = message.split(" - ", 3)
+                if len(parts) == 4:
+                    message = f"{parts[0]} - {parts[1]} - {parts[2]} - {prefix}{parts[3]}"
+                else:
+                    # Fallback: just prepend
+                    message = prefix + message
+
+        for pattern, replacement in SENSITIVE_PATTERNS:
+            message = pattern.sub(replacement, message)
+        return message
+
+
+class JSONFormatter(logging.Formatter):
+    """JSON formatter for structured logging with trace context."""
+
+    def __init__(self, include_trace_context: bool = True) -> None:
+        """Initialize the JSON formatter.
+
+        Args:
+            include_trace_context: Whether to include request_id and agent_id fields.
+        """
+        super().__init__()
+        self.include_trace_context = include_trace_context
+
+    def format(self, record: logging.LogRecord) -> str:
+        """Format log record as JSON with sensitive data masked.
+
+        Args:
+            record: The log record to format.
+
+        Returns:
+            JSON-formatted log message with sensitive data masked.
+        """
+        log_data: dict[str, str | None] = {
+            "timestamp": datetime.now(timezone.utc).isoformat(),
+            "level": record.levelname,
+            "logger": record.name,
+            "message": record.getMessage(),
+        }
+
+        # Add trace context if available
+        if self.include_trace_context:
+            request_id, agent_id = _get_trace_context()
+            if request_id:
+                log_data["request_id"] = request_id
+            if agent_id:
+                log_data["agent_id"] = agent_id
+
+        if record.exc_info:
+            log_data["exception"] = self.formatException(record.exc_info)
+
+        # Mask sensitive data
+        json_str = json.dumps(log_data)
+        for pattern, replacement in SENSITIVE_PATTERNS:
+            json_str = pattern.sub(replacement, json_str)
+
+        return json_str
+
+
+def configure_logging(
+    level: str = "INFO",
+    json_format: bool = False,
+    mask_sensitive: bool = True,
+    include_trace_context: bool = True,
+) -> None:
+    """Configure logging for the application.
+
+    Args:
+        level: Logging level (DEBUG, INFO, WARNING, ERROR).
+        json_format: Use JSON format for structured logging.
+        mask_sensitive: Mask sensitive data in logs.
+        include_trace_context: Include [req=xxx][agent=yyy] in log messages.
+    """
+    # Get root logger
+    root_logger = logging.getLogger()
+    root_logger.setLevel(level)
+
+    # Remove existing handlers
+    for handler in root_logger.handlers[:]:
+        root_logger.removeHandler(handler)
+
+    # Create console handler
+    console_handler = logging.StreamHandler()
+    console_handler.setLevel(level)
+
+    # Choose formatter
+    if json_format:
+        formatter: logging.Formatter = JSONFormatter(
+            include_trace_context=include_trace_context
+        )
+    elif mask_sensitive:
+        formatter = SecureFormatter(
+            fmt="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+            datefmt="%Y-%m-%d %H:%M:%S",
+            include_trace_context=include_trace_context,
+        )
+    else:
+        formatter = logging.Formatter(
+            fmt="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+            datefmt="%Y-%m-%d %H:%M:%S",
+        )
+
+    console_handler.setFormatter(formatter)
+    root_logger.addHandler(console_handler)

spatial_memory/core/metrics.py

@@ -0,0 +1,192 @@
+"""Prometheus metrics for Spatial Memory MCP Server.
+
+This module provides optional Prometheus metrics. If prometheus_client is not
+installed, no-op stubs are provided so the code works without metrics.
+
+Usage:
+    from spatial_memory.core.metrics import (
+        record_request,
+        record_search_similarity,
+        record_embedding_latency,
+        update_memory_count,
+    )
+
+    with record_request("recall", "success"):
+        # ... do work
+        pass
+
+    record_search_similarity(0.85)
+    record_embedding_latency(0.234, model="openai")
+"""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from prometheus_client import Counter as CounterType
+    from prometheus_client import Gauge as GaugeType
+    from prometheus_client import Histogram as HistogramType
+
+try:
+    from prometheus_client import Counter, Gauge, Histogram
+
+    PROMETHEUS_AVAILABLE = True
+except ImportError:
+    PROMETHEUS_AVAILABLE = False
+    Counter = None  # type: ignore
+    Histogram = None  # type: ignore
+    Gauge = None  # type: ignore
+
+# Metrics definitions (only created if prometheus_client available)
+if PROMETHEUS_AVAILABLE:
+    # Request metrics
+    REQUESTS_TOTAL: CounterType = Counter(
+        "spatial_memory_requests_total",
+        "Total number of requests",
+        ["tool", "status"],
+    )
+    REQUEST_DURATION: HistogramType = Histogram(
+        "spatial_memory_request_duration_seconds",
+        "Request duration in seconds",
+        ["tool"],
+        buckets=(0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0),
+    )
+
+    # Memory metrics
+    MEMORIES_TOTAL: GaugeType = Gauge(
+        "spatial_memory_memories_total",
+        "Total number of memories",
+        ["namespace"],
+    )
+    INDEX_STATUS: GaugeType = Gauge(
+        "spatial_memory_index_status",
+        "Index status (1=exists, 0=missing)",
+        ["index_type"],
+    )
+
+    # Search metrics
+    SEARCH_SIMILARITY: HistogramType = Histogram(
+        "spatial_memory_search_similarity_score",
+        "Search result similarity scores",
+        buckets=(0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.95, 1.0),
+    )
+
+    # Embedding metrics
+    EMBEDDING_LATENCY: HistogramType = Histogram(
+        "spatial_memory_embedding_latency_seconds",
+        "Embedding generation latency in seconds",
+        ["model"],
+        buckets=(0.01, 0.05, 0.1, 0.25, 0.5, 1.0, 2.0, 5.0),
+    )
+else:
+    # No-op stubs when prometheus_client is not available
+    REQUESTS_TOTAL = None  # type: ignore
+    REQUEST_DURATION = None  # type: ignore
+    MEMORIES_TOTAL = None  # type: ignore
+    INDEX_STATUS = None  # type: ignore
+    SEARCH_SIMILARITY = None  # type: ignore
+    EMBEDDING_LATENCY = None  # type: ignore
+
+
+@contextmanager
+def record_request(tool: str, status: str = "success") -> Generator[None, None, None]:
+    """Context manager to record request metrics.
+
+    Args:
+        tool: Name of the tool being called.
+        status: Status of the request (success, error, etc.).
+
+    Yields:
+        None
+
+    Example:
+        with record_request("recall", "success"):
+            # ... do work
+            pass
+    """
+    if not PROMETHEUS_AVAILABLE:
+        yield
+        return
+
+    start = time.monotonic()
+    try:
+        yield
+    except Exception:
+        status = "error"
+        raise
+    finally:
+        duration = time.monotonic() - start
+        REQUESTS_TOTAL.labels(tool=tool, status=status).inc()
+        REQUEST_DURATION.labels(tool=tool).observe(duration)
+
+
+def record_search_similarity(similarity: float) -> None:
+    """Record a search result similarity score.
+
+    Args:
+        similarity: Similarity score between 0.0 and 1.0.
+
+    Example:
+        record_search_similarity(0.85)
+    """
+    if PROMETHEUS_AVAILABLE:
+        SEARCH_SIMILARITY.observe(similarity)
+
+
+def record_embedding_latency(duration: float, model: str = "local") -> None:
+    """Record embedding generation latency.
+
+    Args:
+        duration: Time taken to generate embeddings in seconds.
+        model: Model identifier (e.g., "local", "openai").
+
+    Example:
+        record_embedding_latency(0.234, model="openai")
+    """
+    if PROMETHEUS_AVAILABLE:
+        EMBEDDING_LATENCY.labels(model=model).observe(duration)
+
+
+def update_memory_count(namespace: str, count: int) -> None:
+    """Update memory count for a namespace.
+
+    Args:
+        namespace: The namespace identifier.
+        count: Total number of memories in the namespace.
+
+    Example:
+        update_memory_count("default", 1000)
+    """
+    if PROMETHEUS_AVAILABLE:
+        MEMORIES_TOTAL.labels(namespace=namespace).set(count)
+
+
+def update_index_status(index_type: str, exists: bool) -> None:
+    """Update index status.
+
+    Args:
+        index_type: Type of index (e.g., "vector", "fts", "scalar").
+        exists: Whether the index exists.
+
+    Example:
+        update_index_status("vector", True)
+    """
+    if PROMETHEUS_AVAILABLE:
+        INDEX_STATUS.labels(index_type=index_type).set(1 if exists else 0)
+
+
+def is_available() -> bool:
+    """Check if Prometheus metrics are available.
+
+    Returns:
+        True if prometheus_client is installed, False otherwise.
+
+    Example:
+        if is_available():
+            print("Metrics are available")
+    """
+    return PROMETHEUS_AVAILABLE