java-codebase-rag 0.9.0__py3-none-any.whl → 0.9.1__py3-none-any.whl

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/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/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.dist-info → java_codebase_rag-0.9.1.dist-info}/METADATA

@@ -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.dist-info → java_codebase_rag-0.9.1.dist-info}/RECORD

@@ -5,7 +5,7 @@ chunk_heuristics.py,sha256=aQk2NOKxzUdqoUAJUO3G3LE0MN_bYZWNLQ0tkmj5uts,1813
 graph_enrich.py,sha256=Fxp7JurH9wWTOTX1nErjH4H6Jw5bB68YCqR17rsX6gw,62829
 graph_types.py,sha256=P6RVdEiqZlaSgTyXepnGbtMoZq0m_MaP2q6M6i7Uxlc,4539
 index_common.py,sha256=HT6FKHFJ084eFvd3fR1j8z8gf4eWoPHVW8GXLpw464I,285
-java_index_flow_lancedb.py,sha256=OFrvOV2xWPvBkQ1kS6nxcOjdJbYAGcEirzeevrOU1PM,25598
+java_index_flow_lancedb.py,sha256=kDkhnRxl-Kz0YNMY27DTk3gwa8H2nWyGDCQC8hxx-jI,28970
 java_index_v1_common.py,sha256=nF1KrSqboF_RRvWerG9knRRFmWwsrG_CvhgnsoZ8KqA,1154
 java_ontology.py,sha256=ooqr8GucOINpzhdEQ3QzVe5A9GfiR0nUTlySDehn9GA,17129
 ladybug_queries.py,sha256=rVnVEHwWwE4USeX7tICYEl1SSiSxJGnBNe1VX66R3Xk,100531
@@ -15,29 +15,29 @@ path_filtering.py,sha256=R--XzI51LXBu5IBKMCnJWbkNr6I5d-SDmltyQQnWco0,17674
 pr_analysis.py,sha256=zrmZZD5yotJtM02Kif6_jgI_oeformOao793akp0N6Y,18394
 resolve_service.py,sha256=tC5FQsGmqhqn0EOexVDlRq5egnzKDTrI7CMzk0nPpG8,25135
 search_lancedb.py,sha256=1sGSZ6H8J9hGcKAwHpzB7F0_in3A_sEtOI5LvlZYIRI,43864
-server.py,sha256=yNpJX_0D1xXY1FHGGBrb4JMOg2THw4_2Ao8YCdijBp0,35944
+server.py,sha256=nOK3DOr-i3PJnVmE7neye_tptJumgXfIXPlTIn01MPQ,37065
 java_codebase_rag/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 java_codebase_rag/_fdlimit.py,sha256=vkwjsPbZfxzZ2DZTPWO5DxtuNlLzOADzIq07iYX7GCU,2465
 java_codebase_rag/_stdio.py,sha256=TDNbpt2EP0_Zd622ihdlKwlMfxkKHWOfgLVcU6TcNbo,1458
 java_codebase_rag/cli.py,sha256=8WKk-_Zl1Sp7wHX4Y5f3bkIy8hHRH_QWrNuX_3WYTm4,45269
 java_codebase_rag/cli_format.py,sha256=CT7-xdwZ0bMCdP68_UOwkvm-mnLluU3LutlM-mDNk60,1839
 java_codebase_rag/cli_progress.py,sha256=q6Wh97yzLGs1B8UFk_WAKivfQu7Y5RnUUE-T2YHWkIs,3237
-java_codebase_rag/config.py,sha256=Yl7Nf0O_ZOZTPtyPMap83jRM7qNhWwfbVJTKSEdERi4,24916
+java_codebase_rag/config.py,sha256=1EAlKtQx7LUo-gPoE6BOVLz0YEl1OpWFRABkF06zjgk,26464
 java_codebase_rag/installer.py,sha256=c-_tR1Ct_O3yhmruFRnDoM1Wilz-igD9qdkDPsNqLMI,79934
 java_codebase_rag/jrag.py,sha256=cVUWrKOtkwe2r4u1nknzmrMQMy16bdt1n-dZQpiSFWs,191811
 java_codebase_rag/jrag_envelope.py,sha256=5jD3p2O-p9acAKHqoif5FFoqgP7Q2ziCuKSiPCHoywc,47505
 java_codebase_rag/jrag_hints.py,sha256=k2PFE4s3lZgBYHMdZcTjx1-w28nfQcBtQEVsSxI_DvE,9262
 java_codebase_rag/jrag_render.py,sha256=1nUyamL-MOOXDlKvssp-LsBgEtznUnGj_cPteSwuS90,31953
-java_codebase_rag/lance_optimize.py,sha256=_90eajcIpGUNm3GWtx6AYDERNHaMsaMZs7zu2KqNPEU,10347
-java_codebase_rag/pipeline.py,sha256=TkHb7DybFlpHje30aYbuFt5jWFSJBdCQ-7hc4jFaNhI,16481
+java_codebase_rag/lance_optimize.py,sha256=HI3aFebP1fenLL6Cav1jMG5kLXSrHLndO_MIJY6qQVo,11977
+java_codebase_rag/pipeline.py,sha256=L65mjK-IxkVWazUNVyIvzfQNi24behQu-Kdc7o_HwEk,17754
 java_codebase_rag/progress.py,sha256=2IxdMALDM0wAQCyJrrfZ975zM_85C-4BfHxf4AtYifE,23212
 java_codebase_rag/install_data/agents/explorer-rag-cli.md,sha256=mMij_BIQM4agaYhGVYjC-fQSe3We1HFeBvc5JBjJj6A,10071
 java_codebase_rag/install_data/agents/explorer-rag-enhanced.md,sha256=gZsNFbuK0lSnOIlplbbS_muz2ozokJqvFUv65QM0NDM,10406
 java_codebase_rag/install_data/skills/explore-codebase/SKILL.md,sha256=A-v2dueVnxwBzBlxoRjZ2zOJk8DranLQ1TElwn94h0s,11529
 java_codebase_rag/install_data/skills/explore-codebase-cli/SKILL.md,sha256=V5gIKKGkgk2KlAFf9JqPareQDMt1iQOI7cwhfOqJL1c,11348
-java_codebase_rag-0.9.0.dist-info/licenses/LICENSE,sha256=gxvtiHtuviR_q8ZAjWw-QTcF3DyPzg6ZY-lQrr8OPpw,1068
-java_codebase_rag-0.9.0.dist-info/METADATA,sha256=pIQfvYY6sWdmq-PIXm0vOUzNobTUtW14Z-Zp-53qaf8,20088
-java_codebase_rag-0.9.0.dist-info/WHEEL,sha256=K260EYznzXsJYBQGqmI8VTxEdiZYNvDZwW9cBh9-_MA,91
-java_codebase_rag-0.9.0.dist-info/entry_points.txt,sha256=cj3QTc11UYVQnj9T3orc4daiIGaCYrXP149vKbH2R4U,168
-java_codebase_rag-0.9.0.dist-info/top_level.txt,sha256=8vC-VN3cMwz5vhkSTaeJ1a1bDeqLWEfrTks1CvEvIg0,273
-java_codebase_rag-0.9.0.dist-info/RECORD,,
+java_codebase_rag-0.9.1.dist-info/licenses/LICENSE,sha256=gxvtiHtuviR_q8ZAjWw-QTcF3DyPzg6ZY-lQrr8OPpw,1068
+java_codebase_rag-0.9.1.dist-info/METADATA,sha256=v3Zl2nJ80PhrygxQD-aojw6z1CSZ_uRrOdkIWuA65No,20088
+java_codebase_rag-0.9.1.dist-info/WHEEL,sha256=K260EYznzXsJYBQGqmI8VTxEdiZYNvDZwW9cBh9-_MA,91
+java_codebase_rag-0.9.1.dist-info/entry_points.txt,sha256=cj3QTc11UYVQnj9T3orc4daiIGaCYrXP149vKbH2R4U,168
+java_codebase_rag-0.9.1.dist-info/top_level.txt,sha256=8vC-VN3cMwz5vhkSTaeJ1a1bDeqLWEfrTks1CvEvIg0,273
+java_codebase_rag-0.9.1.dist-info/RECORD,,

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(

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(