spatial-memory-mcp 1.0.2__py3-none-any.whl → 1.5.3__py3-none-any.whl

spatial_memory/core/circuit_breaker.py

@@ -0,0 +1,297 @@
+"""Circuit breaker pattern implementation for fault tolerance.
+
+The circuit breaker prevents cascading failures by fast-failing requests
+when a service is unhealthy, allowing time for recovery.
+
+State transitions:
+    CLOSED (normal) -> OPEN (failures >= threshold)
+    OPEN -> HALF_OPEN (after reset_timeout)
+    HALF_OPEN -> CLOSED (probe succeeds)
+    HALF_OPEN -> OPEN (probe fails)
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+import time
+from collections.abc import Callable
+from enum import Enum
+from typing import TypeVar
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+class CircuitState(Enum):
+    """Circuit breaker states.
+
+    CLOSED: Normal operation, requests pass through.
+    OPEN: Circuit is tripped, requests are rejected immediately.
+    HALF_OPEN: Testing recovery, limited requests allowed.
+    """
+
+    CLOSED = "closed"
+    OPEN = "open"
+    HALF_OPEN = "half_open"
+
+
+class CircuitOpenError(Exception):
+    """Raised when circuit is open and call is rejected.
+
+    This exception indicates that the circuit breaker is preventing
+    requests from reaching a failing service.
+    """
+
+    def __init__(
+        self,
+        message: str = "Circuit breaker is open",
+        time_until_retry: float | None = None,
+    ) -> None:
+        """Initialize CircuitOpenError.
+
+        Args:
+            message: Error description.
+            time_until_retry: Seconds until the circuit will transition to HALF_OPEN.
+        """
+        super().__init__(message)
+        self.time_until_retry = time_until_retry
+
+
+class CircuitBreaker:
+    """Circuit breaker for protecting external service calls.
+
+    Monitors failures and opens the circuit when failures exceed a threshold,
+    preventing further calls until a reset timeout has elapsed.
+
+    Example:
+        breaker = CircuitBreaker(failure_threshold=5, reset_timeout=60.0)
+
+        try:
+            result = breaker.call(my_api_call, arg1, arg2)
+        except CircuitOpenError:
+            # Service is unhealthy, use fallback
+            result = fallback_value
+
+    Thread Safety:
+        This class is thread-safe. All state transitions and counters
+        are protected by a lock.
+    """
+
+    def __init__(
+        self,
+        failure_threshold: int = 5,
+        reset_timeout: float = 60.0,
+        half_open_max_calls: int = 1,
+        name: str | None = None,
+    ) -> None:
+        """Initialize the circuit breaker.
+
+        Args:
+            failure_threshold: Number of consecutive failures before opening circuit.
+            reset_timeout: Seconds to wait before transitioning from OPEN to HALF_OPEN.
+            half_open_max_calls: Maximum concurrent calls allowed in HALF_OPEN state.
+            name: Optional name for logging purposes.
+        """
+        if failure_threshold < 1:
+            raise ValueError("failure_threshold must be at least 1")
+        if reset_timeout <= 0:
+            raise ValueError("reset_timeout must be positive")
+        if half_open_max_calls < 1:
+            raise ValueError("half_open_max_calls must be at least 1")
+
+        self._failure_threshold = failure_threshold
+        self._reset_timeout = reset_timeout
+        self._half_open_max_calls = half_open_max_calls
+        self._name = name or "circuit_breaker"
+
+        self._state = CircuitState.CLOSED
+        self._failure_count = 0
+        self._last_failure_time: float | None = None
+        self._half_open_calls = 0
+        self._lock = threading.Lock()
+
+        # Statistics
+        self._total_calls = 0
+        self._total_failures = 0
+        self._total_rejections = 0
+
+    @property
+    def state(self) -> CircuitState:
+        """Get current circuit state.
+
+        This property also handles automatic transition from OPEN to HALF_OPEN
+        when the reset timeout has elapsed.
+
+        Returns:
+            Current circuit state.
+        """
+        with self._lock:
+            return self._get_state_unlocked()
+
+    def _get_state_unlocked(self) -> CircuitState:
+        """Get state without acquiring lock (must be called with lock held)."""
+        if self._state == CircuitState.OPEN:
+            if self._should_transition_to_half_open():
+                self._transition_to_half_open()
+        return self._state
+
+    def _should_transition_to_half_open(self) -> bool:
+        """Check if reset timeout has elapsed (must be called with lock held)."""
+        if self._last_failure_time is None:
+            return False
+        elapsed = time.monotonic() - self._last_failure_time
+        return elapsed >= self._reset_timeout
+
+    def _transition_to_half_open(self) -> None:
+        """Transition to HALF_OPEN state (must be called with lock held)."""
+        logger.info(f"[{self._name}] Circuit transitioning from OPEN to HALF_OPEN")
+        self._state = CircuitState.HALF_OPEN
+        self._half_open_calls = 0
+
+    @property
+    def failure_count(self) -> int:
+        """Get current failure count."""
+        with self._lock:
+            return self._failure_count
+
+    @property
+    def stats(self) -> dict[str, int | str | float | None]:
+        """Get circuit breaker statistics.
+
+        Returns:
+            Dictionary with state, counters, and timing information.
+        """
+        with self._lock:
+            state = self._get_state_unlocked()
+            time_until_retry = None
+            if state == CircuitState.OPEN and self._last_failure_time is not None:
+                elapsed = time.monotonic() - self._last_failure_time
+                time_until_retry = max(0.0, self._reset_timeout - elapsed)
+
+            return {
+                "state": state.value,
+                "failure_count": self._failure_count,
+                "failure_threshold": self._failure_threshold,
+                "total_calls": self._total_calls,
+                "total_failures": self._total_failures,
+                "total_rejections": self._total_rejections,
+                "reset_timeout": self._reset_timeout,
+                "time_until_retry": time_until_retry,
+            }
+
+    def call(self, func: Callable[..., T], *args: object, **kwargs: object) -> T:
+        """Execute function with circuit breaker protection.
+
+        Args:
+            func: Function to execute.
+            *args: Positional arguments for the function.
+            **kwargs: Keyword arguments for the function.
+
+        Returns:
+            Return value of the function.
+
+        Raises:
+            CircuitOpenError: If circuit is OPEN and not ready for retry.
+            Exception: Any exception raised by the function.
+        """
+        with self._lock:
+            self._total_calls += 1
+            current_state = self._get_state_unlocked()
+
+            if current_state == CircuitState.OPEN:
+                self._total_rejections += 1
+                time_until_retry = None
+                if self._last_failure_time is not None:
+                    elapsed = time.monotonic() - self._last_failure_time
+                    time_until_retry = max(0.0, self._reset_timeout - elapsed)
+                raise CircuitOpenError(
+                    f"[{self._name}] Circuit is OPEN, rejecting call",
+                    time_until_retry=time_until_retry,
+                )
+
+            if current_state == CircuitState.HALF_OPEN:
+                if self._half_open_calls >= self._half_open_max_calls:
+                    self._total_rejections += 1
+                    raise CircuitOpenError(
+                        f"[{self._name}] Circuit is HALF_OPEN, max probe calls reached",
+                        time_until_retry=0.0,
+                    )
+                self._half_open_calls += 1
+
+        # Execute function outside the lock
+        try:
+            result = func(*args, **kwargs)
+            self._on_success()
+            return result
+        except Exception as e:
+            self._on_failure(e)
+            raise
+
+    def _on_success(self) -> None:
+        """Handle successful call.
+
+        In CLOSED state: Reset failure count.
+        In HALF_OPEN state: Transition to CLOSED.
+        """
+        with self._lock:
+            if self._state == CircuitState.HALF_OPEN:
+                logger.info(
+                    f"[{self._name}] Probe succeeded, circuit transitioning to CLOSED"
+                )
+                self._state = CircuitState.CLOSED
+                self._failure_count = 0
+                self._half_open_calls = 0
+            elif self._state == CircuitState.CLOSED:
+                self._failure_count = 0
+
+    def _on_failure(self, error: Exception) -> None:
+        """Handle failed call.
+
+        In CLOSED state: Increment failure count, open circuit if threshold reached.
+        In HALF_OPEN state: Transition back to OPEN.
+
+        Args:
+            error: The exception that was raised.
+        """
+        with self._lock:
+            self._total_failures += 1
+            self._failure_count += 1
+            self._last_failure_time = time.monotonic()
+
+            if self._state == CircuitState.HALF_OPEN:
+                logger.warning(
+                    f"[{self._name}] Probe failed ({error!r}), "
+                    f"circuit transitioning back to OPEN"
+                )
+                self._state = CircuitState.OPEN
+                self._half_open_calls = 0
+            elif self._state == CircuitState.CLOSED:
+                if self._failure_count >= self._failure_threshold:
+                    logger.warning(
+                        f"[{self._name}] Failure threshold reached "
+                        f"({self._failure_count}/{self._failure_threshold}), "
+                        f"circuit transitioning to OPEN"
+                    )
+                    self._state = CircuitState.OPEN
+
+    def reset(self) -> None:
+        """Manually reset circuit to CLOSED state.
+
+        This clears all failure counters and transitions the circuit
+        to CLOSED state regardless of current state.
+        """
+        with self._lock:
+            logger.info(f"[{self._name}] Circuit manually reset to CLOSED")
+            self._state = CircuitState.CLOSED
+            self._failure_count = 0
+            self._half_open_calls = 0
+            self._last_failure_time = None
+
+    def __repr__(self) -> str:
+        """Return string representation."""
+        return (
+            f"CircuitBreaker(name={self._name!r}, state={self.state.value}, "
+            f"failures={self.failure_count}/{self._failure_threshold})"
+        )

spatial_memory/core/database.py

@@ -35,7 +35,7 @@ from filelock import FileLock, Timeout as FileLockTimeout
 
 from spatial_memory.core.connection_pool import ConnectionPool
 from spatial_memory.core.errors import FileLockError, MemoryNotFoundError, StorageError, ValidationError
-from spatial_memory.core.utils import utc_now
+from spatial_memory.core.utils import to_aware_utc, utc_now
 
 # Import centralized validation functions
 from spatial_memory.core.validation import (
@@ -391,6 +391,15 @@ def with_process_lock(func: F) -> F:
 # Health Metrics
 # ============================================================================
 
+@dataclass
+class IdempotencyRecord:
+    """Record for idempotency key tracking."""
+    key: str
+    memory_id: str
+    created_at: Any  # datetime
+    expires_at: Any  # datetime
+
+
 @dataclass
 class IndexStats:
     """Statistics for a single index."""
@@ -3045,3 +3054,160 @@ class Database:
             return self.table.version
         except Exception as e:
             raise StorageError(f"Failed to get current version: {e}") from e
+
+    # ========================================================================
+    # Idempotency Key Management
+    # ========================================================================
+
+    def _ensure_idempotency_table(self) -> None:
+        """Ensure the idempotency keys table exists."""
+        if self._db is None:
+            raise StorageError("Database not connected")
+
+        existing_tables_result = self._db.list_tables()
+        if hasattr(existing_tables_result, 'tables'):
+            existing_tables = existing_tables_result.tables
+        else:
+            existing_tables = existing_tables_result
+
+        if "idempotency_keys" not in existing_tables:
+            schema = pa.schema([
+                pa.field("key", pa.string()),
+                pa.field("memory_id", pa.string()),
+                pa.field("created_at", pa.timestamp("us")),
+                pa.field("expires_at", pa.timestamp("us")),
+            ])
+            self._db.create_table("idempotency_keys", schema=schema)
+            logger.info("Created idempotency_keys table")
+
+    @property
+    def idempotency_table(self) -> LanceTable:
+        """Get the idempotency keys table, creating if needed."""
+        if self._db is None:
+            self.connect()
+        self._ensure_idempotency_table()
+        assert self._db is not None
+        return self._db.open_table("idempotency_keys")
+
+    def get_by_idempotency_key(self, key: str) -> IdempotencyRecord | None:
+        """Look up an idempotency record by key.
+
+        Args:
+            key: The idempotency key to look up.
+
+        Returns:
+            IdempotencyRecord if found and not expired, None otherwise.
+
+        Raises:
+            StorageError: If database operation fails.
+        """
+        if not key:
+            return None
+
+        try:
+            safe_key = _sanitize_string(key)
+            results = (
+                self.idempotency_table.search()
+                .where(f"key = '{safe_key}'")
+                .limit(1)
+                .to_list()
+            )
+
+            if not results:
+                return None
+
+            record = results[0]
+            now = utc_now()
+
+            # Check if expired (convert DB naive datetime to aware for comparison)
+            expires_at = record.get("expires_at")
+            if expires_at is not None:
+                expires_at_aware = to_aware_utc(expires_at)
+                if expires_at_aware < now:
+                    # Expired - clean it up and return None
+                    logger.debug(f"Idempotency key '{key}' has expired")
+                    return None
+
+            return IdempotencyRecord(
+                key=record["key"],
+                memory_id=record["memory_id"],
+                created_at=record["created_at"],
+                expires_at=record["expires_at"],
+            )
+
+        except Exception as e:
+            raise StorageError(f"Failed to look up idempotency key: {e}") from e
+
+    @with_process_lock
+    @with_write_lock
+    def store_idempotency_key(
+        self,
+        key: str,
+        memory_id: str,
+        ttl_hours: float = 24.0,
+    ) -> None:
+        """Store an idempotency key mapping.
+
+        Args:
+            key: The idempotency key.
+            memory_id: The memory ID that was created.
+            ttl_hours: Time-to-live in hours (default: 24 hours).
+
+        Raises:
+            ValidationError: If inputs are invalid.
+            StorageError: If database operation fails.
+        """
+        if not key:
+            raise ValidationError("Idempotency key cannot be empty")
+        if not memory_id:
+            raise ValidationError("Memory ID cannot be empty")
+        if ttl_hours <= 0:
+            raise ValidationError("TTL must be positive")
+
+        now = utc_now()
+        expires_at = now + timedelta(hours=ttl_hours)
+
+        record = {
+            "key": key,
+            "memory_id": memory_id,
+            "created_at": now,
+            "expires_at": expires_at,
+        }
+
+        try:
+            self.idempotency_table.add([record])
+            logger.debug(
+                f"Stored idempotency key '{key}' -> memory '{memory_id}' "
+                f"(expires in {ttl_hours}h)"
+            )
+        except Exception as e:
+            raise StorageError(f"Failed to store idempotency key: {e}") from e
+
+    @with_process_lock
+    @with_write_lock
+    def cleanup_expired_idempotency_keys(self) -> int:
+        """Remove expired idempotency keys.
+
+        Returns:
+            Number of keys removed.
+
+        Raises:
+            StorageError: If cleanup fails.
+        """
+        try:
+            now = utc_now()
+            count_before = self.idempotency_table.count_rows()
+
+            # Delete expired keys
+            predicate = f"expires_at < timestamp '{now.isoformat()}'"
+            self.idempotency_table.delete(predicate)
+
+            count_after = self.idempotency_table.count_rows()
+            deleted = count_before - count_after
+
+            if deleted > 0:
+                logger.info(f"Cleaned up {deleted} expired idempotency keys")
+
+            return deleted
+        except Exception as e:
+            raise StorageError(f"Failed to cleanup idempotency keys: {e}") from e

spatial_memory/core/embeddings.py

@@ -12,6 +12,11 @@ from typing import TYPE_CHECKING, Any, Literal, TypeVar
 
 import numpy as np
 
+from spatial_memory.core.circuit_breaker import (
+    CircuitBreaker,
+    CircuitOpenError,
+    CircuitState,
+)
 from spatial_memory.core.errors import ConfigurationError, EmbeddingError
 
 if TYPE_CHECKING:
@@ -158,6 +163,7 @@ class EmbeddingService:
 
     Supports local sentence-transformers models and optional OpenAI API.
     Uses ONNX Runtime by default for 2-3x faster inference.
+    Optionally uses a circuit breaker for fault tolerance with external services.
     """
 
     def __init__(
@@ -165,6 +171,10 @@ class EmbeddingService:
         model_name: str = "all-MiniLM-L6-v2",
         openai_api_key: str | Any | None = None,
         backend: EmbeddingBackend = "auto",
+        circuit_breaker: CircuitBreaker | None = None,
+        circuit_breaker_enabled: bool = True,
+        circuit_breaker_failure_threshold: int = 5,
+        circuit_breaker_reset_timeout: float = 60.0,
     ) -> None:
         """Initialize the embedding service.
 
@@ -174,6 +184,14 @@ class EmbeddingService:
                 Can be a string or a SecretStr (pydantic).
             backend: Inference backend. 'auto' uses ONNX if available (default),
                 'onnx' forces ONNX Runtime, 'pytorch' forces PyTorch.
+            circuit_breaker: Optional pre-configured circuit breaker instance.
+                If provided, other circuit breaker parameters are ignored.
+            circuit_breaker_enabled: Whether to enable circuit breaker for OpenAI calls.
+                Defaults to True. Only applies to OpenAI models.
+            circuit_breaker_failure_threshold: Number of consecutive failures before
+                opening the circuit. Default is 5.
+            circuit_breaker_reset_timeout: Seconds to wait before attempting recovery.
+                Default is 60.0 seconds.
         """
         self.model_name = model_name
         # Handle both plain strings and SecretStr (pydantic)
@@ -203,6 +221,23 @@ class EmbeddingService:
                     "OpenAI API key required for OpenAI embedding models"
                 )
 
+        # Circuit breaker for OpenAI API calls (optional)
+        if circuit_breaker is not None:
+            self._circuit_breaker: CircuitBreaker | None = circuit_breaker
+        elif circuit_breaker_enabled and self.use_openai:
+            self._circuit_breaker = CircuitBreaker(
+                failure_threshold=circuit_breaker_failure_threshold,
+                reset_timeout=circuit_breaker_reset_timeout,
+                name=f"embedding_service_{model_name}",
+            )
+            logger.info(
+                f"Circuit breaker enabled for embedding service "
+                f"(threshold={circuit_breaker_failure_threshold}, "
+                f"timeout={circuit_breaker_reset_timeout}s)"
+            )
+        else:
+            self._circuit_breaker = None
+
     def _load_local_model(self) -> None:
         """Load local sentence-transformers model with ONNX or PyTorch backend."""
         if self._model is not None:
@@ -300,6 +335,26 @@ class EmbeddingService:
             self._load_local_model()
         return self._active_backend or "pytorch"
 
+    @property
+    def circuit_state(self) -> CircuitState | None:
+        """Get the current circuit breaker state.
+
+        Returns:
+            CircuitState if circuit breaker is enabled, None otherwise.
+        """
+        if self._circuit_breaker is None:
+            return None
+        return self._circuit_breaker.state
+
+    @property
+    def circuit_breaker(self) -> CircuitBreaker | None:
+        """Get the circuit breaker instance.
+
+        Returns:
+            CircuitBreaker if enabled, None otherwise.
+        """
+        return self._circuit_breaker
+
     def embed(self, text: str) -> np.ndarray:
         """Generate embedding for a single text.
 
@@ -320,7 +375,7 @@ class EmbeddingService:
 
         # Generate embedding (outside lock to allow concurrent generation)
         if self.use_openai:
-            embedding = self._embed_openai([text])[0]
+            embedding = self._embed_openai_with_circuit_breaker([text])[0]
         else:
             embedding = self._embed_local([text])[0]
 
@@ -352,7 +407,7 @@ class EmbeddingService:
             return []
 
         if self.use_openai:
-            return self._embed_openai(texts)
+            return self._embed_openai_with_circuit_breaker(texts)
         else:
             return self._embed_local(texts)
 
@@ -387,6 +442,41 @@ class EmbeddingService:
             masked_error = _mask_api_key(str(e))
             raise EmbeddingError(f"Failed to generate embeddings: {masked_error}") from e
 
+    def _embed_openai_with_circuit_breaker(self, texts: list[str]) -> list[np.ndarray]:
+        """Generate embeddings using OpenAI API with circuit breaker protection.
+
+        Wraps the OpenAI embedding call with a circuit breaker to prevent
+        cascading failures when the API is unavailable.
+
+        Args:
+            texts: List of texts to embed.
+
+        Returns:
+            List of embedding vectors.
+
+        Raises:
+            EmbeddingError: If circuit is open or embedding generation fails.
+        """
+        if self._circuit_breaker is None:
+            # No circuit breaker, call directly
+            return self._embed_openai(texts)
+
+        try:
+            return self._circuit_breaker.call(self._embed_openai, texts)
+        except CircuitOpenError as e:
+            logger.warning(
+                f"Circuit breaker is open for embedding service, "
+                f"time until retry: {e.time_until_retry:.1f}s"
+                if e.time_until_retry is not None
+                else "Circuit breaker is open for embedding service"
+            )
+            raise EmbeddingError(
+                f"Embedding service temporarily unavailable (circuit open). "
+                f"Try again in {e.time_until_retry:.0f} seconds."
+                if e.time_until_retry is not None
+                else "Embedding service temporarily unavailable (circuit open)."
+            ) from e
+
     @retry_on_api_error(max_attempts=3, backoff=1.0)
     def _embed_openai(self, texts: list[str]) -> list[np.ndarray]:
         """Generate embeddings using OpenAI API with retry logic.