From e360a3a0404aae387eaac607c299658ab41e993e Mon Sep 17 00:00:00 2001 From: KeilerHirsch Date: Fri, 10 Jul 2026 14:08:22 +0200 Subject: [PATCH] =?UTF-8?q?fix(mcp):=20answer=20initialize=20immediately?= =?UTF-8?q?=20=E2=80=94=20run=20startup=20preflight=20in=20a=20background?= =?UTF-8?q?=20thread?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The stdio loop ran _refresh_sqlite_integrity_status() and _refresh_vector_disabled_flag() before reading the first request. PRAGMA quick_check reads every page of chroma.sqlite3, so on multi-GB palaces the probe alone (measured: 20.3s on a 1.72 GB / 326k-drawer palace, 40-46s under disk/lock contention) starves the MCP client's 60s connect timeout — even though the initialize response itself never touches the database. The HTTP transport already starts without the synchronous probe. Move both probes to a daemon thread (mcp-startup-preflight). The #1222 intent is preserved: the probe still starts at startup and logs its warning as soon as it finishes. Consumers that need the verdict (_ensure_sqlite_integrity_status via the tool-call integrity gate) serialize on a new _sqlite_integrity_refresh_lock with double-checked locking, so a tool call arriving mid-probe waits for the in-flight verdict instead of running a second O(database size) quick_check — and never proceeds unverified. Measured on the 1.72 GB palace with the >512 MB startup gate disabled (MEMPALACE_STARTUP_INTEGRITY_MAX_MB=0, full quick_check in flight): initialize 1.4s (was 20-46s); first tool call after probe completion 3.4s with sqlite_integrity checked=true ok=true. Complements c54531a: the oversized-palace skip still applies to the background probe, but the handshake no longer depends on it. --- mempalace/mcp_server.py | 54 +++++++++++++++++++--- tests/test_mcp_server.py | 96 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 143 insertions(+), 7 deletions(-) diff --git a/mempalace/mcp_server.py b/mempalace/mcp_server.py index 73654a8..55a9fad 100644 --- a/mempalace/mcp_server.py +++ b/mempalace/mcp_server.py @@ -342,6 +342,11 @@ _last_request_time: float = time.monotonic() _sqlite_integrity_checked = False _sqlite_integrity_errors: list[str] = [] _sqlite_integrity_check_error = "" +# Serializes quick_check runs between the async startup preflight thread and +# lazy consumers on the protocol thread (double-checked in +# _ensure_sqlite_integrity_status) so the O(database size) probe never runs +# twice concurrently. +_sqlite_integrity_refresh_lock = threading.Lock() _SQLITE_INTEGRITY_ERROR_CODE = -32002 _SQLITE_INTEGRITY_ALLOWED_TOOLS = frozenset( { @@ -526,6 +531,12 @@ def _refresh_sqlite_integrity_status() -> None: SQLite-layer corruption (#1818). """ + with _sqlite_integrity_refresh_lock: + _refresh_sqlite_integrity_status_locked() + + +def _refresh_sqlite_integrity_status_locked() -> None: + # Probe body; callers must hold _sqlite_integrity_refresh_lock. global _sqlite_integrity_checked global _sqlite_integrity_errors global _sqlite_integrity_check_error @@ -583,8 +594,14 @@ def _refresh_sqlite_integrity_status() -> None: def _ensure_sqlite_integrity_status() -> None: - if not _sqlite_integrity_checked: - _refresh_sqlite_integrity_status() + if _sqlite_integrity_checked: + return + with _sqlite_integrity_refresh_lock: + # Double-checked: the startup preflight thread may have finished the + # probe while this caller waited on the lock — don't pay the + # O(database size) quick_check twice. + if not _sqlite_integrity_checked: + _refresh_sqlite_integrity_status_locked() def _sqlite_integrity_payload() -> dict: @@ -5357,6 +5374,21 @@ def _serve_http(host: str, port: int) -> None: logger.info("MemPalace MCP HTTP server shutting down") +def _startup_preflight() -> None: + """Startup SQLite integrity + HNSW capacity probes, off the protocol thread. + + Runs the same checks the stdio loop used to run synchronously before + reading the first request. Failures must never take down the server: the + lazy consumers (_ensure_sqlite_integrity_status, _get_client) re-run or + re-check on demand, so an exception here only loses the early warning. + """ + try: + _ensure_sqlite_integrity_status() + _refresh_vector_disabled_flag() + except Exception: + logger.exception("startup preflight failed") + + def _run_stdio_loop() -> None: _restore_stdout() @@ -5373,11 +5405,19 @@ def _run_stdio_loop() -> None: logger.info("MemPalace MCP Server starting...") - # Pre-flight: probe HNSW capacity before any tool call so the warning - # is visible at startup rather than on first use (#1222). Pure - # filesystem read; never opens a chromadb client. - _refresh_sqlite_integrity_status() - _refresh_vector_disabled_flag() + # Pre-flight in a background thread: PRAGMA quick_check reads every page + # of chroma.sqlite3 (20s+ on multi-GB palaces) and running it before the + # protocol loop starves the client's initialize timeout, even though the + # handshake itself never touches the database. The #1222 intent (warnings + # visible at startup rather than on first use) is preserved — the probe + # starts now and logs as soon as it finishes; tool calls that need the + # verdict serialize on _sqlite_integrity_refresh_lock via + # _ensure_sqlite_integrity_status instead of re-running the probe. + threading.Thread( + target=_startup_preflight, + name="mcp-startup-preflight", + daemon=True, + ).start() # Opt-in: pre-load the embedder so the first chromadb-write tool call # does not pay the ONNX/CoreML cold-load tax under the MCP client diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index dce308c..4697d0e 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -5341,3 +5341,99 @@ class TestListDrawersDateFilters: since = datetime(2026, 1, 2) assert _filed_at_in_window("2026-01-02T08:00:00Z", since, None) is True + + +# ── MCP stdio startup: async preflight ─────────────────────────────────── + + +def test_startup_preflight_does_not_block_initialize(monkeypatch): + """The startup integrity probe is O(database size) (PRAGMA quick_check + reads every page of chroma.sqlite3 — 20s+ on multi-GB palaces) and used + to run before the protocol loop, starving the client's initialize + timeout. It now runs on the mcp-startup-preflight thread; the handshake + must answer immediately while the probe is still in flight.""" + import threading + import time + + from mempalace import mcp_server + + probe_started = threading.Event() + release_probe = threading.Event() + + def slow_probe(): + probe_started.set() + release_probe.wait(10) + mcp_server._sqlite_integrity_checked = True + + monkeypatch.setattr(mcp_server, "_refresh_sqlite_integrity_status_locked", slow_probe) + monkeypatch.setattr(mcp_server, "_sqlite_integrity_checked", False) + monkeypatch.setattr(mcp_server, "_refresh_vector_disabled_flag", lambda: None) + + preflight = threading.Thread(target=mcp_server._startup_preflight, daemon=True) + preflight.start() + try: + assert probe_started.wait(5), "preflight thread never started the probe" + + started = time.monotonic() + response = mcp_server.handle_request( + { + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": {"protocolVersion": "2024-11-05"}, + } + ) + elapsed = time.monotonic() - started + + assert response["result"]["serverInfo"]["name"] == "mempalace" + assert elapsed < 1.0, f"initialize blocked {elapsed:.2f}s behind the startup probe" + finally: + release_probe.set() + preflight.join(5) + + +def test_ensure_sqlite_integrity_status_joins_inflight_probe(monkeypatch): + """A lazy consumer (tool-call integrity gate) arriving while the startup + preflight probe is still running must wait for that probe's verdict on + _sqlite_integrity_refresh_lock — not run a second O(database size) + quick_check concurrently, and not proceed without a verdict.""" + import threading + + from mempalace import mcp_server + + probe_calls = [] + probe_started = threading.Event() + release_probe = threading.Event() + + def slow_probe(): + probe_calls.append(1) + probe_started.set() + release_probe.wait(10) + mcp_server._sqlite_integrity_checked = True + + monkeypatch.setattr(mcp_server, "_refresh_sqlite_integrity_status_locked", slow_probe) + monkeypatch.setattr(mcp_server, "_sqlite_integrity_checked", False) + + background = threading.Thread( + target=mcp_server._refresh_sqlite_integrity_status, daemon=True + ) + background.start() + assert probe_started.wait(5), "background probe never started" + + consumer_done = threading.Event() + + def consumer(): + mcp_server._ensure_sqlite_integrity_status() + consumer_done.set() + + consumer_thread = threading.Thread(target=consumer, daemon=True) + consumer_thread.start() + try: + assert not consumer_done.wait(0.3), "consumer bypassed the in-flight probe" + release_probe.set() + assert consumer_done.wait(5), "consumer never unblocked after the probe finished" + assert probe_calls == [1], "quick_check probe ran more than once" + finally: + release_probe.set() + background.join(5) + consumer_thread.join(5)