java-codebase-rag 0.9.0__tar.gz → 0.9.1__tar.gz

{java_codebase_rag-0.9.0/java_codebase_rag.egg-info → java_codebase_rag-0.9.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: java-codebase-rag
-Version: 0.9.0
+Version: 0.9.1
 Summary: MCP server for semantic + structural search over Java codebases
 Author: HumanBean17
 License-Expression: MIT

{java_codebase_rag-0.9.0 → java_codebase_rag-0.9.1}/java_codebase_rag/config.py

@@ -47,15 +47,39 @@ ENV_RUN_HEAVY = "JAVA_CODEBASE_RAG_RUN_HEAVY"
 COCOINDEX_MAX_INFLIGHT_COMPONENTS_ENV = "COCOINDEX_MAX_INFLIGHT_COMPONENTS"
 COCOINDEX_DEFAULT_MAX_INFLIGHT_COMPONENTS = "256"
 
+# Lance native DataFusion hash-join memory pool ceiling (FairSpillPool). The
+# lance default is ~100 MiB, tuned for query workloads — too small for the
+# single big ``merge_insert`` cocoindex emits at the end of a flow component.
+# On ``--full-reprocess`` (all rows match the existing table → bulk-update
+# path) the hash join builds on a large side and exhausts the pool somewhere
+# around 75k-100k chunks: "Resources exhausted: Failed to allocate ... for
+# HashJoinInput ... N MiB remain available for the total pool". cocoindex is a
+# bare pass-through to lancedb (it never sets a Session/memory_limit), so it
+# inherits this default — we raise it here. FairSpillPool is a *reservation
+# ceiling*, not a pre-allocation: setting 1 GiB does not reserve 1 GiB upfront,
+# it just allows the join to grow before spilling/erroring, so it is safe on
+# memory-constrained hosts. An operator can still override via their own
+# ``LANCE_MEM_POOL_SIZE`` (subprocess_env copies os.environ, and apply is via
+# ``setdefault`` so the operator value wins). Increment is unaffected (tiny
+# batch → tiny hash table); only the full-reprocess write path is at risk.
+LANCE_MEM_POOL_SIZE_ENV = "LANCE_MEM_POOL_SIZE"
+LANCE_DEFAULT_MEM_POOL_SIZE = "1073741824"  # 1 GiB
+
 
 def cocoindex_subprocess_env_defaults() -> dict[str, str]:
-    """Env defaults applied to every CocoIndex subprocess to bound concurrency.
+    """Env defaults applied to every CocoIndex subprocess.
+
+    Bounds CocoIndex concurrency (``COCOINDEX_MAX_INFLIGHT_COMPONENTS``; see
+    :issue:`306`) and raises the Lance hash-join memory ceiling
+    (``LANCE_MEM_POOL_SIZE``) so a large full-reprocess does not exhaust the
+    default ~100 MiB pool mid-``merge_insert``.
 
     Apply with ``env.setdefault(...)`` so a caller-provided (operator) value
-    always wins. See :issue:`306`.
+    always wins.
     """
     return {
-        COCOINDEX_MAX_INFLIGHT_COMPONENTS_ENV: COCOINDEX_DEFAULT_MAX_INFLIGHT_COMPONENTS
+        COCOINDEX_MAX_INFLIGHT_COMPONENTS_ENV: COCOINDEX_DEFAULT_MAX_INFLIGHT_COMPONENTS,
+        LANCE_MEM_POOL_SIZE_ENV: LANCE_DEFAULT_MEM_POOL_SIZE,
     }
 
 _DEFAULT_EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"

{java_codebase_rag-0.9.0 → java_codebase_rag-0.9.1}/java_codebase_rag/lance_optimize.py

@@ -184,6 +184,31 @@ async def optimize_lance_tables(
 
                 if last_exc is None:
                     results[name] = "ok"
+                    # Best-effort BTREE scalar index on the primary key ("id").
+                    # cocoindex's merge_insert defaults to use_index=True but
+                    # never creates a scalar PK index itself (declaring
+                    # primary_key in the schema does NOT auto-build a lance
+                    # index), so without this every merge_insert — increment
+                    # included — is a forced full scan of the PK column,
+                    # O(existing rows). On a large repo that scan dominates
+                    # increment wall-clock; with the index present the join does
+                    # lookups (~O(batch*log N)). Failure is non-fatal (the table
+                    # is still correct, just un-indexed) and never alters the
+                    # "ok" status, mirroring the FTS block below. ``replace=True``
+                    # keeps it idempotent across runs; table.optimize() above
+                    # maintains it on subsequent runs.
+                    try:
+                        from lancedb.index import BTree
+                        await table.create_index("id", config=BTree(), replace=True)
+                    except Exception as exc:
+                        low = str(exc).lower()
+                        if not any(
+                            w in low for w in ("exist", "duplicate", "already", "same name")
+                        ) and not quiet:
+                            print(
+                                f"java-codebase-rag: optimize: {name} id-index skipped: {exc}",
+                                file=sys.stderr,
+                            )
                     # Best-effort FTS index at index time (PR-SEARCH-3) so hybrid
                     # search works on all tables (java/sql/yaml) without a
                     # first-query race. Failure is non-fatal — the lazy

{java_codebase_rag-0.9.0 → java_codebase_rag-0.9.1}/java_codebase_rag/pipeline.py

@@ -132,6 +132,26 @@ def run_cocoindex_update(
     on_progress: Callable[[ProgressEvent], None] | None = None,
     on_progress_console: object | None = None,
 ) -> subprocess.CompletedProcess[str]:
+    if full_reprocess:
+        # A full reprocess rebuilds every row, so DROP the Lance target tables
+        # first and let cocoindex recreate them via the fast INSERT path. The
+        # in-place alternative (cocoindex's bulk-update merge_insert) emits
+        # ~one deletion-vector + version commit PER matched row — O(rows) of
+        # tiny file IO that scales to multi-minute hangs on large repos
+        # (measured ~83s sys time / 3474 deletion files for 3475 chunks on
+        # Shopizer; drop+recreate is ~3.6s sys / 0 deletions, ~3.7x faster and
+        # hang-free). Output is identical either way (full recompute); only the
+        # write path differs. Drop failure is non-fatal — if it somehow fails,
+        # the update falls back to the slow in-place path. The same fix is
+        # applied on the async server path (``server.run_refresh_pipeline``).
+        drop = run_cocoindex_drop(env, quiet=quiet)
+        if drop.returncode != 0 and not is_cocoindex_preflight_blocker(drop):
+            print(
+                "java-codebase-rag: drop-before-reprocess failed "
+                f"(exit {drop.returncode}); falling back to in-place update: "
+                f"{(drop.stderr or '').strip()[:200]}",
+                file=sys.stderr,
+            )
     result = _run_cocoindex_update_impl(
         env,
         full_reprocess=full_reprocess,

{java_codebase_rag-0.9.0 → java_codebase_rag-0.9.1/java_codebase_rag.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: java-codebase-rag
-Version: 0.9.0
+Version: 0.9.1
 Summary: MCP server for semantic + structural search over Java codebases
 Author: HumanBean17
 License-Expression: MIT

{java_codebase_rag-0.9.0 → java_codebase_rag-0.9.1}/java_codebase_rag.egg-info/SOURCES.txt

@@ -101,6 +101,7 @@ tests/test_microservice_scope.py
 tests/test_outgoing_call_extraction.py
 tests/test_packaging_metadata.py
 tests/test_path_filtering.py
+tests/test_pipeline.py
 tests/test_pr_analysis.py
 tests/test_progress.py
 tests/test_resolve_routes_messaging_layer_c.py

{java_codebase_rag-0.9.0 → java_codebase_rag-0.9.1}/java_index_flow_lancedb.py

@@ -112,6 +112,38 @@ _NUM_TXN_BEFORE_OPTIMIZE = 10**12
 # parent clamps to total on the terminal event anyway).
 _VECTORS_TICK_EVERY = 25
 
+# Bounded concurrency for the per-file drain in app_main. cocoindex's embedder
+# is ``@coco.fn.as_async(batching=True, runner=GPU, max_batch_size=64)`` — but
+# its batching layer only coalesces calls that are in flight SIMULTANEOUSLY. A
+# serial ``async for … await`` loop keeps just one file's chunks (avg 1–3) in
+# flight, so real batches stay tiny and MPS idles between them (measured ~138
+# chunks/s vs the ~235 chunks/s ceiling at batch=64 for all-MiniLM-L6-v2).
+# Draining many files at once with a semaphore puts their chunks in flight
+# together → the embedder coalesces them into full batches → MPS climbs toward
+# the ceiling. Measured on Shopizer (1167 files / 3475 chunks): full init drops
+# from ~46.7s (serial) to ~36.0s (32) / ~34.3s (64), with identical row output.
+#
+# This stays inside ONE component, so the earlier mount_each→app_main win is
+# preserved: still exactly ONE merge_insert per table at commit. Memoization
+# (``@coco.fn(memo=True)``) and the lock-guarded tick counter are safe under
+# concurrency; ``parse_java`` uses a per-thread tree-sitter Parser (already
+# routed via ``asyncio.to_thread``) and ``splitter.split`` is synchronous so the
+# event loop cannot reenter it.
+#
+# Default 64 is sized to MATCH the embedder's hardcoded max_batch_size=64 (the
+# decorator above; not a constructor arg, so not raisable from the flow): ~64
+# files in flight reliably fills a 64-chunk batch and saturates MPS. Going higher
+# buys nothing — the batch is already capped — and lower underfills it. Memory
+# is NOT the limiting factor here: cocoindex buffers ALL staged rows until the
+# single final merge_insert regardless of concurrency, so peak RSS is set by
+# total chunk count (the commit buffer), not by how many files process at once.
+# Set to ``1`` for the old serial behavior; raise/lower only if you have also
+# changed the effective batch size or are constraining the commit buffer itself.
+_FILE_CONCURRENCY = max(
+    1,
+    int(os.environ.get("JAVA_CODEBASE_RAG_FILE_CONCURRENCY", "64") or "64"),
+)
+
 # Thread-safe counter: cocoindex may call process_*_file concurrently
 # (mount_each parallelism is implementation-defined). A module-level lock guards
 # both the counter and the emission so two threads never interleave a tick.
@@ -533,6 +565,29 @@ async def process_yaml_file(
         )
 
 
+async def _drain_files_concurrently(
+    files: Any, process_fn: Any, table: Any, sem: asyncio.Semaphore
+) -> None:
+    """Run ``process_fn(file, table)`` over every file with bounded concurrency.
+
+    Replaces the serial ``async for … await process_*_file`` loop so the
+    embedder's batching layer sees many files' chunks in flight at once (see
+    ``_FILE_CONCURRENCY``). Materializes the async iterable up front — file
+    handles are lightweight and cocoindex already realized the collection when
+    the walker mounted, so this is not a second walk. An empty collection is a
+    no-op (e.g. SQL/YAML tables on a repo with none).
+    """
+    items = [f async for _, f in files.items()]
+    if not items:
+        return
+
+    async def _one(_file: Any) -> None:
+        async with sem:
+            await process_fn(_file, table)
+
+    await asyncio.gather(*(_one(f) for f in items))
+
+
 @coco.fn
 async def app_main() -> None:
     java_schema = await lancedb.TableSchema.from_class(
@@ -646,12 +701,17 @@ async def app_main() -> None:
     # unchanged files still skip re-embedding on incremental; _RowHandler.
     # reconcile skips rows whose fingerprint is unchanged → increment carries
     # only changed rows in its single merge_insert.
-    async for _key, _file in java_files.items():
-        await process_java_file(_file, java_table)
-    async for _key, _file in sql_files.items():
-        await process_sql_file(_file, sql_table)
-    async for _key, _file in yaml_files.items():
-        await process_yaml_file(_file, yaml_table)
+    #
+    # PERF (concurrency): drain files with a bounded semaphore instead of a
+    # serial ``async for … await``. See ``_FILE_CONCURRENCY`` — this is what
+    # lets the embedder's batching layer fill real batches (embedding dominates
+    # init cost, and serial files starve it). One shared semaphore bounds total
+    # in-flight work; tables are drained in order (java dominates, sql/yaml are
+    # usually near-empty).
+    _sem = asyncio.Semaphore(_FILE_CONCURRENCY)
+    await _drain_files_concurrently(java_files, process_java_file, java_table, _sem)
+    await _drain_files_concurrently(sql_files, process_sql_file, sql_table, _sem)
+    await _drain_files_concurrently(yaml_files, process_yaml_file, yaml_table, _sem)
 
 
 app = coco.App(

{java_codebase_rag-0.9.0 → java_codebase_rag-0.9.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "java-codebase-rag"
-version = "0.9.0"
+version = "0.9.1"
 description = "MCP server for semantic + structural search over Java codebases"
 readme = "README.md"
 requires-python = ">=3.11"

{java_codebase_rag-0.9.0 → java_codebase_rag-0.9.1}/server.py

@@ -349,6 +349,31 @@ async def run_refresh_pipeline(
             )
     proc: asyncio.subprocess.Process | None = None
     out_b, err_b = b"", b""
+    # DROP the Lance target tables so the update takes the fast INSERT path
+    # instead of cocoindex's in-place bulk-update, which emits ~one deletion-
+    # vector + version commit PER matched row — O(rows) of tiny file IO that
+    # hangs for many minutes on large repos. Drop+recreate is identical output
+    # for a full rebuild (the very thing --full-reprocess means). Same fix on
+    # the sync path: pipeline.run_cocoindex_update. Drop failure is non-fatal:
+    # the update falls back to the slow in-place path.
+    try:
+        drop_proc = await asyncio.create_subprocess_exec(
+            str(cocoindex_bin),
+            "drop",
+            _COCOINDEX_TARGET,
+            "-f",
+            cwd=str(flow_path.parent),
+            env=_cocoindex_subprocess_env(root),
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+        )
+        await drop_proc.communicate()
+    except Exception as exc:
+        print(
+            f"java-codebase-rag: drop-before-reprocess failed ({exc!s}); "
+            "falling back to in-place update",
+            file=sys.stderr,
+        )
     if quiet:
         try:
             proc = await asyncio.create_subprocess_exec(

{java_codebase_rag-0.9.0 → java_codebase_rag-0.9.1}/tests/test_config.py

@@ -506,6 +506,10 @@ def test_cocoindex_subprocess_env_defaults_uses_real_inflight_env_var() -> None:
     assert defaults["COCOINDEX_MAX_INFLIGHT_COMPONENTS"] == "256"
     # The bogus name from the broken #293 fix must NOT leak back in.
     assert "COCOINDEX_SOURCE_MAX_INFLIGHT_ROWS" not in defaults
+    # Lance native hash-join pool ceiling, raised from the ~100 MiB default so a
+    # large full-reprocess merge_insert does not exhaust it mid-commit. Applied
+    # via setdefault, so an operator's own value still wins.
+    assert defaults["LANCE_MEM_POOL_SIZE"] == "1073741824"
 
 
 class TestConfigSourcePointer:

{java_codebase_rag-0.9.0 → java_codebase_rag-0.9.1}/tests/test_lance_optimize.py

@@ -123,31 +123,92 @@ async def test_optimize_builds_fts_index_after_success(monkeypatch, tmp_path) ->
     table = _FakeTable(name, [None])  # optimize succeeds on first try
     conn = _FakeConnection(table_names={name}, tables={name: table})
     _install_fake_lancedb(monkeypatch, conn)
-    # Make ``from lancedb.index import FTS`` resolve against a stand-in module.
-    index_mod = types.ModuleType("lancedb.index")
-
-    class FTS:  # stands in for lancedb.index.FTS config object
-        pass
-
-    index_mod.FTS = FTS
-    monkeypatch.setitem(sys.modules, "lancedb.index", index_mod)
+    fts_cls, _btree_cls = _install_fake_lancedb_index(monkeypatch)
 
     results = await lance_optimize.optimize_lance_tables(tmp_path, quiet=True)
 
     assert results[name] == "ok"
-    assert len(table.create_index_calls) == 1, (
-        f"expected FTS create_index once after optimize, got {table.create_index_calls}"
-    )
-    call = table.create_index_calls[0]
-    assert call["args"] and call["args"][0] == "text", (
-        f"FTS index must target the 'text' column, got args={call['args']}"
+    # optimize now builds TWO indices (BTREE on "id" + FTS on "text"); find the
+    # FTS one by column rather than asserting on call count / order.
+    fts_calls = [
+        c for c in table.create_index_calls if c["args"] and c["args"][0] == "text"
+    ]
+    assert len(fts_calls) == 1, (
+        f"expected one FTS create_index on 'text', got {table.create_index_calls}"
     )
+    call = fts_calls[0]
     assert call["kwargs"].get("replace") is True, (
         f"FTS index must be built with replace=True, got kwargs={call['kwargs']}"
     )
+    assert isinstance(call["kwargs"].get("config"), fts_cls), (
+        f"FTS index config must be a lancedb.index.FTS, got kwargs={call['kwargs']}"
+    )
     assert conn.closed is True
 
 
+async def test_optimize_builds_btree_pk_index_after_success(monkeypatch, tmp_path) -> None:
+    """A successful optimize builds a BTREE scalar index on the ``id`` PK.
+
+    cocoindex's ``merge_insert`` defaults to ``use_index=True`` but never creates
+    a scalar PK index itself (declaring ``primary_key`` does NOT auto-build a
+    lance index), so without this every ``merge_insert`` — increment included —
+    is a forced full scan of the PK column (O(existing rows)). The BTREE index
+    flips that to ~O(batch*log N) lookups. Guards the path against silent
+    regression under the broad ``except`` that swallows index-build failures.
+    """
+    import types
+
+    from java_codebase_rag import lance_optimize
+
+    name = lance_optimize.LANCE_TABLE_NAMES[0]
+    table = _FakeTable(name, [None])  # optimize succeeds on first try
+    conn = _FakeConnection(table_names={name}, tables={name: table})
+    _install_fake_lancedb(monkeypatch, conn)
+    _fts_cls, btree_cls = _install_fake_lancedb_index(monkeypatch)
+
+    results = await lance_optimize.optimize_lance_tables(tmp_path, quiet=True)
+
+    assert results[name] == "ok"
+    id_calls = [
+        c for c in table.create_index_calls if c["args"] and c["args"][0] == "id"
+    ]
+    assert len(id_calls) == 1, (
+        f"expected one BTREE create_index on 'id', got {table.create_index_calls}"
+    )
+    call = id_calls[0]
+    assert call["kwargs"].get("replace") is True, (
+        f"BTREE PK index must be built with replace=True, got kwargs={call['kwargs']}"
+    )
+    assert isinstance(call["kwargs"].get("config"), btree_cls), (
+        f"BTREE PK index config must be a lancedb.index.BTree, got kwargs={call['kwargs']}"
+    )
+
+
+def _install_fake_lancedb_index(monkeypatch):
+    """Make ``from lancedb.index import FTS, BTree`` resolve to stand-in classes.
+
+    Returns the ``(FTS, BTree)`` config stand-ins so a test can assert the right
+    config object was passed. Both the FTS path (PR-SEARCH-3) and the BTREE PK
+    path build their index config via a local ``from lancedb.index import ...``;
+    without a stand-in those imports hit the broad ``except`` and silently no-op,
+    leaving the path unexercised.
+    """
+    import types
+
+    index_mod = types.ModuleType("lancedb.index")
+
+    class FTS:  # stands in for lancedb.index.FTS config object
+        pass
+
+    class BTree:  # stands in for lancedb.index.BTree config object
+        pass
+
+    index_mod.FTS = FTS
+    index_mod.BTree = BTree
+    monkeypatch.setitem(sys.modules, "lancedb.index", index_mod)
+    return FTS, BTree
+
+
 async def test_optimize_does_not_retry_non_conflict_error(monkeypatch, tmp_path) -> None:
     """A non-conflict exception is re-raised (captured per-table), never retried."""
     from java_codebase_rag import lance_optimize

java_codebase_rag-0.9.1/tests/test_pipeline.py

@@ -0,0 +1,103 @@
+"""Tests for ``java_codebase_rag.pipeline`` subprocess helpers.
+
+Focus: ``run_cocoindex_update`` drops the Lance target tables before a *full
+reprocess* so the update takes the fast INSERT path. The in-place alternative
+(cocoindex's bulk-update ``merge_insert``) emits ~one deletion-vector + version
+commit per matched row — O(rows) of tiny file IO that hangs for many minutes on
+large repos. Drop+recreate is identical output for a full rebuild.
+"""
+from __future__ import annotations
+
+import subprocess
+
+from java_codebase_rag import pipeline
+
+
+def _ok() -> subprocess.CompletedProcess:
+    return subprocess.CompletedProcess(args=[], returncode=0, stdout="", stderr="")
+
+
+def _stub_impl(monkeypatch, seen: dict) -> None:
+    """Replace the impl + post-optimize with no-op stubs (no cocoindex/lancedb)."""
+
+    def fake_impl(env, **kwargs):
+        seen["update"] = seen.get("update", 0) + 1
+        seen["full_reprocess"] = kwargs.get("full_reprocess")
+        return _ok()
+
+    monkeypatch.setattr(pipeline, "_run_cocoindex_update_impl", fake_impl)
+    monkeypatch.setattr(pipeline, "_maybe_run_serialized_optimize", lambda *a, **k: None)
+
+
+def test_full_reprocess_drops_tables_first(monkeypatch) -> None:
+    """full_reprocess=True drops exactly once before the update (INSERT path)."""
+    seen: dict = {}
+    _stub_impl(monkeypatch, seen)
+    drops: list[dict] = []
+
+    def fake_drop(env, *, quiet):
+        drops.append(env)
+        return _ok()
+
+    monkeypatch.setattr(pipeline, "run_cocoindex_drop", fake_drop)
+
+    pipeline.run_cocoindex_update({"X": "1"}, full_reprocess=True, quiet=True)
+
+    assert len(drops) == 1, "full_reprocess must drop exactly once before update"
+    assert drops[0] == {"X": "1"}, "drop must receive the same env as the update"
+    assert seen["update"] == 1
+
+
+def test_increment_does_not_drop(monkeypatch) -> None:
+    """full_reprocess=False (increment) must NOT drop — it would lose the table."""
+    seen: dict = {}
+    _stub_impl(monkeypatch, seen)
+    drops: list[dict] = []
+
+    def fake_drop(env, *, quiet):
+        drops.append(env)
+        return _ok()
+
+    monkeypatch.setattr(pipeline, "run_cocoindex_drop", fake_drop)
+
+    pipeline.run_cocoindex_update({}, full_reprocess=False, quiet=True)
+
+    assert drops == [], "increment must not drop the tables"
+    assert seen["update"] == 1
+
+
+def test_drop_failure_falls_back_to_inplace(monkeypatch, capsys) -> None:
+    """A non-preflight drop failure does not abort — the update still runs in-place."""
+    seen: dict = {}
+    _stub_impl(monkeypatch, seen)
+    monkeypatch.setattr(
+        pipeline,
+        "run_cocoindex_drop",
+        lambda env, *, quiet: subprocess.CompletedProcess(
+            args=[], returncode=2, stdout="", stderr="boom"
+        ),
+    )
+
+    pipeline.run_cocoindex_update({}, full_reprocess=True, quiet=True)
+
+    assert seen["update"] == 1, "update must still run after a non-fatal drop failure"
+    assert "drop-before-reprocess failed" in capsys.readouterr().err
+
+
+def test_drop_preflight_blocker_is_silent(monkeypatch, capsys) -> None:
+    """A preflight drop stub (cocoindex not installed, e.g. graph-only) is not noisy."""
+    seen: dict = {}
+    _stub_impl(monkeypatch, seen)
+    monkeypatch.setattr(
+        pipeline,
+        "run_cocoindex_drop",
+        lambda env, *, quiet: subprocess.CompletedProcess(
+            args=["cocoindex"], returncode=127, stdout="", stderr="not found"
+        ),
+    )
+
+    pipeline.run_cocoindex_update({}, full_reprocess=True, quiet=True)
+
+    assert seen["update"] == 1
+    # 127 preflight is expected on graph-only installs and must NOT warn.
+    assert "drop-before-reprocess failed" not in capsys.readouterr().err