Skip to content

Commit 84b9ad9

Browse files
committed
Make MCP indexing observable and tools available immediately
Addresses dev-team findings about CodeRAG-as-MCP-tool feeling broken at broad-root (e.g. /home) scale: "enabled" but not "usable", opaque status, and tools that aren't reachable in time inside Hermes/Claude Code/Codex. - Serve the MCP protocol immediately. run_mcp() moved warm-up (which downloads/loads the embedding model on first run) and the initial index into one background bootstrap thread, so mcp.run() — and thus tools/list — is reached at once. The toolset is registered synchronously in build_mcp(), so clients see all five tools right away instead of timing out behind the model download. This is the CodeRAG-side cause of the "tools registered but unusable" race. - Live, pollable index progress. New thread-safe IndexProgress (types.py) threaded from Indexer.index(live=...) up through CodeRAG.index() to the index_status tool, which now returns a `progress` object: state (idle/scanning/indexing/optimizing/ready/failed), files_discovered, files_to_index, files_indexed, chunks, current_path, elapsed, last_error. files_discovered ticks up during the long pre-embed scan, so a big index reads as "scanning" instead of a stuck 0. reindex drives the same object. - Earlier partial results. During a live (MCP background) index, buffered rows are committed every ~5s so dense search returns hits before the 8192-chunk flush boundary. Gated on live is not None, so the CLI, watcher and tests keep today's single-flush-at-end batching (defaults unchanged). - MCP best-practice polish. ToolAnnotations(readOnlyHint=True) on search_code/search_files/get_file/index_status; reindex marked non-read-only. - Concurrency safety. Serving before warm-up means a query can arrive mid bootstrap; guard CodeRAG's lazy provider/store/searcher/indexer construction with a reentrant build lock so two threads can't build a second conflicting LanceStore. Tests: live-progress + failure-path tests in test_indexer.py; live progress, annotations, and reindex-progress assertions in test_mcp.py. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015SPsWy8a63EpYMFDJjVE1e
1 parent 16b4658 commit 84b9ad9

6 files changed

Lines changed: 331 additions & 61 deletions

File tree

coderag/api.py

Lines changed: 37 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515

1616
from coderag._lines import split_lines
1717
from coderag.config import Config
18-
from coderag.types import IndexStats, SearchHit
18+
from coderag.types import IndexProgress, IndexStats, SearchHit
1919

2020
if TYPE_CHECKING: # avoid import-time cost / cycles
2121
from coderag.embeddings import EmbeddingProvider
@@ -39,6 +39,12 @@ def __init__(self, config: Optional[Config] = None) -> None:
3939
# surface, the MCP server's background index, and the live watcher) can't
4040
# interleave a file's delete-before-add sequence. Reads (search) are unaffected.
4141
self._index_lock = threading.Lock()
42+
# Guards the lazy construction of the collaborators below. The MCP server now serves
43+
# the protocol before warm-up finishes, so a query can land while the background
44+
# bootstrap is still building the store/provider — without this lock two threads could
45+
# each construct a second (conflicting) LanceStore. Reentrant because the properties
46+
# depend on each other (e.g. ``store`` reads ``provider`` while holding the lock).
47+
self._build_lock = threading.RLock()
4248

4349
# --- lazily constructed collaborators ---
4450

@@ -47,27 +53,34 @@ def provider(self) -> "EmbeddingProvider":
4753
if self._provider is None:
4854
from coderag.embeddings import get_provider
4955

50-
self._provider = get_provider(self.config)
56+
with self._build_lock:
57+
if self._provider is None:
58+
self._provider = get_provider(self.config)
5159
return self._provider
5260

5361
@property
5462
def store(self) -> "LanceStore":
5563
if self._store is None:
5664
from coderag.store.lance_store import LanceStore
5765

58-
self.config.store_dir.mkdir(parents=True, exist_ok=True)
59-
self._store = LanceStore(self.config.store_dir, self.provider.dim)
60-
# Clears the store when the embedding model/dim changed; a re-index then
61-
# repopulates the now-empty tables (there is no separate cache to rebuild).
62-
self._store.bootstrap(self.provider.dim, self.provider.model_id)
66+
with self._build_lock:
67+
if self._store is None:
68+
self.config.store_dir.mkdir(parents=True, exist_ok=True)
69+
store = LanceStore(self.config.store_dir, self.provider.dim)
70+
# Clears the store when the embedding model/dim changed; a re-index then
71+
# repopulates the now-empty tables (no separate cache to rebuild).
72+
store.bootstrap(self.provider.dim, self.provider.model_id)
73+
self._store = store
6374
return self._store
6475

6576
@property
6677
def indexer(self) -> "Indexer":
6778
if self._indexer is None:
6879
from coderag.indexer import Indexer
6980

70-
self._indexer = Indexer(self.config, self.provider, self.store)
81+
with self._build_lock:
82+
if self._indexer is None:
83+
self._indexer = Indexer(self.config, self.provider, self.store)
7184
return self._indexer
7285

7386
@property
@@ -76,27 +89,34 @@ def searcher(self) -> "HybridSearcher":
7689
from coderag.retrieval.rerank import get_reranker
7790
from coderag.retrieval.search import HybridSearcher
7891

79-
self._searcher = HybridSearcher(
80-
self.config,
81-
self.provider,
82-
self.store,
83-
reranker=get_reranker(self.config),
84-
)
92+
with self._build_lock:
93+
if self._searcher is None:
94+
self._searcher = HybridSearcher(
95+
self.config,
96+
self.provider,
97+
self.store,
98+
reranker=get_reranker(self.config),
99+
)
85100
return self._searcher
86101

87102
# --- public operations ---
88103

89104
def index(
90-
self, path: Optional[Union[str, Path]] = None, *, full: bool = False
105+
self,
106+
path: Optional[Union[str, Path]] = None,
107+
*,
108+
full: bool = False,
109+
live: Optional[IndexProgress] = None,
91110
) -> IndexStats:
92111
"""Incrementally index ``path`` (defaults to the configured watched dir).
93112
94113
Only files whose content hash changed are re-embedded. Pass ``full=True`` to
95-
force a clean rebuild.
114+
force a clean rebuild. Pass ``live`` (an :class:`IndexProgress`) to receive live,
115+
pollable progress — the MCP server uses this so ``index_status`` reflects the run.
96116
"""
97117
target = Path(path).expanduser() if path else self.config.watched_dir
98118
with self._index_lock:
99-
return self.indexer.index(target, full=full)
119+
return self.indexer.index(target, full=full, live=live)
100120

101121
def search(self, query: str, top_k: Optional[int] = None) -> List[SearchHit]:
102122
"""Hybrid (dense + lexical) search over the indexed codebase."""

coderag/indexer.py

Lines changed: 44 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -23,13 +23,19 @@
2323
from coderag.chunking.languages import detect_language
2424
from coderag.config import Config
2525
from coderag.embeddings import EmbeddingProvider
26-
from coderag.types import Chunk, IndexStats
26+
from coderag.types import Chunk, IndexProgress, IndexStats
2727

2828
if TYPE_CHECKING:
2929
from coderag.store.lance_store import LanceStore
3030

3131
logger = logging.getLogger(__name__)
3232

33+
# During a long initial index, commit buffered rows at least this often so dense search can
34+
# return partial results before the steady-state 8192-chunk flush boundary. Only applied when
35+
# a live progress object is supplied (the MCP background index), so the CLI/watcher keep their
36+
# single-flush-at-end batching.
37+
_PARTIAL_FLUSH_SECS = 5.0
38+
3339

3440
class _ProgressReporter:
3541
"""Live, human-facing indexing progress, written to stderr (stdout stays clean).
@@ -98,11 +104,21 @@ def index(
98104
*,
99105
full: bool = False,
100106
progress: bool = False,
107+
live: Optional[IndexProgress] = None,
101108
) -> IndexStats:
109+
"""Index ``target`` incrementally.
110+
111+
``progress`` enables the human-facing stderr narration; ``live`` is an optional
112+
machine-readable :class:`IndexProgress` the caller can poll concurrently (used by the
113+
MCP server's background index so ``index_status`` reflects live state). Both default
114+
off, so existing callers (CLI/watcher/tests) are unaffected.
115+
"""
102116
root = self.config.watched_dir.resolve()
103117
target = (target or self.config.watched_dir).resolve()
104118
prune = target == root # only a full-root pass removes vanished files
105119
rep = _ProgressReporter(progress)
120+
if live is not None:
121+
live.begin("scanning")
106122

107123
stats = IndexStats()
108124
if full:
@@ -121,6 +137,8 @@ def index(
121137
stats.files_skipped += 1
122138
else:
123139
work.append(item)
140+
if live is not None:
141+
live.saw_file(len(walked), len(work))
124142
rep.update(
125143
f"Scanning {target}{len(walked)} file(s) seen, "
126144
f"{len(work)} to index, {stats.files_skipped} unchanged/skipped…"
@@ -139,10 +157,20 @@ def index(
139157
# 2. (Re)index changed files. Chunking + embedding (the CPU/network cost) may run
140158
# in parallel across files (config.index_workers); the store writes stay on this
141159
# single thread to preserve the delete-before-add invariant and single writer.
142-
for added, removed in self._embed_and_write(work, reporter=rep):
160+
if live is not None and work:
161+
live.set_state("indexing")
162+
last_flush = time.monotonic()
163+
for item, added, removed in self._embed_and_write(work, reporter=rep):
143164
stats.chunks_added += added
144165
stats.chunks_removed += removed
145166
stats.files_indexed += 1
167+
if live is not None:
168+
live.wrote_file(item.rel, added)
169+
# Commit periodically so dense search picks up partials during a long initial
170+
# index, instead of waiting for the 8192-chunk boundary or the final persist.
171+
if time.monotonic() - last_flush > _PARTIAL_FLUSH_SECS:
172+
self.store.flush()
173+
last_flush = time.monotonic()
146174

147175
# 3. Prune files that disappeared from disk (full-root passes only).
148176
if prune:
@@ -157,6 +185,8 @@ def index(
157185
# never triggers a whole-index rebuild.
158186
changed = stats.files_indexed > 0 or stats.files_removed > 0
159187
if prune and changed:
188+
if live is not None:
189+
live.set_state("optimizing")
160190
self.store.optimize()
161191
else:
162192
self.store.flush()
@@ -168,6 +198,8 @@ def index(
168198
f"✓ Indexed {stats.files_indexed} file(s) — "
169199
f"{stats.total_files} total / {stats.total_chunks} chunks."
170200
)
201+
if live is not None:
202+
live.finish("ready")
171203
return stats
172204

173205
# --- internals ---
@@ -220,13 +252,17 @@ def _maybe_work(
220252

221253
def _embed_and_write(
222254
self, work: List[_Work], *, reporter: _ProgressReporter
223-
) -> Iterator[Tuple[int, int]]:
255+
) -> Iterator[Tuple[_Work, int, int]]:
224256
"""Chunk+embed each file (optionally across worker threads) and apply the writes.
225257
226258
Embedding is the expensive, parallelizable step and touches no shared mutable
227259
state, so it runs in a thread pool when ``index_workers > 1``. The store writes are
228260
drained here on the single calling thread, so the no-duplicate (delete-before-add)
229261
invariant and the single-writer store are preserved.
262+
263+
Yields ``(item, chunks_added, chunks_removed)`` per file — the ``_Work`` item is
264+
surfaced so the caller can report the current path (the worker pool completes out of
265+
order, so positional zipping back to ``work`` is not possible).
230266
"""
231267
if not work:
232268
return
@@ -239,14 +275,17 @@ def _embed_and_write(
239275
with ThreadPoolExecutor(max_workers=workers) as pool:
240276
futures = {pool.submit(self._prepare, item): item for item in work}
241277
for fut in as_completed(futures):
278+
item = futures[fut]
242279
chunks, vectors = fut.result()
243-
yield self._write(futures[fut], chunks, vectors)
280+
added, removed = self._write(item, chunks, vectors)
281+
yield item, added, removed
244282
done += 1
245283
reporter.update(f"Embedding {done}/{total} file(s)…")
246284
else:
247285
for item in work:
248286
chunks, vectors = self._prepare(item)
249-
yield self._write(item, chunks, vectors)
287+
added, removed = self._write(item, chunks, vectors)
288+
yield item, added, removed
250289
done += 1
251290
reporter.update(f"Embedding {done}/{total} file(s)…")
252291

0 commit comments

Comments
 (0)