fix(mcp): answer initialize immediately — run startup preflight in a background thread

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.
This commit is contained in:
KeilerHirsch 2026-07-10 14:08:22 +02:00
parent d7819d4b85
commit e360a3a040
2 changed files with 143 additions and 7 deletions

View File

@ -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

View File

@ -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)