From ce8e294924de14b2cfb51e2e13daa5639fce730e Mon Sep 17 00:00:00 2001 From: Dmitry Teryaev Date: Sun, 12 Jul 2026 00:28:10 +0300 Subject: [PATCH 1/2] perf(mcp): lazy-load vector backend so warm `jrag search` skips torch import MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `mcp_v2` eagerly imported `search_lancedb` at module load, which imports `sentence_transformers`/`torch` at its own top level (~2.8s). The warm `jrag watch` client reconstructs `SearchOutput` from `mcp_v2` (`watch/client._reconstruct`) without ever calling the vector backend, so every warm CLI read paid that ~2.8s import — masking the daemon's win and leaving warm `jrag search` at ~2.6s instead of near-instant. Convert the eager import to lazy: `run_search`/`TABLES` start as sentinels and are populated by `_ensure_vector_backend()` on the first `search_v2` call. The guard preserves existing semantics: - tests that monkeypatch `mcp_v2.run_search` (non-None) are not overwritten; - graph-only installs (ImportError) leave `run_search=None` -> lexical fallback. Measured on tests/bank-chat-system (130 Java files): - `SearchOutput` import: 2.85s -> 0.16s - warm `jrag search`: 2.6s -> 0.15s steady (17x); cold path unchanged - results byte-identical warm vs cold Tests: run_search monkeypatchers (185), lightweight watch (101), heavy warm-path (6), all 6 golden IPC byte-identity (warm == cold) — all green. Co-Authored-By: Claude --- src/java_codebase_rag/mcp/mcp_v2.py | 41 ++++++++++++++++++++++------- 1 file changed, 32 insertions(+), 9 deletions(-) diff --git a/src/java_codebase_rag/mcp/mcp_v2.py b/src/java_codebase_rag/mcp/mcp_v2.py index 7310054..861d464 100644 --- a/src/java_codebase_rag/mcp/mcp_v2.py +++ b/src/java_codebase_rag/mcp/mcp_v2.py @@ -50,15 +50,37 @@ from java_codebase_rag.mcp.mcp_hints import MCP_HINTS_STRUCTURED_FIELD_DESCRIPTION # The vector stack (lancedb/torch, reached via search_lancedb) is optional — it is absent on -# graph-only installs (macOS Intel). Import eagerly when available so ``run_search``/``TABLES`` -# exist as module attributes (tests monkeypatch ``mcp_v2.run_search``; callers use ``TABLES``); -# fall back to sentinels on ImportError so importing this module never fails and ``search_v2`` -# can return a clean "vector search unavailable" envelope instead of crashing. -try: - from java_codebase_rag.search.search_lancedb import TABLES, run_search -except ImportError: # graph-only install: no torch/lancedb - TABLES = {} - run_search = None +# graph-only installs (macOS Intel). It is imported LAZILY on the first ``search_v2`` call via +# ``_ensure_vector_backend``, NOT at module load. Reason: ``search_lancedb`` imports +# ``sentence_transformers``/``torch`` at its own top level (~2.8s), and the warm ``jrag watch`` +# client reconstructs ``SearchOutput`` from this module (see ``watch/client._reconstruct``) +# without ever calling the vector backend — an eager import here would force every warm CLI +# read to pay that ~2.8s, masking the daemon's win. ``run_search``/``TABLES`` still exist as +# module attributes (sentinels until first use); tests monkeypatch ``mcp_v2.run_search`` to a +# non-None value, which ``_ensure_vector_backend`` treats as authoritative and will not overwrite. +TABLES: dict = {} +run_search: Any = None +_VECTOR_BACKEND_LOADED = False + + +def _ensure_vector_backend() -> None: + """Populate ``run_search``/``TABLES`` from ``search_lancedb`` on first use. + + A no-op once the backend has been loaded OR once ``run_search`` is non-None (tests + monkeypatch it before calling ``search_v2``). On graph-only installs the import fails and + ``run_search`` stays ``None`` — ``search_v2`` then takes the lexical fallback path. + """ + global run_search, TABLES, _VECTOR_BACKEND_LOADED + if run_search is not None or _VECTOR_BACKEND_LOADED: + return + _VECTOR_BACKEND_LOADED = True + try: + from java_codebase_rag.search.search_lancedb import TABLES as _TABLES, run_search as _run_search + except ImportError: # graph-only install: no torch/lancedb — leave run_search=None / TABLES={} + pass + else: + run_search = _run_search + TABLES = _TABLES __all__ = [ "search_v2", "find_v2", @@ -913,6 +935,7 @@ def search_v2( if nf and (err := _nodefilter_applicability_error("symbol", nf)): _log_fail_loud("applicability") return SearchOutput(success=False, message=err, advisories=[], limit=None, offset=None) + _ensure_vector_backend() advisories: list[str] = [] lexical_mode = run_search is None if lexical_mode: From b844cb6f8dacba106b4df31fe8e96a30d6ce8fb0 Mon Sep 17 00:00:00 2001 From: Dmitry Teryaev Date: Sun, 12 Jul 2026 00:58:23 +0300 Subject: [PATCH 2/2] fix(mcp): sentinel + lock for lazy vector backend (review follow-up) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address two code-review findings on the lazy-load change: 1. Critical — sentinel conflation. The guard `run_search is not None` could not distinguish "unloaded" (`None` sentinel) from a test that forces lexical mode via `monkeypatch.setattr(mcp_v2, "run_search", None)`. The loader ran anyway, overwrote the patch with the real backend, and `lexical_mode` flipped to False — regressing `test_search_v2_lexical_dispatch_end_to_end` (Lance "Table not found"). Fix: a distinct `_NOT_LOADED = object()` sentinel for the unloaded state, so `None` unambiguously means "lexical/disabled" and is authoritative in both monkeypatch directions (None → lexical; non-None → vector). 2. Important — thread race. The MCP server dispatches `search_v2` via `asyncio.to_thread` (server.py:659), so two concurrent first-batch searches could race the unguarded sentinel mutation: Thread B sees `_VECTOR_BACKEND_LOADED=True` but `run_search` still `None` and wrongly takes the lexical path / reads `TABLES={}`. Fix: double-checked locking with a `threading.Lock`, mirroring the existing `_st_lock` pattern. The sentinel also lets us drop `_VECTOR_BACKEND_LOADED` entirely (the sentinel itself tracks state), and a non-ImportError now propagates and leaves `run_search` as `_NOT_LOADED` so the next call retries instead of poisoning the process. Verified: previously-failing lexical test passes; 206 lexical + run_search- monkeypatcher tests pass; 12 heavy warm + golden IPC byte-identity tests (warm == cold) pass; SearchOutput import stays ~0.14s. Co-Authored-By: Claude --- src/java_codebase_rag/mcp/mcp_v2.py | 48 ++++++++++++++++++----------- 1 file changed, 30 insertions(+), 18 deletions(-) diff --git a/src/java_codebase_rag/mcp/mcp_v2.py b/src/java_codebase_rag/mcp/mcp_v2.py index 861d464..07df593 100644 --- a/src/java_codebase_rag/mcp/mcp_v2.py +++ b/src/java_codebase_rag/mcp/mcp_v2.py @@ -55,32 +55,44 @@ # ``sentence_transformers``/``torch`` at its own top level (~2.8s), and the warm ``jrag watch`` # client reconstructs ``SearchOutput`` from this module (see ``watch/client._reconstruct``) # without ever calling the vector backend — an eager import here would force every warm CLI -# read to pay that ~2.8s, masking the daemon's win. ``run_search``/``TABLES`` still exist as -# module attributes (sentinels until first use); tests monkeypatch ``mcp_v2.run_search`` to a -# non-None value, which ``_ensure_vector_backend`` treats as authoritative and will not overwrite. +# read to pay that ~2.8s, masking the daemon's win. +# +# ``run_search`` uses a distinct ``_NOT_LOADED`` sentinel (not ``None``) for the unloaded state, +# so ``None`` unambiguously means "lexical/disabled". This lets tests force the lexical path by +# monkeypatching ``run_search=None`` (``_ensure_vector_backend`` treats that as authoritative and +# will not overwrite it), while a non-None monkeypatch drives the vector path. ``TABLES`` is +# ``{}`` until first successful load (only read in the vector path, where it is guaranteed set). +_NOT_LOADED = object() +run_search: Any = _NOT_LOADED TABLES: dict = {} -run_search: Any = None -_VECTOR_BACKEND_LOADED = False +_vector_backend_lock = threading.Lock() def _ensure_vector_backend() -> None: """Populate ``run_search``/``TABLES`` from ``search_lancedb`` on first use. - A no-op once the backend has been loaded OR once ``run_search`` is non-None (tests - monkeypatch it before calling ``search_v2``). On graph-only installs the import fails and - ``run_search`` stays ``None`` — ``search_v2`` then takes the lexical fallback path. + No-op once ``run_search`` has left the ``_NOT_LOADED`` state — real callable (vector path), + ``None`` (graph-only install or a test-forced lexical mode), or a test monkeypatch. + Thread-safe (double-checked locking): the MCP server dispatches ``search_v2`` via + ``asyncio.to_thread``, so two concurrent first-batch searches could otherwise race the + sentinel mutation. On ImportError (graph-only) ``run_search`` becomes ``None`` and + ``search_v2`` takes the lexical fallback path. A non-ImportError propagates (caught by + ``search_v2``'s outer handler) and leaves ``run_search`` as ``_NOT_LOADED`` so the next + call retries rather than poisoning the process. """ - global run_search, TABLES, _VECTOR_BACKEND_LOADED - if run_search is not None or _VECTOR_BACKEND_LOADED: + global run_search, TABLES + if run_search is not _NOT_LOADED: return - _VECTOR_BACKEND_LOADED = True - try: - from java_codebase_rag.search.search_lancedb import TABLES as _TABLES, run_search as _run_search - except ImportError: # graph-only install: no torch/lancedb — leave run_search=None / TABLES={} - pass - else: - run_search = _run_search - TABLES = _TABLES + with _vector_backend_lock: + if run_search is not _NOT_LOADED: + return + try: + from java_codebase_rag.search.search_lancedb import TABLES as _TABLES, run_search as _run_search + except ImportError: # graph-only install: no torch/lancedb — lexical fallback + run_search = None + else: + run_search = _run_search + TABLES = _TABLES __all__ = [ "search_v2", "find_v2",