sari 0.0.1__py3-none-any.whl

mcp/session.py

@@ -0,0 +1,234 @@
+import json
+import logging
+import asyncio
+import inspect
+from typing import Dict, Any, Optional
+from .registry import Registry, SharedState
+from app.workspace import WorkspaceManager
+
+logger = logging.getLogger(__name__)
+
+class Session:
+    """
+    Handles a single client connection.
+    Parses JSON-RPC, manages workspace binding via Registry.
+    """
+    def __init__(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
+        self.reader = reader
+        self.writer = writer
+        self.workspace_root: Optional[str] = None
+        self.shared_state: Optional[SharedState] = None
+        self.registry = Registry.get_instance()
+        self.running = True
+
+    async def handle_connection(self):
+        try:
+            while self.running:
+                # Read Headers
+                headers = {}
+                line_count = 0
+                while True:
+                    line = await self.reader.readline()
+                    if not line:
+                        self.running = False
+                        break
+                    
+                    line_str = line.decode("utf-8").strip()
+                    line_count += 1
+                    
+                    if not line_str:
+                        break
+                    
+                    # Protocol Check: First line must be Content-Length
+                    if line_count == 1:
+                        if line_str.startswith("{"):
+                            logger.error("Received JSONL instead of HTTP-style framed message")
+                            await self.send_error(None, -32700, "JSONL not supported. Use Content-Length framing.")
+                            self.running = False
+                            break
+                        
+                        if not line_str.lower().startswith("content-length:"):
+                            logger.error(f"First header must be Content-Length, got: {line_str!r}")
+                            await self.send_error(None, -32700, "Invalid protocol framing: Content-Length header required first")
+                            self.running = False
+                            break
+
+                    if ":" in line_str:
+                        k, v = line_str.split(":", 1)
+                        headers[k.strip().lower()] = v.strip()
+                    else:
+                        # Malformed header or missing Content-Length
+                        logger.error(f"Malformed header line: {line_str!r}")
+                        await self.send_error(None, -32700, "Invalid protocol framing")
+                        self.running = False
+                        break
+
+                if not self.running:
+                    break
+
+                try:
+                    content_length = int(headers.get("content-length", 0))
+                except (ValueError, TypeError):
+                    logger.error(f"Invalid Content-Length value: {headers.get('content-length')!r}")
+                    await self.send_error(None, -32700, "Invalid Content-Length value")
+                    self.running = False
+                    break
+
+                if content_length <= 0:
+                    logger.error("Received message without Content-Length (JSONL is not supported)")
+                    await self.send_error(None, -32700, "Content-Length header required (JSONL is not supported)")
+                    # Since protocol framing is broken, we must terminate
+                    self.running = False
+                    break
+
+                body = await self.reader.readexactly(content_length)
+                if not body:
+                    break
+
+                try:
+                    request_str = body.decode("utf-8")
+                    request = json.loads(request_str)
+                    await self.process_request(request)
+                except json.JSONDecodeError:
+                    logger.error(f"Invalid JSON received: {body[:100]!r}")
+                    # Try to extract ID manually for better correlation if possible
+                    msg_id = None
+                    try:
+                         # Simple regex for "id": 123 or "id": "abc"
+                         import re
+                         match = re.search(r'"id"\s*:\s*("(?:\\"|[^"])*"|\d+|null)', request_str)
+                         if match:
+                             msg_id = json.loads(match.group(1))
+                    except Exception:
+                        pass
+                    await self.send_error(msg_id, -32700, "Parse error")
+                except Exception as e:
+                    logger.error(f"Error processing request: {e}", exc_info=True)
+                    # We might have parsed the ID already if it's not a Parse error
+                    msg_id = None
+                    try:
+                        msg_id = json.loads(body.decode("utf-8")).get("id")
+                    except Exception:
+                        pass
+                    await self.send_error(msg_id, -32603, str(e))
+
+        except (asyncio.IncompleteReadError, ConnectionResetError):
+            logger.info("Connection closed by client")
+        finally:
+            self.cleanup()
+            try:
+                res = self.writer.close()
+                if inspect.isawaitable(res):
+                    await res
+            except Exception:
+                pass
+            try:
+                await self.writer.wait_closed()
+            except Exception:
+                pass
+
+    async def process_request(self, request: Dict[str, Any]):
+        method = request.get("method")
+        params = request.get("params", {})
+        msg_id = request.get("id")
+
+        if method == "initialize":
+            await self.handle_initialize(request)
+        elif method == "initialized":
+            # Just forward to server if bound
+            if self.shared_state:
+                loop = asyncio.get_event_loop()
+                await loop.run_in_executor(
+                    None,
+                    self.shared_state.server.handle_initialized,
+                    params
+                )
+        elif method == "shutdown":
+            # Respond to shutdown but keep connection open for exit
+            response = {"jsonrpc": "2.0", "id": msg_id, "result": None}
+            await self.send_json(response)
+        elif method == "exit":
+            self.running = False
+        else:
+            # Forward other requests to the bound server
+            if not self.shared_state:
+                await self.send_error(msg_id, -32002, "Server not initialized. Send 'initialize' first.")
+                return
+
+            # Execute in thread pool to not block async loop
+            # Since LocalSearchMCPServer is synchronous
+            loop = asyncio.get_event_loop()
+            response = await loop.run_in_executor(
+                None, 
+                self.shared_state.server.handle_request, 
+                request
+            )
+            
+            if response:
+                await self.send_json(response)
+
+    async def handle_initialize(self, request: Dict[str, Any]):
+        params = request.get("params", {})
+        msg_id = request.get("id")
+        
+        root_uri = params.get("rootUri") or params.get("rootPath")
+        if not root_uri:
+            # Fallback for clients that omit rootUri/rootPath
+            root_uri = WorkspaceManager.resolve_workspace_root()
+
+        # Handle file:// prefix
+        if root_uri.startswith("file://"):
+            workspace_root = root_uri[7:]
+        else:
+            workspace_root = root_uri
+
+        # If already bound to a different workspace, release it
+        if self.workspace_root and self.workspace_root != workspace_root:
+            self.registry.release(self.workspace_root)
+            self.shared_state = None
+
+        self.workspace_root = workspace_root
+        self.shared_state = self.registry.get_or_create(self.workspace_root)
+        
+        # Delegate specific initialize logic to the server instance
+        # We need to construct the result based on server's response
+        # LocalSearchMCPServer.handle_initialize returns the result dict directly
+        try:
+            result = self.shared_state.server.handle_initialize(params)
+            response = {
+                "jsonrpc": "2.0",
+                "id": msg_id,
+                "result": result
+            }
+            await self.send_json(response)
+        except Exception as e:
+            # Rollback: release the workspace if initialization failed
+            self.registry.release(self.workspace_root)
+            self.workspace_root = None
+            self.shared_state = None
+            await self.send_error(msg_id, -32000, str(e))
+
+    async def send_json(self, data: Dict[str, Any]):
+        body = json.dumps(data).encode("utf-8")
+        header = f"Content-Length: {len(body)}\r\n\r\n".encode("ascii")
+        res = self.writer.write(header + body)
+        if inspect.isawaitable(res):
+            await res
+        await self.writer.drain()
+
+    async def send_error(self, msg_id: Any, code: int, message: str):
+        response = {
+            "jsonrpc": "2.0",
+            "id": msg_id,
+            "error": {
+                "code": code,
+                "message": message
+            }
+        }
+        await self.send_json(response)
+
+    def cleanup(self):
+        if self.workspace_root:
+            self.registry.release(self.workspace_root)
+            self.workspace_root = None
+            self.shared_state = None

mcp/telemetry.py

@@ -0,0 +1,112 @@
+#!/usr/bin/env python3
+"""
+Telemetry and logging for Local Search MCP Server.
+"""
+import sys
+import queue
+import threading
+import time
+from pathlib import Path
+from datetime import datetime
+from typing import Optional
+try:
+    from app.indexer import _redact
+except ImportError:
+    # Fallback if imports fail (e.g. running script standalone without path)
+    # But usually app is in path.
+    def _redact(t): return t
+
+class TelemetryLogger:
+    """Handles logging and telemetry for MCP server."""
+    
+    def __init__(self, log_dir: Optional[Path] = None):
+        """
+        Initialize telemetry logger.
+        
+        Args:
+            log_dir: Directory for log files. If None, uses global log dir.
+        """
+        self.log_dir = Path(log_dir) if log_dir else None
+        self._queue: Optional[queue.Queue] = None
+        self._writer_thread: Optional[threading.Thread] = None
+        self._stop_event = threading.Event()
+        self._drop_count = 0
+        self._backlog_limit = 1000
+
+        if self.log_dir:
+            self._queue = queue.Queue()
+            self._writer_thread = threading.Thread(target=self._writer_loop, daemon=True)
+            self._writer_thread.start()
+    
+    def log_error(self, message: str) -> None:
+        """Log error message to stderr and file."""
+        print(f"[sari] ERROR: {message}", file=sys.stderr, flush=True)
+        self._enqueue(f"[ERROR] {message}")
+    
+    def log_info(self, message: str) -> None:
+        """Log info message to stderr and file."""
+        print(f"[sari] INFO: {message}", file=sys.stderr, flush=True)
+        self._enqueue(f"[INFO] {message}")
+    
+    def log_telemetry(self, message: str) -> None:
+        """
+        Log telemetry to file.
+        
+        Args:
+            message: Telemetry message to log
+        """
+        self._enqueue(message)
+
+    def _enqueue(self, message: str) -> None:
+        if not self._queue:
+            return
+        if self._queue.qsize() > self._backlog_limit:
+            self._drop_count += 1
+            return
+        self._queue.put(message)
+
+    def _writer_loop(self) -> None:
+        if not self.log_dir:
+            return
+        while not self._stop_event.is_set() or (self._queue and not self._queue.empty()):
+            try:
+                msg = self._queue.get(timeout=0.2) if self._queue else None
+            except queue.Empty:
+                continue
+            if msg is None:
+                continue
+            self._write_to_file(msg)
+            if self._queue:
+                self._queue.task_done()
+
+    def _write_to_file(self, message: str) -> None:
+        """Helper to write message with timestamp to log file."""
+        if not self.log_dir:
+            return
+        
+        # Redact secrets before writing to disk
+        message = _redact(message)
+        
+        try:
+            self.log_dir.mkdir(parents=True, exist_ok=True)
+            log_file = self.log_dir / "sari.log"
+            
+            timestamp = datetime.now().astimezone().isoformat()
+            with open(log_file, "a", encoding="utf-8") as f:
+                f.write(f"[{timestamp}] {message}\n")
+        except Exception as e:
+            print(f"[sari] ERROR: Failed to log to file: {e}", file=sys.stderr, flush=True)
+
+    def stop(self, timeout: float = 2.0) -> None:
+        if not self._queue or not self._writer_thread:
+            return
+        self._stop_event.set()
+        self._writer_thread.join(timeout=timeout)
+
+    def get_queue_depth(self) -> int:
+        if not self._queue:
+            return 0
+        return self._queue.qsize()
+
+    def get_drop_count(self) -> int:
+        return self._drop_count

mcp/test_cli.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+"""
+Unit tests for Sari CLI HTTP helpers.
+"""
+import io
+import json
+import os
+import sys
+import tempfile
+from contextlib import redirect_stdout
+from pathlib import Path
+from unittest.mock import patch
+
+SCRIPT_DIR = Path(__file__).parent
+sys.path.insert(0, str(SCRIPT_DIR))
+
+from cli import _get_http_host_port, cmd_search, cmd_status
+
+
+def _set_env(key: str, value: str):
+    if value is None:
+        os.environ.pop(key, None)
+    else:
+        os.environ[key] = value
+
+
+def test_get_http_host_port_prefers_server_json():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        data_dir = Path(tmpdir) / ".codex" / "tools" / "sari" / "data"
+        data_dir.mkdir(parents=True, exist_ok=True)
+        server_json = data_dir / "server.json"
+        server_json.write_text(json.dumps({"host": "127.0.0.1", "port": 47788}))
+
+        prev = os.environ.get("DECKARD_WORKSPACE_ROOT")
+        _set_env("DECKARD_WORKSPACE_ROOT", tmpdir)
+        try:
+            host, port = _get_http_host_port()
+            assert host == "127.0.0.1"
+            assert port == 47788
+        finally:
+            _set_env("DECKARD_WORKSPACE_ROOT", prev)
+
+
+def test_cmd_status_prints_json():
+    with patch("cli._request_http", return_value={"ok": True}) as mock_req:
+        buf = io.StringIO()
+        with redirect_stdout(buf):
+            rc = cmd_status(None)
+        assert rc == 0
+        mock_req.assert_called_once_with("/status", {})
+        out = buf.getvalue().strip()
+        assert out == json.dumps({"ok": True}, ensure_ascii=False, indent=2)
+
+
+def test_cmd_search_prints_json():
+    args = type("Args", (), {"query": "AuthService", "repo": "demo", "limit": 7})
+    with patch("cli._request_http", return_value={"ok": True, "q": "AuthService"}) as mock_req:
+        buf = io.StringIO()
+        with redirect_stdout(buf):
+            rc = cmd_search(args)
+        assert rc == 0
+        mock_req.assert_called_once_with("/search", {"q": "AuthService", "limit": 7, "repo": "demo"})
+        out = buf.getvalue().strip()
+        assert out == json.dumps({"ok": True, "q": "AuthService"}, ensure_ascii=False, indent=2)
+
+
+def run_tests():
+    tests = [
+        test_get_http_host_port_prefers_server_json,
+        test_cmd_status_prints_json,
+        test_cmd_search_prints_json,
+    ]
+    passed = 0
+    failed = 0
+    for test in tests:
+        try:
+            test()
+            print(f"✓ {test.__name__}")
+            passed += 1
+        except Exception:
+            failed += 1
+            raise
+    print(f"\n{passed} passed, {failed} failed")
+    return failed == 0
+
+
+if __name__ == "__main__":
+    success = run_tests()
+    sys.exit(0 if success else 1)

mcp/test_daemon.py

@@ -0,0 +1,124 @@
+import asyncio
+import json
+import socket
+import subprocess
+import sys
+import time
+import os
+import signal
+from pathlib import Path
+
+DAEMON_PORT = 47780
+DAEMON_HOST = "127.0.0.1"
+
+def wait_for_port(port, timeout=5):
+    start = time.time()
+    while time.time() - start < timeout:
+        try:
+            with socket.create_connection((DAEMON_HOST, port), timeout=0.1):
+                return True
+        except (ConnectionRefusedError, OSError):
+            time.sleep(0.1)
+    return False
+
+def send_rpc(sock, method, params=None, msg_id=1):
+    req = {
+        "jsonrpc": "2.0",
+        "method": method,
+        "id": msg_id,
+        "params": params or {}
+    }
+    body = json.dumps(req).encode('utf-8')
+    header = f"Content-Length: {len(body)}\r\n\r\n".encode('ascii')
+    sock.sendall(header + body)
+    
+    # Read response
+    f = sock.makefile('rb')
+    # Read headers
+    headers = {}
+    while True:
+        line = f.readline()
+        if not line or line == b"\r\n":
+            break
+        line_str = line.decode('utf-8').strip()
+        if ":" in line_str:
+            k, v = line_str.split(":", 1)
+            headers[k.strip().lower()] = v.strip()
+    
+    content_length = int(headers.get("content-length", 0))
+    if content_length > 0:
+        return json.loads(f.read(content_length).decode('utf-8'))
+    return None
+
+def test_daemon():
+    print("Starting daemon...")
+    env = os.environ.copy()
+    env["DECKARD_DAEMON_PORT"] = str(DAEMON_PORT)
+    
+    # Run as module from repo root
+    repo_root = Path(__file__).parent.parent
+    
+    proc = subprocess.Popen(
+        [sys.executable, "-m", "mcp.daemon"],
+        cwd=str(repo_root),
+        env=env,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True
+    )
+    
+    try:
+        if not wait_for_port(DAEMON_PORT):
+            print("Daemon failed to start")
+            print(proc.stderr.read())
+            sys.exit(1)
+            
+        print("Daemon started.")
+        
+        # Client 1: WS1
+        s1 = socket.create_connection((DAEMON_HOST, DAEMON_PORT))
+        print("Client 1 connected")
+        res1 = send_rpc(s1, "initialize", {"rootUri": "file:///tmp/test_ws1"})
+        print(f"Client 1 init result: {res1}")
+        assert "result" in res1
+        
+        # Client 2: WS1 (Should share indexer)
+        s2 = socket.create_connection((DAEMON_HOST, DAEMON_PORT))
+        print("Client 2 connected")
+        res2 = send_rpc(s2, "initialize", {"rootUri": "file:///tmp/test_ws1"})
+        print(f"Client 2 init result: {res2}")
+        assert "result" in res2
+        
+        # Client 3: WS2 (New indexer)
+        s3 = socket.create_connection((DAEMON_HOST, DAEMON_PORT))
+        print("Client 3 connected")
+        res3 = send_rpc(s3, "initialize", {"rootUri": "file:///tmp/test_ws2"})
+        print(f"Client 3 init result: {res3}")
+        assert "result" in res3
+        
+        # Verify functionality - e.g. tools/list
+        res_list = send_rpc(s1, "tools/list", {}, msg_id=2)
+        assert len(res_list["result"]["tools"]) > 0
+        print("Client 1 tools list OK")
+
+        # Clean up
+        s1.close()
+        s2.close()
+        s3.close()
+        print("Clients disconnected")
+        
+        time.sleep(1) # Allow daemon to log disconnects
+        
+    finally:
+        print("Stopping daemon...")
+        proc.terminate()
+        try:
+            outs, errs = proc.communicate(timeout=2)
+            print("Daemon stdout:", outs)
+            print("Daemon stderr:", errs)
+        except subprocess.TimeoutExpired:
+            proc.kill()
+            print("Daemon killed")
+
+if __name__ == "__main__":
+    test_daemon()