From 4773e1f8cf72a1e938b3fc65c1dc65520e7d57e2 Mon Sep 17 00:00:00 2001 From: Coding-Dev-Tools Date: Thu, 3 Sep 2026 19:17:45 -0400 Subject: [PATCH 1/2] 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 --- engraphis/backends/embedder_st.py | 12 +++++ engraphis/mcp_server.py | 87 +++++++++++++++++++++++++------ scripts/init.py | 33 ++++++++++++ tests/test_backends_factories.py | 40 ++++++++++++++ tests/test_init.py | 19 +++++++ tests/test_mcp_server.py | 39 ++++++++++++++ 6 files changed, 215 insertions(+), 15 deletions(-) diff --git a/engraphis/backends/embedder_st.py b/engraphis/backends/embedder_st.py index c6445d99..bd546d0e 100644 --- a/engraphis/backends/embedder_st.py +++ b/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 diff --git a/engraphis/mcp_server.py b/engraphis/mcp_server.py index bb9177e8..bf4ee163 100644 --- a/engraphis/mcp_server.py +++ b/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__": diff --git a/scripts/init.py b/scripts/init.py index ac144cee..5610356b 100644 --- a/scripts/init.py +++ b/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: diff --git a/tests/test_backends_factories.py b/tests/test_backends_factories.py index 84061e69..210472ce 100644 --- a/tests/test_backends_factories.py +++ b/tests/test_backends_factories.py @@ -115,6 +115,46 @@ def __init__( } +def test_get_embedder_prefers_local_cache(monkeypatch): + import sentence_transformers + calls = [] + + class FakeST: + def __init__(self, model_name, **kwargs): + calls.append((model_name, kwargs)) + + def get_embedding_dimension(self): + return 384 + + monkeypatch.setattr(sentence_transformers, "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): + import sentence_transformers + 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.setattr(sentence_transformers, "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 diff --git a/tests/test_init.py b/tests/test_init.py index 204b3da6..7388eb83 100644 --- a/tests/test_init.py +++ b/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 + diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 98836b61..34b8099b 100644 --- a/tests/test_mcp_server.py +++ b/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" + From 42735bf32ab3ac610f53bd4c0bb578d3cf48ec56 Mon Sep 17 00:00:00 2001 From: Coding-Dev-Tools Date: Thu, 3 Sep 2026 19:29:49 -0400 Subject: [PATCH 2/2] test(backends): mock sentence_transformers via sys.modules for offline test suites --- tests/test_backends_factories.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/tests/test_backends_factories.py b/tests/test_backends_factories.py index 210472ce..27ab3e6e 100644 --- a/tests/test_backends_factories.py +++ b/tests/test_backends_factories.py @@ -116,7 +116,6 @@ def __init__( def test_get_embedder_prefers_local_cache(monkeypatch): - import sentence_transformers calls = [] class FakeST: @@ -126,7 +125,11 @@ def __init__(self, model_name, **kwargs): def get_embedding_dimension(self): return 384 - monkeypatch.setattr(sentence_transformers, "SentenceTransformer", FakeST) + 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 @@ -134,7 +137,6 @@ def get_embedding_dimension(self): def test_get_embedder_falls_back_when_not_cached(monkeypatch): - import sentence_transformers calls = [] class FakeST: @@ -146,7 +148,11 @@ def __init__(self, model_name, **kwargs): def get_embedding_dimension(self): return 384 - monkeypatch.setattr(sentence_transformers, "SentenceTransformer", FakeST) + 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