sari 0.0.1__py3-none-any.whl

mcp/proxy.py

@@ -0,0 +1,304 @@
+import sys
+import json
+import socket
+import threading
+import os
+import time
+import subprocess
+import logging
+import sys
+import tempfile
+from pathlib import Path
+
+# Add project root to sys.path for absolute imports
+SCRIPT_DIR = Path(__file__).parent
+REPO_ROOT = SCRIPT_DIR.parent
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from mcp.telemetry import TelemetryLogger
+from app.workspace import WorkspaceManager
+
+try:
+    import fcntl  # type: ignore
+except Exception:
+    fcntl = None
+
+# Configure logging to stderr so it doesn't interfere with MCP STDIO
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    stream=sys.stderr
+)
+logger = logging.getLogger("mcp-proxy")
+telemetry = TelemetryLogger(WorkspaceManager.get_global_log_dir())
+
+
+def _log_info(message: str) -> None:
+    logger.info(message)
+    try:
+        telemetry.log_info(message)
+    except Exception:
+        pass
+
+
+def _log_error(message: str) -> None:
+    logger.error(message)
+    try:
+        telemetry.log_error(message)
+    except Exception:
+        pass
+
+DEFAULT_HOST = "127.0.0.1"
+DEFAULT_PORT = 47779
+_HEADER_SEP = b"\r\n\r\n"
+_MODE_FRAMED = "framed"
+_MODE_JSONL = "jsonl"
+
+def _lock_file(lock_file) -> None:
+    if fcntl is not None:
+        fcntl.flock(lock_file, fcntl.LOCK_EX)
+        return
+    try:
+        import msvcrt
+        lock_file.seek(0)
+        msvcrt.locking(lock_file.fileno(), msvcrt.LK_LOCK, 1)
+    except Exception:
+        pass
+
+def _unlock_file(lock_file) -> None:
+    if fcntl is not None:
+        fcntl.flock(lock_file, fcntl.LOCK_UN)
+        return
+    try:
+        import msvcrt
+        lock_file.seek(0)
+        msvcrt.locking(lock_file.fileno(), msvcrt.LK_UNLCK, 1)
+    except Exception:
+        pass
+
+def start_daemon_if_needed(host, port):
+    """Checks if daemon is running, if not starts it."""
+    try:
+        with socket.create_connection((host, port), timeout=0.1):
+            return True
+    except (ConnectionRefusedError, OSError):
+        pass
+
+    lock_path = os.path.join(tempfile.gettempdir(), f"sari-daemon-{host}-{port}.lock")
+    with open(lock_path, "w") as lock_file:
+        try:
+            # Acquire exclusive lock (blocking)
+            _lock_file(lock_file)
+            
+            # Double-check if daemon started while waiting for lock
+            try:
+                with socket.create_connection((host, port), timeout=0.1):
+                    return True
+            except (ConnectionRefusedError, OSError):
+                pass
+
+            _log_info("Daemon not running, starting...")
+            
+            # Assume we are in mcp/proxy.py, so parent of parent is repo root
+            repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+            env = os.environ.copy()
+            # Do not infer workspace root from marker; rely on explicit roots/env/config.
+            
+            # Detach process
+            subprocess.Popen(
+                [sys.executable, "-m", "mcp.daemon"],
+                cwd=repo_root,
+                env=env,
+                start_new_session=True,
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL
+            )
+            
+            # Wait for it to come up
+            for _ in range(20):
+                try:
+                    with socket.create_connection((host, port), timeout=0.1):
+                        _log_info("Daemon started successfully.")
+                        return True
+                except (ConnectionRefusedError, OSError):
+                    time.sleep(0.1)
+            
+            _log_error("Failed to start daemon.")
+            return False
+            
+        finally:
+            _unlock_file(lock_file)
+
+def forward_socket_to_stdout(sock, mode_holder):
+    try:
+        f = sock.makefile("rb")
+        while True:
+            # Read Headers
+            headers = {}
+            while True:
+                line = f.readline()
+                if not line:
+                    break
+                line_str = line.decode("utf-8").strip()
+                if not line_str:
+                    break
+                if ":" in line_str:
+                    k, v = line_str.split(":", 1)
+                    headers[k.strip().lower()] = v.strip()
+            
+            if not headers and not line:
+                break
+            
+            content_length = int(headers.get("content-length", 0))
+            if content_length <= 0:
+                continue
+                
+            body = f.read(content_length)
+            if not body:
+                break
+                
+            mode = mode_holder.get("mode") or _MODE_FRAMED
+            if mode == _MODE_JSONL:
+                sys.stdout.buffer.write(body + b"\n")
+                sys.stdout.buffer.flush()
+            else:
+                header = f"Content-Length: {len(body)}\r\n\r\n".encode("ascii")
+                sys.stdout.buffer.write(header + body)
+                sys.stdout.buffer.flush()
+    except Exception as e:
+        _log_error(f"Error forwarding socket to stdout: {e}")
+    finally:
+        # If socket closes, we should probably exit
+        os._exit(0)
+
+def _read_mcp_message(stdin):
+    """Read one MCP framed message (Content-Length) or JSONL fallback."""
+    line = stdin.readline()
+    if not line:
+        return None
+    while line in (b"\n", b"\r\n"):
+        line = stdin.readline()
+        if not line:
+            return None
+
+    if line.lstrip().startswith((b"{", b"[")):
+        return line.rstrip(b"\r\n"), _MODE_JSONL
+
+    headers = [line]
+    while True:
+        h = stdin.readline()
+        if not h:
+            return None
+        if h in (b"\n", b"\r\n"):
+            break
+        headers.append(h)
+
+    content_length = None
+    for h in headers:
+        parts = h.decode("utf-8", errors="ignore").split(":", 1)
+        if len(parts) != 2:
+            continue
+        key = parts[0].strip().lower()
+        if key == "content-length":
+            try:
+                content_length = int(parts[1].strip())
+            except ValueError:
+                pass
+            break
+
+    if content_length is None:
+        return None
+
+    body = stdin.read(content_length)
+    if not body:
+        return None
+    return body, _MODE_FRAMED
+
+
+def forward_stdin_to_socket(sock, mode_holder):
+    try:
+        stdin = sys.stdin.buffer
+        while True:
+            res = _read_mcp_message(stdin)
+            if res is None:
+                break
+            msg, mode = res
+            # Inject rootUri for initialize when client omits it (per-connection workspace)
+            try:
+                req = json.loads(msg.decode("utf-8"))
+
+                def _inject(obj):
+                    if not isinstance(obj, dict) or obj.get("method") != "initialize":
+                        return obj, False
+                    params = obj.get("params") or {}
+                    if params.get("rootUri") or params.get("rootPath"):
+                        return obj, False
+                    ws = os.environ.get("DECKARD_WORKSPACE_ROOT") or os.environ.get("LOCAL_SEARCH_WORKSPACE_ROOT")
+                    if not ws:
+                        return obj, False
+                    params = dict(params)
+                    params["rootUri"] = f"file://{ws}"
+                    obj = dict(obj)
+                    obj["params"] = params
+                    _log_info(f"Injected rootUri for initialize: {ws}")
+                    return obj, True
+
+                injected = False
+                if isinstance(req, dict):
+                    req, injected = _inject(req)
+                elif isinstance(req, list):
+                    new_list = []
+                    for item in req:
+                        item2, did = _inject(item)
+                        injected = injected or did
+                        new_list.append(item2)
+                    req = new_list
+
+                if injected:
+                    msg = json.dumps(req).encode("utf-8")
+            except Exception:
+                pass
+            if mode_holder.get("mode") is None:
+                mode_holder["mode"] = mode
+            
+            header = f"Content-Length: {len(msg)}\r\n\r\n".encode("ascii")
+            sock.sendall(header + msg)
+    except Exception as e:
+        _log_error(f"Error forwarding stdin to socket: {e}")
+        sock.close()
+        sys.exit(1)
+
+def main():
+    # Log startup context for diagnostics
+    _log_info(
+        "Proxy startup: cwd=%s argv=%s env.DECKARD_WORKSPACE_ROOT=%s env.LOCAL_SEARCH_WORKSPACE_ROOT=%s"
+        % (
+            os.getcwd(),
+            sys.argv,
+            os.environ.get("DECKARD_WORKSPACE_ROOT"),
+            os.environ.get("LOCAL_SEARCH_WORKSPACE_ROOT"),
+        )
+    )
+    host = os.environ.get("DECKARD_DAEMON_HOST", DEFAULT_HOST)
+    port = int(os.environ.get("DECKARD_DAEMON_PORT", DEFAULT_PORT))
+
+    if not start_daemon_if_needed(host, port):
+        sys.exit(1)
+
+    try:
+        sock = socket.create_connection((host, port))
+    except Exception as e:
+        _log_error(f"Could not connect to daemon: {e}")
+        sys.exit(1)
+
+    # Start threads for bidirectional forwarding
+    mode_holder = {"mode": None}
+    t1 = threading.Thread(target=forward_socket_to_stdout, args=(sock, mode_holder), daemon=True)
+    t1.start()
+
+    forward_stdin_to_socket(sock, mode_holder)
+
+if __name__ == "__main__":
+    main()

mcp/registry.py

@@ -0,0 +1,218 @@
+#!/usr/bin/env python3
+"""
+Workspace Registry for Sari Daemon.
+
+Manages shared state (server instance) per workspace with refcount-based lifecycle.
+When all clients disconnect from a workspace (refcount=0), resources are cleaned up.
+"""
+import logging
+import threading
+from pathlib import Path
+from typing import Dict, Optional
+
+from app.http_server import serve_forever
+from app.registry import ServerRegistry
+
+logger = logging.getLogger("sari.registry")
+
+
+class SharedState:
+    """Holds the server instance and reference count for a workspace.
+    
+    Multiple clients connected to the same workspace share this state,
+    avoiding duplicate indexing and DB connections.
+    
+    Also manages the dedicated HTTP server for this workspace.
+    """
+    
+    def __init__(self, workspace_root: str):
+        from .server import LocalSearchMCPServer
+        self.workspace_root = workspace_root
+        self.server = LocalSearchMCPServer(workspace_root)
+        self.ref_count = 0
+        self.lock = threading.Lock()
+        
+        # HTTP Server State
+        self.httpd = None
+        self.http_port = 0
+        self.http_thread = None
+        
+        # Initialize Core Server (Loads Config)
+        try:
+            self.server._ensure_initialized()
+            cfg = self.server.cfg
+            
+            # Start HTTP Server
+            self.httpd, self.http_port = serve_forever(
+                host=cfg.server_host, 
+                port=cfg.server_port, 
+                db=self.server.db, 
+                indexer=self.server.indexer, 
+                version=self.server.SERVER_VERSION,
+                workspace_root=self.workspace_root
+            )
+            logger.info(f"Started HTTP Server for {workspace_root} on port {self.http_port}")
+            
+        except Exception as e:
+            logger.error(f"Failed to start server components for {workspace_root}: {e}")
+            # We don't raise here to allow partial functionality (MCP might work without HTTP?)
+            # But usually if init fails, MCP fails too.
+            pass
+            
+        logger.info(f"Created SharedState for workspace: {workspace_root}")
+
+    def acquire(self) -> int:
+        """Increment refcount (client connected)."""
+        with self.lock:
+            self.ref_count += 1
+            logger.debug(f"Acquired {self.workspace_root} (refcount={self.ref_count})")
+            return self.ref_count
+
+    def release(self) -> int:
+        """Decrement refcount (client disconnected)."""
+        with self.lock:
+            self.ref_count -= 1
+            logger.debug(f"Released {self.workspace_root} (refcount={self.ref_count})")
+            return self.ref_count
+
+    def shutdown(self) -> None:
+        """Stop indexer, close DB, and shutdown HTTP server."""
+        logger.info(f"Shutting down SharedState for {self.workspace_root}")
+        
+        # Unregister from Global Registry
+        try:
+            ServerRegistry().unregister(self.workspace_root)
+        except Exception as e:
+            logger.error(f"Failed to unregister workspace: {e}")
+
+        # Shutdown HTTP Server
+        if self.httpd:
+            logger.info("Shutting down HTTP server...")
+            self.httpd.shutdown()
+            self.httpd.server_close()
+            
+        self.server.shutdown()
+
+
+class Registry:
+    """Singleton registry to manage shared server instances.
+    
+    Provides refcount-based lifecycle management:
+    - get_or_create(): Get or create shared state, refcount++
+    - release(): refcount--, cleanup if refcount==0
+    """
+    _instance: Optional["Registry"] = None
+    _lock = threading.Lock()
+
+    def __init__(self):
+        self._workspaces: Dict[str, SharedState] = {}
+        self._registry_lock = threading.Lock()
+
+    @classmethod
+    def get_instance(cls) -> "Registry":
+        """Get the singleton Registry instance."""
+        with cls._lock:
+            if cls._instance is None:
+                cls._instance = Registry()
+                logger.info("Registry singleton created")
+        return cls._instance
+    
+    @classmethod
+    def reset_instance(cls) -> None:
+        """Reset singleton for testing purposes."""
+        with cls._lock:
+            if cls._instance is not None:
+                cls._instance.shutdown_all()
+            cls._instance = None
+
+    def get_or_create(self, workspace_root: str) -> SharedState:
+        """Get existing or create new SharedState for a workspace.
+        
+        Automatically increments refcount.
+        
+        Args:
+            workspace_root: Absolute path to workspace root
+            
+        Returns:
+            SharedState for the workspace
+        """
+        resolved_root = str(Path(workspace_root).resolve())
+        
+        with self._registry_lock:
+            if resolved_root not in self._workspaces:
+                self._workspaces[resolved_root] = SharedState(resolved_root)
+                logger.info(f"Registered new workspace: {resolved_root}")
+            
+            state = self._workspaces[resolved_root]
+            state.acquire()
+            return state
+
+    def release(self, workspace_root: str) -> None:
+        """Release SharedState for a workspace.
+        
+        Decrements refcount. If refcount reaches 0, cleans up resources.
+        
+        Args:
+            workspace_root: Absolute path to workspace root
+        """
+        resolved_root = str(Path(workspace_root).resolve())
+        
+        with self._registry_lock:
+            if resolved_root not in self._workspaces:
+                logger.warning(f"Attempted to release unknown workspace: {resolved_root}")
+                return
+            
+            state = self._workspaces[resolved_root]
+            remaining = state.release()
+            
+            if remaining <= 0:
+                state.shutdown()
+                del self._workspaces[resolved_root]
+                logger.info(f"Unregistered workspace: {resolved_root}")
+
+    def get(self, workspace_root: str) -> Optional[SharedState]:
+        """Get SharedState without modifying refcount.
+        
+        Args:
+            workspace_root: Absolute path to workspace root
+            
+        Returns:
+            SharedState if exists, None otherwise
+        """
+        resolved_root = str(Path(workspace_root).resolve())
+        with self._registry_lock:
+            return self._workspaces.get(resolved_root)
+
+    def list_workspaces(self) -> Dict[str, int]:
+        """List all active workspaces with their refcounts.
+        
+        Returns:
+            Dict mapping workspace_root to refcount
+        """
+        with self._registry_lock:
+            return {ws: state.ref_count for ws, state in self._workspaces.items()}
+
+    def active_count(self) -> int:
+        """Get number of active workspaces.
+        
+        Returns:
+            Number of workspaces with refcount > 0
+        """
+        with self._registry_lock:
+            return len(self._workspaces)
+
+    def shutdown_all(self) -> None:
+        """Shutdown all workspaces (for daemon stop)."""
+        with self._registry_lock:
+            workspace_count = len(self._workspaces)
+            for workspace_root, state in list(self._workspaces.items()):
+                logger.info(f"Shutting down workspace: {workspace_root}")
+                state.shutdown()
+            self._workspaces.clear()
+        logger.info(f"Registry shutdown complete ({workspace_count} workspaces)")
+
+
+# Convenience function for getting the registry
+def get_registry() -> Registry:
+    """Get the global Registry singleton."""
+    return Registry.get_instance()