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:
parent
d7819d4b85
commit
e360a3a040
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Reference in New Issue