opencode-semantic-memory 0.1.0__py3-none-any.whl

opencode_memory/http_server.py

@@ -0,0 +1,464 @@
+"""HTTP server for opencode-memory with streamable HTTP transport.
+
+This allows multiple OpenCode sessions to share a single memory server,
+avoiding the overhead of loading the embedding model multiple times.
+
+Usage:
+    # Start the server (once, e.g., via systemd user service)
+    python -m opencode_memory.http_server
+
+    # Configure OpenCode to connect via remote MCP
+    # In opencode.json:
+    # {
+    #   "mcp": {
+    #     "memory": {
+    #       "type": "remote",
+    #       "url": "http://localhost:9824/mcp"
+    #     }
+    #   }
+    # }
+
+Environment variables:
+    OPENCODE_MEMORY_HOST: Host to bind to (default: 127.0.0.1)
+    OPENCODE_MEMORY_PORT: Port to bind to (default: 9824)
+    OPENCODE_MEMORY_API_KEY: Optional API key for authentication
+    OPENCODE_MEMORY_RATE_LIMIT: Requests per minute per client (default: 60)
+"""
+
+import asyncio
+import logging
+import os
+import secrets
+import sys
+import time
+from collections import defaultdict
+from contextlib import asynccontextmanager
+from typing import Any
+
+from mcp.server.streamable_http import StreamableHTTPServerTransport
+
+from opencode_memory import metrics
+from opencode_memory.daemon import background_throttle
+from opencode_memory.server import MemoryServer
+
+try:
+    import uvicorn
+    from starlette.applications import Starlette
+    from starlette.responses import JSONResponse
+except ImportError:
+    print("HTTP server requires additional dependencies:")
+    print("  pip install starlette uvicorn")
+    sys.exit(1)
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger(__name__)
+
+DEFAULT_HOST = os.environ.get("OPENCODE_MEMORY_HOST", "127.0.0.1")
+DEFAULT_PORT = int(os.environ.get("OPENCODE_MEMORY_PORT", "9824"))
+API_KEY = os.environ.get("OPENCODE_MEMORY_API_KEY")  # Optional auth
+RATE_LIMIT = int(os.environ.get("OPENCODE_MEMORY_RATE_LIMIT", "60"))  # per minute
+
+
+class RateLimiter:
+    """Simple in-memory rate limiter using sliding window."""
+
+    def __init__(self, requests_per_minute: int = 60):
+        self.requests_per_minute = requests_per_minute
+        self.window_seconds = 60
+        self._requests: dict[str, list[float]] = defaultdict(list)
+
+    def _get_client_id(self, scope: dict) -> str:
+        """Get client identifier from request."""
+        # Use X-Forwarded-For if behind proxy, otherwise use client IP
+        headers = dict(scope.get("headers", []))
+        forwarded = headers.get(b"x-forwarded-for", b"").decode()
+        if forwarded:
+            client_ip = forwarded.split(",")[0].strip()
+        else:
+            client = scope.get("client", ("unknown", 0))
+            client_ip = client[0] if client else "unknown"
+        return client_ip
+
+    def is_allowed(self, scope: dict) -> tuple[bool, int]:
+        """Check if request is allowed. Returns (allowed, retry_after_seconds)."""
+        client_id = self._get_client_id(scope)
+        now = time.time()
+        window_start = now - self.window_seconds
+
+        # Clean old requests
+        self._requests[client_id] = [ts for ts in self._requests[client_id] if ts > window_start]
+
+        # Check limit
+        if len(self._requests[client_id]) >= self.requests_per_minute:
+            # Calculate retry-after
+            oldest = min(self._requests[client_id])
+            retry_after = int(oldest + self.window_seconds - now) + 1
+            return False, retry_after
+
+        # Record request
+        self._requests[client_id].append(now)
+        return True, 0
+
+    def cleanup(self) -> None:
+        """Remove stale entries (call periodically)."""
+        now = time.time()
+        window_start = now - self.window_seconds
+        stale_clients = [
+            c for c, reqs in self._requests.items() if all(ts < window_start for ts in reqs)
+        ]
+        for client in stale_clients:
+            del self._requests[client]
+
+
+_rate_limiter = RateLimiter(RATE_LIMIT)
+
+_memory_server: MemoryServer | None = None
+
+
+class ClientTracker:
+    """Track active clients based on recent request activity."""
+
+    def __init__(self, active_window_seconds: int = 300):
+        """
+        Args:
+            active_window_seconds: Consider client active if seen within this window (default 5 min)
+        """
+        self.active_window = active_window_seconds
+        self._clients: dict[str, dict] = {}  # client_id -> {last_seen, request_count, first_seen}
+
+    def record_request(self, scope: dict) -> None:
+        """Record a request from a client."""
+        client_id = self._get_client_id(scope)
+        now = time.time()
+
+        if client_id not in self._clients:
+            self._clients[client_id] = {
+                "first_seen": now,
+                "last_seen": now,
+                "request_count": 1,
+            }
+        else:
+            self._clients[client_id]["last_seen"] = now
+            self._clients[client_id]["request_count"] += 1
+
+    def _get_client_id(self, scope: dict) -> str:
+        """Get client identifier from request.
+
+        Uses X-OpenCode-Session header if present (set by opencode),
+        otherwise falls back to IP address.
+        """
+        headers = dict(scope.get("headers", []))
+
+        # Check for opencode session header (most reliable)
+        session_id = headers.get(b"x-opencode-session", b"").decode()
+        if session_id:
+            return f"session:{session_id[:16]}"  # Truncate for readability
+
+        # Fall back to IP (without port - ephemeral ports change per request)
+        forwarded = headers.get(b"x-forwarded-for", b"").decode()
+        if forwarded:
+            return forwarded.split(",")[0].strip()
+
+        client = scope.get("client", ("unknown", 0))
+        return client[0] if client else "unknown"
+
+    def get_active_clients(self) -> dict:
+        """Get info about currently active clients."""
+        now = time.time()
+        cutoff = now - self.active_window
+
+        active = []
+        stale = []
+
+        for client_id, info in self._clients.items():
+            if info["last_seen"] > cutoff:
+                seconds_ago = int(now - info["last_seen"])
+                active.append(
+                    {
+                        "client": client_id,
+                        "last_seen_seconds_ago": seconds_ago,
+                        "requests": info["request_count"],
+                        "session_duration_seconds": int(now - info["first_seen"]),
+                    }
+                )
+            else:
+                stale.append(client_id)
+
+        # Clean up stale entries
+        for client_id in stale:
+            del self._clients[client_id]
+
+        return {
+            "active_count": len(active),
+            "active_window_seconds": self.active_window,
+            "clients": sorted(active, key=lambda x: x["last_seen_seconds_ago"]),
+        }
+
+
+_client_tracker = ClientTracker()
+
+
+def get_active_clients() -> dict | None:
+    """Get active client info. Returns None if tracker not initialized."""
+    return _client_tracker.get_active_clients()
+
+
+def _check_auth(scope: dict) -> bool:
+    """Check API key authentication if configured."""
+    if not API_KEY:
+        return True  # No auth configured
+
+    headers = dict(scope.get("headers", []))
+
+    # Check Authorization header (Bearer token)
+    auth = headers.get(b"authorization", b"").decode()
+    if auth.startswith("Bearer "):
+        token = auth[7:]
+        if secrets.compare_digest(token, API_KEY):
+            return True
+
+    # Check X-API-Key header
+    api_key = headers.get(b"x-api-key", b"").decode()
+    if api_key:
+        if secrets.compare_digest(api_key, API_KEY):
+            return True
+
+    metrics.auth_failures.inc()
+    return False
+
+
+async def handle_mcp(scope: dict[str, Any], receive: Any, send: Any) -> None:
+    """Handle MCP requests via streamable HTTP transport - stateless mode."""
+    # Check authentication
+    if not _check_auth(scope):
+        response = JSONResponse(
+            {"error": "Unauthorized", "message": "Invalid or missing API key"},
+            status_code=401,
+        )
+        await response(scope, receive, send)
+        return
+
+    # Track client activity
+    _client_tracker.record_request(scope)
+
+    # Check rate limit
+    allowed, retry_after = _rate_limiter.is_allowed(scope)
+    if not allowed:
+        metrics.rate_limit_rejections.inc()
+        response = JSONResponse(
+            {
+                "error": "Rate limited",
+                "message": f"Too many requests. Retry after {retry_after} seconds.",
+                "retry_after": retry_after,
+            },
+            status_code=429,
+            headers={"Retry-After": str(retry_after)},
+        )
+        await response(scope, receive, send)
+        return
+
+    if _memory_server is None:
+        response = JSONResponse({"error": "Server not initialized"}, status_code=503)
+        await response(scope, receive, send)
+        return
+
+    transport = StreamableHTTPServerTransport(
+        mcp_session_id=None,
+        is_json_response_enabled=True,
+    )
+    ready_event = asyncio.Event()
+
+    async def run_server():
+        async with transport.connect() as (read_stream, write_stream):
+            ready_event.set()
+            await _memory_server.server.run(
+                read_stream,
+                write_stream,
+                _memory_server.server.create_initialization_options(),
+                stateless=True,
+            )
+
+    async with background_throttle.priority_context():
+        server_task = asyncio.create_task(run_server())
+        await ready_event.wait()
+        try:
+            await transport.handle_request(scope, receive, send)
+        finally:
+            await transport.terminate()
+            server_task.cancel()
+            try:
+                await server_task
+            except asyncio.CancelledError:
+                pass
+
+
+async def _prewarm_embeddings(server: MemoryServer) -> None:
+    """Pre-warm the embedding model in background."""
+    import asyncio
+
+    await asyncio.sleep(0.5)  # Let server finish startup first
+    try:
+        logger.info("Pre-warming embedding model...")
+        await server.embeddings.embed_async("warmup")
+        logger.info("Embedding model ready")
+    except Exception as e:
+        logger.warning(f"Failed to pre-warm embeddings: {e}")
+
+
+@asynccontextmanager
+async def lifespan(app):
+    """Manage server lifecycle."""
+    global _memory_server
+
+    logger.info("Starting opencode-memory HTTP server...")
+
+    _memory_server = MemoryServer(enable_daemon=True)
+
+    if _memory_server.daemon:
+        await _memory_server.daemon.start()
+
+    asyncio.create_task(_prewarm_embeddings(_memory_server))
+
+    logger.info(f"Server ready at http://{DEFAULT_HOST}:{DEFAULT_PORT}")
+    logger.info(f"MCP endpoint: http://{DEFAULT_HOST}:{DEFAULT_PORT}/mcp")
+
+    yield
+
+    logger.info("Shutting down...")
+
+    if _memory_server:
+        # Wait for any pending background embeddings
+        still_pending = await MemoryServer.wait_for_pending_embeddings(timeout=5.0)
+        if still_pending > 0:
+            logger.warning(f"Shutdown with {still_pending} embeddings still pending")
+
+        await _memory_server.enricher.close()
+        if _memory_server.daemon:
+            await _memory_server.daemon.stop()
+
+
+class MCPApp:
+    """ASGI app that routes to MCP or health endpoints."""
+
+    async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None:
+        if scope["type"] != "http":
+            return
+
+        path = scope.get("path", "")
+
+        if path == "/health":
+            # Health check doesn't require auth - returns basic health status
+            health_status = {"status": "ok", "server": "opencode-memory"}
+
+            # Add daemon health info if available
+            if _memory_server and _memory_server.daemon:
+                health_status["daemon"] = {
+                    "running": _memory_server.daemon.is_running,
+                }
+
+            # Check for degraded state (e.g., embedding queue backed up)
+            if _memory_server:
+                pending = len(MemoryServer._pending_embeddings)
+                if pending > 50:
+                    health_status["status"] = "degraded"
+                    health_status["warning"] = f"Embedding queue backed up: {pending} pending"
+
+            status_code = 200 if health_status["status"] == "ok" else 503
+            response = JSONResponse(health_status, status_code=status_code)
+            await response(scope, receive, send)
+        elif path == "/stats":
+            # Stats endpoint for monitoring
+            if not _check_auth(scope):
+                response = JSONResponse({"error": "Unauthorized"}, status_code=401)
+                await response(scope, receive, send)
+                return
+
+            stats = {
+                "status": "ok",
+                "server": "opencode-memory",
+                "rate_limit": {
+                    "requests_per_minute": RATE_LIMIT,
+                },
+                "clients": _client_tracker.get_active_clients(),
+            }
+            if _memory_server:
+                status = _memory_server._get_status()
+                stats["embedding_queue"] = status.get("embedding_queue")
+                stats["storage"] = status.get("storage")
+                stats["memories"] = status.get("memories")
+                stats["cache"] = status.get("cache")
+                stats["links"] = status.get("links")
+                if _memory_server.daemon:
+                    stats["daemon"] = {
+                        "running": _memory_server.daemon.is_running,
+                    }
+
+            response = JSONResponse(stats)
+            await response(scope, receive, send)
+        elif path == "/metrics":
+            # Prometheus metrics endpoint
+            if _memory_server:
+                status = _memory_server._get_status()
+                metrics.update_from_status(status)
+
+            from starlette.responses import Response
+
+            response = Response(
+                content=metrics.registry.render(),
+                media_type="text/plain; version=0.0.4; charset=utf-8",
+            )
+            await response(scope, receive, send)
+        elif path == "/mcp":
+            await handle_mcp(scope, receive, send)
+        else:
+            response = JSONResponse({"error": "Not found"}, status_code=404)
+            await response(scope, receive, send)
+
+
+def create_app() -> Starlette:
+    """Create the Starlette application."""
+    return Starlette(
+        routes=[],
+        lifespan=lifespan,
+        middleware=[],
+    )
+
+
+def create_asgi_app():
+    """Create ASGI app with MCP handling."""
+    starlette_app = create_app()
+
+    async def app(scope: dict[str, Any], receive: Any, send: Any) -> None:
+        if scope["type"] == "lifespan":
+            await starlette_app(scope, receive, send)
+        else:
+            mcp_app = MCPApp()
+            await mcp_app(scope, receive, send)
+
+    return app
+
+
+def main(host: str = DEFAULT_HOST, port: int = DEFAULT_PORT) -> None:
+    """Run the HTTP server."""
+    import argparse
+
+    parser = argparse.ArgumentParser(description="OpenCode Memory HTTP Server")
+    parser.add_argument("--host", default=host, help=f"Host to bind to (default: {host})")
+    parser.add_argument("--port", type=int, default=port, help=f"Port to bind to (default: {port})")
+    args = parser.parse_args()
+
+    app = create_asgi_app()
+
+    uvicorn.run(
+        app,
+        host=args.host,
+        port=args.port,
+        log_level="info",
+    )
+
+
+if __name__ == "__main__":
+    main()

opencode_memory/ingestion/__init__.py

@@ -0,0 +1,7 @@
+"""Ingestion modules for opencode-memory."""
+
+from opencode_memory.ingestion.embeddings import EmbeddingEngine
+from opencode_memory.ingestion.parser import MarkdownParser
+from opencode_memory.ingestion.watcher import FileWatcher
+
+__all__ = ["EmbeddingEngine", "MarkdownParser", "FileWatcher"]

opencode_memory/ingestion/embeddings.py

@@ -0,0 +1,211 @@
+"""Local embedding generation using sentence-transformers.
+
+All processing is done locally - no data is sent to any external service.
+The model is downloaded once from HuggingFace and cached locally.
+
+Uses ProcessPoolExecutor to avoid GIL blocking during CPU-intensive embedding.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import multiprocessing
+import os
+from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
+from typing import TYPE_CHECKING, Any
+
+# Use 'spawn' instead of 'fork' to avoid deadlocks in multi-threaded processes
+# This is the default on macOS and Windows, but not on Linux
+try:
+    multiprocessing.set_start_method("spawn", force=False)
+except RuntimeError:
+    # Already set, ignore
+    pass
+
+if TYPE_CHECKING:
+    from sentence_transformers import SentenceTransformer
+
+DEFAULT_EMBEDDING_DIM = 384
+DEFAULT_MODEL_NAME = "all-MiniLM-L6-v2"
+
+_thread_executor: ThreadPoolExecutor | None = None
+_process_executor: ProcessPoolExecutor | None = None
+_subprocess_model: Any = None
+_atexit_registered: bool = False
+
+
+def _shutdown_executors() -> None:
+    """Shutdown executors on process exit to prevent resource leaks."""
+    global _thread_executor, _process_executor
+    if _thread_executor is not None:
+        _thread_executor.shutdown(wait=False)
+        _thread_executor = None
+    if _process_executor is not None:
+        _process_executor.shutdown(wait=False)
+        _process_executor = None
+
+
+def _get_thread_executor() -> ThreadPoolExecutor:
+    """Get or create the thread pool executor (for model loading)."""
+    global _thread_executor, _atexit_registered
+    if _thread_executor is None:
+        _thread_executor = ThreadPoolExecutor(max_workers=2, thread_name_prefix="embedding")
+        if not _atexit_registered:
+            import atexit
+
+            atexit.register(_shutdown_executors)
+            _atexit_registered = True
+    return _thread_executor
+
+
+def _get_process_executor() -> ProcessPoolExecutor:
+    """Get or create the process pool executor (for CPU-intensive embedding)."""
+    global _process_executor, _atexit_registered
+    if _process_executor is None:
+        _process_executor = ProcessPoolExecutor(max_workers=2)
+        if not _atexit_registered:
+            import atexit
+
+            atexit.register(_shutdown_executors)
+            _atexit_registered = True
+    return _process_executor
+
+
+def _configure_thread_limits() -> None:
+    """Configure thread limits before any torch/numpy imports."""
+    num_threads = max(1, (os.cpu_count() or 4) // 4)  # ~5-6 threads per worker
+    os.environ["OMP_NUM_THREADS"] = str(num_threads)
+    os.environ["MKL_NUM_THREADS"] = str(num_threads)
+    os.environ["OPENBLAS_NUM_THREADS"] = str(num_threads)
+    os.environ["VECLIB_MAXIMUM_THREADS"] = str(num_threads)
+    os.environ["NUMEXPR_NUM_THREADS"] = str(num_threads)
+
+
+def _load_model(model_name: str) -> Any:
+    """Load the sentence transformer model (deferred import)."""
+    _configure_thread_limits()  # Must be before torch import
+    os.environ["HF_HUB_OFFLINE"] = "1"
+    os.environ["TRANSFORMERS_OFFLINE"] = "1"
+
+    import torch
+
+    num_threads = max(1, (os.cpu_count() or 4) // 4)
+    torch.set_num_threads(num_threads)
+    torch.set_num_interop_threads(2)
+
+    from sentence_transformers import SentenceTransformer
+
+    return SentenceTransformer(model_name)
+
+
+def _embed_in_subprocess(text: str, model_name: str = DEFAULT_MODEL_NAME) -> list[float]:
+    """Embed text in subprocess (loads model if needed)."""
+    global _subprocess_model
+    if _subprocess_model is None:
+        _subprocess_model = _load_model(model_name)
+    embedding = _subprocess_model.encode(text, convert_to_numpy=True)
+    return embedding.tolist()
+
+
+def _embed_batch_in_subprocess(
+    texts: list[str], model_name: str = DEFAULT_MODEL_NAME
+) -> list[list[float]]:
+    """Embed batch in subprocess (loads model if needed)."""
+    global _subprocess_model
+    if _subprocess_model is None:
+        _subprocess_model = _load_model(model_name)
+    embeddings = _subprocess_model.encode(texts, convert_to_numpy=True)
+    return embeddings.tolist()
+
+
+class EmbeddingEngine:
+    """Generate embeddings using a local sentence-transformers model.
+
+    Uses ProcessPoolExecutor for all operations to avoid loading the model
+    in the main process. This saves ~200-400MB of memory since the model
+    only lives in the subprocess workers.
+    """
+
+    def __init__(self, model_name: str = DEFAULT_MODEL_NAME, use_process_pool: bool = True):
+        self.model_name = model_name
+        self._model: SentenceTransformer | None = None
+        self._dimension: int | None = None
+        self._use_process_pool = use_process_pool
+
+    @property
+    def model(self) -> SentenceTransformer:
+        """Lazy-load the model in main process.
+
+        Note: Prefer embed() which uses subprocess to avoid loading model
+        in main process.
+        """
+        if self._model is None:
+            self._model = _load_model(self.model_name)
+            self._dimension = self._model.get_embedding_dimension()
+        return self._model
+
+    @property
+    def dimension(self) -> int:
+        """Get the embedding dimension without loading the model.
+
+        Returns the known dimension for all-MiniLM-L6-v2 (384) to avoid
+        loading the model just to get the dimension during startup.
+        """
+        if self._dimension is not None:
+            return self._dimension
+        return DEFAULT_EMBEDDING_DIM
+
+    def embed(self, text: str) -> list[float]:
+        """Generate embedding for a single text.
+
+        When use_process_pool=True (default), runs in subprocess to avoid
+        loading model in main process.
+        """
+        if self._use_process_pool:
+            # Run in subprocess synchronously via executor
+            future = _get_process_executor().submit(_embed_in_subprocess, text, self.model_name)
+            return future.result()
+        else:
+            embedding = self.model.encode(text, convert_to_numpy=True)
+            return embedding.tolist()  # type: ignore[no-any-return]
+
+    def embed_batch(self, texts: list[str]) -> list[list[float]]:
+        """Generate embeddings for multiple texts.
+
+        When use_process_pool=True (default), runs in subprocess to avoid
+        loading model in main process.
+        """
+        if self._use_process_pool:
+            future = _get_process_executor().submit(
+                _embed_batch_in_subprocess, texts, self.model_name
+            )
+            return future.result()
+        else:
+            embeddings = self.model.encode(texts, convert_to_numpy=True)
+            return embeddings.tolist()  # type: ignore[no-any-return]
+
+    async def embed_async(self, text: str) -> list[float]:
+        """Generate embedding for a single text (async, runs in process pool).
+
+        Uses ProcessPoolExecutor to avoid GIL blocking the event loop.
+        """
+        loop = asyncio.get_event_loop()
+        if self._use_process_pool:
+            return await loop.run_in_executor(
+                _get_process_executor(), _embed_in_subprocess, text, self.model_name
+            )
+        else:
+            return await loop.run_in_executor(_get_thread_executor(), self.embed, text)
+
+    async def embed_batch_async(self, texts: list[str]) -> list[list[float]]:
+        """Generate embeddings for multiple texts (async, runs in process pool).
+
+        Uses ProcessPoolExecutor to avoid GIL blocking the event loop.
+        """
+        loop = asyncio.get_event_loop()
+        if self._use_process_pool:
+            return await loop.run_in_executor(
+                _get_process_executor(), _embed_batch_in_subprocess, texts, self.model_name
+            )
+        else:
+            return await loop.run_in_executor(_get_thread_executor(), self.embed_batch, texts)