fix(mcp): isolate stdio wire protocol, add cache-first embedder fast path and background warmup (#194)

* fix(mcp): isolate stdio wire protocol, add cache-first embedder fast path and background warmup

- Redirect sys.stdout to sys.stderr in stdio MCP mode so library prints/warnings never corrupt the JSON-RPC wire
- Load cached SentenceTransformer models with local_files_only=True first to eliminate network roundtrips and timeouts
- Make service() singleton initialization thread-safe and warm up in a background daemon thread
- Add embedder health check to 'engraphis-init --check' and provide 'engraphis-init --prefetch' command

* test(backends): mock sentence_transformers via sys.modules for offline test suites

Author: Coding-Dev-Tools

engraphis/backends/embedder_st.py

@@ -319,6 +319,18 @@ def get_embedder(
             }
             if local_files_only:
                 factory_kwargs["local_files_only"] = True
+            else:
+                # Optimize cold start: if the model is already in local Hugging Face cache,
+                # loading with local_files_only=True avoids network roundtrips to huggingface.co.
+                # If cached, it returns immediately; if not, it seamlessly falls through to download.
+                try:
+                    cached_kwargs = dict(factory_kwargs)
+                    cached_kwargs["local_files_only"] = True
+                    emb = SentenceTransformerEmbedder(resolved_model_name, **cached_kwargs)
+                    LAST_EMBEDDER_ERROR = ""
+                    return emb
+                except Exception:
+                    pass
             emb = SentenceTransformerEmbedder(resolved_model_name, **factory_kwargs)
             LAST_EMBEDDER_ERROR = ""
             return emb

engraphis/mcp_server.py

@@ -27,10 +27,13 @@
 import hmac
 import json
 import logging
+import math
+import os
 import re
-import time
 import secrets
-import math
+import sys
+import threading
+import time
 
 from collections import OrderedDict
 from dataclasses import dataclass
@@ -103,22 +106,27 @@ def set_service(svc: MemoryService) -> None:
     _service = svc
 
 
+_service_lock = threading.Lock()
+
+
 def service() -> MemoryService:
     """Lazily build the service so server startup is instant (model loads on first use)."""
     global _service
     if _service is None:
-        _service = MemoryService.create(
-            settings.db_path,
-            embed_model=settings.embed_model or None,
-            embed_revision=getattr(settings, "embed_revision", "") or None,
-            require_immutable_models=bool(getattr(settings, "require_immutable_models", False)),
-            require_exact_backends=bool(getattr(settings, "require_exact_backends", False)),
-            embed_dim=settings.embed_dim if settings.embed_dim is not None else 384,
-            vector_backend=settings.vector_backend,
-            rerank_model=getattr(settings, "rerank_model", "") or None,
-            rerank_revision=getattr(settings, "rerank_revision", "") or None,
-            extractor=settings.extractor,
-        )
+        with _service_lock:
+            if _service is None:
+                _service = MemoryService.create(
+                    settings.db_path,
+                    embed_model=settings.embed_model or None,
+                    embed_revision=getattr(settings, "embed_revision", "") or None,
+                    require_immutable_models=bool(getattr(settings, "require_immutable_models", False)),
+                    require_exact_backends=bool(getattr(settings, "require_exact_backends", False)),
+                    embed_dim=settings.embed_dim if settings.embed_dim is not None else 384,
+                    vector_backend=settings.vector_backend,
+                    rerank_model=getattr(settings, "rerank_model", "") or None,
+                    rerank_revision=getattr(settings, "rerank_revision", "") or None,
+                    extractor=settings.extractor,
+                )
     return _service
 
 
@@ -3033,10 +3041,59 @@ def _eager_exact_backend_check() -> None:
         service()
 
 
+def _start_background_warmup() -> None:
+    """Warm up the memory service in a background daemon thread.
+
+    Allows the initial MCP handshake (initialize, tools/list) to respond in
+    milliseconds while warming SQLite and the embedding model before the agent's
+    first tool invocation. Can be disabled via ENGRAPHIS_MCP_WARMUP=0.
+    """
+    warmup_env = os.environ.get("ENGRAPHIS_MCP_WARMUP", "1").strip().lower()
+    if warmup_env in {"0", "false", "no", "off"}:
+        return
+    thread = threading.Thread(target=service, name="engraphis-warmup", daemon=True)
+    thread.start()
+
+
+async def _safe_run_stdio_async(server: FastMCP) -> None:
+    """Run stdio transport with pure wire protocol isolation.
+
+    In stdio MCP, standard output is exclusively the JSON-RPC wire. Redirect
+    Python's global `sys.stdout` to `sys.stderr` so that any prints, warnings,
+    or dependency output (PyTorch, transformers, tqdm, pydantic) flow safely to
+    stderr without corrupting JSON-RPC messages on the client pipe.
+    """
+    import anyio
+    from io import TextIOWrapper
+    from mcp.server.stdio import stdio_server
+
+    real_stdout_buffer = getattr(sys.stdout, "buffer", None)
+    real_stdin_buffer = getattr(sys.stdin, "buffer", None)
+    if real_stdout_buffer is not None and real_stdin_buffer is not None:
+        sys.stdout = sys.stderr
+        wrapped_stdout = anyio.wrap_file(TextIOWrapper(real_stdout_buffer, encoding="utf-8"))
+        wrapped_stdin = anyio.wrap_file(TextIOWrapper(real_stdin_buffer, encoding="utf-8", errors="replace"))
+        async with stdio_server(stdin=wrapped_stdin, stdout=wrapped_stdout) as (read_stream, write_stream):
+            await server._mcp_server.run(
+                read_stream,
+                write_stream,
+                server._mcp_server.create_initialization_options(),
+            )
+    else:
+        async with stdio_server() as (read_stream, write_stream):
+            await server._mcp_server.run(
+                read_stream,
+                write_stream,
+                server._mcp_server.create_initialization_options(),
+            )
+
+
 def main() -> None:
     """Console entry point (``engraphis-mcp``). Runs Smart MCP over stdio."""
     _eager_exact_backend_check()
-    mcp.run()
+    _start_background_warmup()
+    import anyio
+    anyio.run(lambda: _safe_run_stdio_async(mcp))
 
 
 if __name__ == "__main__":

scripts/init.py

@@ -12,6 +12,7 @@
     engraphis-init --encrypted     # require SQLCipher and provision a private DB key file
     engraphis-init --force         # overwrite the trusted config file
     engraphis-init --check         # doctor: verify install, extras, DB writability
+    engraphis-init --prefetch      # pre-cache embedding model weights for instant MCP startup
 
 Non-interactive by design (no prompts): safe in scripts, CI, and agent shells.
 """
@@ -104,10 +105,38 @@ def cmd_check() -> int:
     except Exception:
         _miss("Engraphis Cloud", "saved session unavailable; reconnect if needed")
 
+    try:
+        from engraphis.backends.embedder_st import get_embedder
+        emb = get_embedder(settings.embed_model or None, dim=settings.embed_dim or 384)
+        emb.embed(["engraphis doctor check"])
+        _ok("embedder functional", f"{type(emb).__name__} ({getattr(emb, 'dim', 384)}d)")
+    except Exception as exc:
+        _fail("embedder functional", f"{type(exc).__name__}: {exc}")
+        failures += 1
+
     print("all good" if failures == 0 else f"{failures} problem(s) found")
     return 0 if failures == 0 else 1
 
 
+def cmd_prefetch() -> int:
+    """Download and warm up the configured embedding model ahead of time."""
+    from engraphis.config import settings
+    model_name = (settings.embed_model or "").strip()
+    if not model_name:
+        print("  [--] No remote embedding model configured; deterministic offline embedder is active.")
+        return 0
+    print(f"engraphis prefetch - model '{model_name}'")
+    try:
+        from engraphis.backends.embedder_st import get_embedder
+        emb = get_embedder(model_name, dim=settings.embed_dim or 384, require_exact=True)
+        emb.embed(["engraphis prefetch warmup"])
+        _ok("model prefetch", f"{type(emb).__name__} ({getattr(emb, 'dim', 384)}d) ready")
+        return 0
+    except Exception as exc:
+        _fail("model prefetch", f"{type(exc).__name__}: {exc}")
+        return 1
+
+
 def _env_content(db_path: Path, token: str, key_path: Optional[Path] = None) -> str:
     lines = [
         "# Engraphis - generated by engraphis-init. Full reference: .env.example",
@@ -243,10 +272,14 @@ def main(argv=None) -> int:
     )
     ap.add_argument("--check", action="store_true",
                     help="doctor mode: verify the installation without writing config")
+    ap.add_argument("--prefetch", action="store_true",
+                    help="pre-cache the configured embedding model for instant MCP startup")
     args = ap.parse_args(argv)
 
     if args.check:
         return cmd_check()
+    if args.prefetch:
+        return cmd_prefetch()
 
     db_path = Path(args.db).expanduser().resolve()
     try:

tests/test_backends_factories.py

@@ -115,6 +115,52 @@ def __init__(
     }
 
 
+def test_get_embedder_prefers_local_cache(monkeypatch):
+    calls = []
+
+    class FakeST:
+        def __init__(self, model_name, **kwargs):
+            calls.append((model_name, kwargs))
+
+        def get_embedding_dimension(self):
+            return 384
+
+    monkeypatch.setitem(
+        sys.modules,
+        "sentence_transformers",
+        SimpleNamespace(SentenceTransformer=FakeST),
+    )
+    emb = get_embedder("sentence-transformers/all-MiniLM-L6-v2", 384)
+    assert emb.dim == 384
+    assert len(calls) == 1
+    assert calls[0][1].get("local_files_only") is True
+
+
+def test_get_embedder_falls_back_when_not_cached(monkeypatch):
+    calls = []
+
+    class FakeST:
+        def __init__(self, model_name, **kwargs):
+            calls.append((model_name, kwargs))
+            if kwargs.get("local_files_only") is True:
+                raise OSError("not in cache")
+
+        def get_embedding_dimension(self):
+            return 384
+
+    monkeypatch.setitem(
+        sys.modules,
+        "sentence_transformers",
+        SimpleNamespace(SentenceTransformer=FakeST),
+    )
+    emb = get_embedder("sentence-transformers/all-MiniLM-L6-v2", 384)
+    assert emb.dim == 384
+    assert len(calls) == 2
+    assert calls[0][1].get("local_files_only") is True
+    assert not calls[1][1].get("local_files_only")
+
+
+
 @pytest.mark.parametrize("revision", [None, "main", "A" * 40, "a" * 39])
 def test_embedder_strict_mode_rejects_mutable_remote_revision_before_load(monkeypatch, revision):
     import engraphis.backends.embedder_st as embedder_st

tests/test_init.py

@@ -229,3 +229,22 @@ def test_doctor_reports_connected_cloud_install(tmp_path, monkeypatch, capsys):
     assert main(["--check"]) == 0
     out = capsys.readouterr().out
     assert "Engraphis Cloud - installation connected" in out
+
+
+def test_doctor_reports_functional_embedder(tmp_path, monkeypatch, capsys):
+    _fresh_settings(monkeypatch, tmp_path)
+    assert main(["--check"]) == 0
+    out = capsys.readouterr().out
+    assert "embedder functional" in out
+
+
+def test_prefetch_command_reports_ready_or_offline(tmp_path, monkeypatch, capsys):
+    _fresh_settings(monkeypatch, tmp_path)
+    # Test prefetch with offline deterministic model
+    monkeypatch.setenv("ENGRAPHIS_EMBED_MODEL", "")
+    import engraphis.config as cfg
+    monkeypatch.setattr(cfg, "settings", cfg.Settings())
+    assert main(["--prefetch"]) == 0
+    out = capsys.readouterr().out
+    assert "deterministic offline embedder is active" in out
+

tests/test_mcp_server.py

@@ -1267,3 +1267,42 @@ def test_classic_remember_persists_subject_key_and_claim_kind_to_chain(monkeypat
     assert head["id"] == payload["id"]
     assert head["subject_key"] == "deploy.timeout"
     assert head["claim_kind"] == "configured_value"
+
+
+def test_service_singleton_is_thread_safe():
+    import engraphis.mcp_server as srv
+    from concurrent.futures import ThreadPoolExecutor
+    # Concurrently call srv.service() across 5 threads
+    with ThreadPoolExecutor(max_workers=5) as executor:
+        instances = list(executor.map(lambda _: srv.service(), range(5)))
+    assert len(instances) == 5
+    for inst in instances[1:]:
+        assert inst is instances[0]
+
+
+def test_background_warmup_honors_env(monkeypatch):
+    import engraphis.mcp_server as server
+    called = []
+    monkeypatch.setattr(server, "service", lambda: called.append(True))
+
+    # Disabled by env
+    monkeypatch.setenv("ENGRAPHIS_MCP_WARMUP", "0")
+    server._start_background_warmup()
+    assert not called
+
+    # Enabled by default or env
+    started_threads = []
+    real_thread = server.threading.Thread
+
+    def fake_thread(*args, **kwargs):
+        t = real_thread(*args, **kwargs)
+        started_threads.append(t)
+        return t
+
+    monkeypatch.setattr(server.threading, "Thread", fake_thread)
+    monkeypatch.setenv("ENGRAPHIS_MCP_WARMUP", "1")
+    server._start_background_warmup()
+    assert len(started_threads) == 1
+    assert started_threads[0].daemon is True
+    assert started_threads[0].name == "engraphis-warmup"
+