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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: java-codebase-rag
-Version: 0.9.0
+Version: 0.9.2
 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.2}/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.2}/java_codebase_rag/jrag.py

@@ -812,11 +812,27 @@ def build_parser() -> argparse.ArgumentParser:
             "CALLS hops from the route handler. Intra-service is an INDEX-TIME data "
             "property: CALLS edges are intra-codebase by construction, and the query "
             "carries no microservice predicate, so the result reflects whatever the "
-            "fixture indexed (no query-time constraint). --max-hops clamped to 1..8."
+            "fixture indexed (no query-time constraint). --depth clamped to 1..8."
         ),
     )
-    flow.add_argument("query", help="Route path (e.g. '/chat/assign'). Resolved with hint_kind=route.")
-    flow.add_argument("--max-hops", type=int, default=5, dest="max_hops", help="Max CALLS hops (clamped 1..8, default 5).")
+    flow.add_argument(
+        "query",
+        help=(
+            "Route path (e.g. '/chat/assign') or Kafka topic name (e.g. "
+            "'banking.chat.compliance.review'). Resolved with hint_kind=route; "
+            "kafka_topic Routes match on topic."
+        ),
+    )
+    # Primary flag is --depth (consistent with callers/callees/impact/decompose).
+    # --max-hops is kept as a hidden back-compat alias (same dest).
+    flow.add_argument(
+        "--depth", type=int, default=5, dest="depth",
+        help="Max CALLS hops (clamped 1..8, default 5).",
+    )
+    flow.add_argument(
+        "--max-hops", type=int, dest="depth",
+        default=argparse.SUPPRESS, help=argparse.SUPPRESS,
+    )
     flow.set_defaults(handler=_cmd_flow)
 
     # ---- Compose traversals + file inspection (PR-JRAG-3b) ----
@@ -3001,7 +3017,7 @@ def _cmd_flow(args: argparse.Namespace) -> int:
         args, reason="trace_request_flow carries no microservice predicate; intra-codebase is an index-time data property"
     )
 
-    max_hops = max(1, min(8, getattr(args, "max_hops", 5)))
+    max_hops = max(1, min(8, getattr(args, "depth", 5)))
     flow_data = graph.trace_request_flow(entry_route_id=node.id, max_hops=max_hops)
 
     root_id = node.id

{java_codebase_rag-0.9.0 → java_codebase_rag-0.9.2}/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.2}/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.2/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.2
 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.2}/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.2}/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.2}/java_ontology.py

@@ -90,6 +90,8 @@ VALID_RESOLVE_REASONS: frozenset[str] = frozenset((
     "short_name",
     "route_template",
     "route_method_path",
+    "route_topic",
+    "route_topic_prefix",
     "client_target",
     "client_target_path",
     "client_name",
@@ -428,6 +430,8 @@ ResolveReason = Literal[
     "short_name",
     "route_template",
     "route_method_path",
+    "route_topic",
+    "route_topic_prefix",
     "client_target",
     "client_target_path",
     "client_name",

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "java-codebase-rag"
-version = "0.9.0"
+version = "0.9.2"
 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.2}/resolve_service.py

@@ -43,11 +43,13 @@ _RESOLVE_REASON_PRIORITY: dict[ResolveReason, int] = {
     "producer_topic_prefix": 1,
     "fqn_suffix": 2,
     "route_template": 2,
+    "route_topic": 2,
     "client_fqn": 2,
     "short_name": 3,
     "client_target": 3,
     "client_name": 3,
     "producer_topic": 3,
+    "route_topic_prefix": 3,
 }
 
 _SYMBOL_RESOLVE_RETURN = (
@@ -295,6 +297,28 @@ def _resolve_route_candidates(
             path_val = str(row.get("path_template") or row.get("path") or "")
             out.append((_node_ref_from_row("route", row), "route_template", len(path_val)))
 
+    # Kafka/topic routes carry their name in ``topic`` (``path``/``path_template``
+    # are empty), so path-based matching above cannot reach them. Match on
+    # ``r.topic`` the same way ``_resolve_producer_candidates`` matches
+    # ``p.topic`` — this lets ``flow``/``callers``/``overview`` resolve a
+    # ``kafka_topic`` Route by topic name. ``_drop_route_mirrors`` below then
+    # discards the no-EXPOSES producer phantom in favour of the server route.
+    rows = g._rows(  # noqa: SLF001
+        f"MATCH (r:Route) WHERE r.topic = $topic{scope} RETURN {_ROUTE_RESOLVE_RETURN} LIMIT $lim",
+        {"topic": identifier, "lim": lim, **scope_params},
+    )
+    for row in rows:
+        out.append((_node_ref_from_row("route", row), "route_topic", len(identifier)))
+
+    if not identifier.startswith("/"):
+        rows = g._rows(  # noqa: SLF001
+            f"MATCH (r:Route) WHERE r.topic STARTS WITH $topic{scope} "
+            f"RETURN {_ROUTE_RESOLVE_RETURN} LIMIT $lim",
+            {"topic": identifier, "lim": lim, **scope_params},
+        )
+        for row in rows:
+            out.append((_node_ref_from_row("route", row), "route_topic_prefix", len(identifier)))
+
     return _drop_route_mirrors(g, out)
 
 

{java_codebase_rag-0.9.0 → java_codebase_rag-0.9.2}/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.2}/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.2}/tests/test_jrag_traversal_direct.py

@@ -22,8 +22,11 @@ Tests (bank-chat fixture):
 13. test_impact_service_post_filter_emits_warning
 14. test_decompose_renders_role_waterfall
 15. test_flow_outbound_intra_service_on_fixture
-16. test_traversal_resolve_ambiguous_stops
-17. test_traversal_rejects_offset
+16. test_flow_follows_kafka_topic_on_fixture
+17. test_flow_depth_flag_and_max_hops_alias
+18. test_callers_topic_disambiguates_with_kind
+19. test_traversal_resolve_ambiguous_stops
+20. test_traversal_rejects_offset
 """
 from __future__ import annotations
 
@@ -620,6 +623,134 @@ def test_flow_outbound_intra_service_on_fixture(
     )
 
 
+def test_flow_follows_kafka_topic_on_fixture(
+    corpus_root: Path, ladybug_db_path: Path
+) -> None:
+    """flow resolves a Kafka topic to its Route AND follows it (async edges).
+
+    Regression: ``_resolve_route_candidates`` matched only on path, so a
+    ``kafka_topic`` Route (name in ``topic``, ``path=''``) was unresolvable and
+    ``jrag flow <topic>`` returned 'none'. Resolution alone is not enough — this
+    test also asserts the follow graph is non-empty and carries an
+    ``ASYNC_CALLS`` edge, proving the kafka-specific inbound arm (topic-matched
+    Producer) actually fires once the route resolves.
+
+    ``banking.chat.incoming`` is the canonical dual-sided topic: produced by
+    ``FollowUpKafkaPublisher`` (inbound ``ASYNC_CALLS``) and consumed by
+    ``ChatKafkaListener`` (outbound ``CALLS`` to ``orchestrationService.handle``).
+    """
+    env = _env_for(corpus_root, ladybug_db_path)
+    proc = _run_jrag(
+        ["flow", "banking.chat.incoming", "--format", "json"], env=env
+    )
+    assert proc.returncode == 0, (
+        f"flow on kafka topic failed: rc={proc.returncode}\nstdout={proc.stdout}\nstderr={proc.stderr}"
+    )
+    payload = json.loads(proc.stdout)
+    assert payload["status"] == "ok", f"expected ok, got {payload}"
+    assert payload.get("root"), "expected root id (the kafka_topic Route)"
+    root_node = payload.get("nodes", {}).get(payload["root"], {})
+    assert root_node.get("kind") == "route", f"expected route root, got {root_node}"
+    # The follow graph must be non-empty AND carry a kafka async inbound edge —
+    # this is the part that proves "flow follows kafka", not just "resolves it".
+    edges = payload.get("edges", [])
+    assert edges, f"expected non-empty follow graph for a kafka topic, got edges={edges}"
+    assert any(e.get("edge_type") == "ASYNC_CALLS" for e in edges), (
+        f"expected an ASYNC_CALLS (topic-matched Producer) inbound edge, got {edges}"
+    )
+
+
+def test_flow_depth_flag_and_max_hops_alias(
+    corpus_root: Path, ladybug_db_path: Path
+) -> None:
+    """flow uses --depth (consistent with callers/callees/impact/decompose).
+
+    --max-hops remains as a hidden back-compat alias (same dest). Both must be
+    accepted, produce the same traversal, and --depth must actually change the
+    traversal depth (more hops => strictly more CALLS edges).
+    """
+    env = _env_for(corpus_root, ladybug_db_path)
+
+    # --depth is the primary flag now.
+    proc_depth = _run_jrag(
+        ["flow", "/chat/assign", "--depth", "2", "--format", "json"], env=env
+    )
+    assert proc_depth.returncode == 0, (
+        f"flow --depth failed: rc={proc_depth.returncode}\nstdout={proc_depth.stdout}"
+    )
+    payload_depth = json.loads(proc_depth.stdout)
+    assert payload_depth["status"] == "ok", f"expected ok, got {payload_depth}"
+
+    # --max-hops is the hidden alias; must still be accepted (back-compat).
+    proc_alias = _run_jrag(
+        ["flow", "/chat/assign", "--max-hops", "2", "--format", "json"], env=env
+    )
+    assert proc_alias.returncode == 0, (
+        f"flow --max-hops alias failed: rc={proc_alias.returncode}\nstdout={proc_alias.stdout}"
+    )
+    payload_alias = json.loads(proc_alias.stdout)
+    assert payload_alias["status"] == "ok", f"expected ok, got {payload_alias}"
+    # Same dest -> same traversal (depth 2 in both cases).
+    assert payload_alias.get("root") == payload_depth.get("root"), (
+        "expected identical root id for --depth and --max-hops"
+    )
+
+    # --depth must actually affect the traversal: depth 1 yields strictly fewer
+    # outbound CALLS edges than depth 5 on this fixture (~3 vs ~20).
+    proc_shallow = _run_jrag(
+        ["flow", "/chat/assign", "--depth", "1", "--format", "json"], env=env
+    )
+    proc_deep = _run_jrag(
+        ["flow", "/chat/assign", "--depth", "5", "--format", "json"], env=env
+    )
+    shallow = json.loads(proc_shallow.stdout)
+    deep = json.loads(proc_deep.stdout)
+    assert len(shallow.get("edges", [])) < len(deep.get("edges", [])), (
+        f"--depth should change traversal size; depth=1 -> {len(shallow.get('edges', []))}, "
+        f"depth=5 -> {len(deep.get('edges', []))}"
+    )
+
+
+def test_callers_topic_disambiguates_with_kind(
+    corpus_root: Path, ladybug_db_path: Path
+) -> None:
+    """A dual-sided topic resolves to BOTH a Producer and a Route -> ambiguous
+    without --kind; --kind route collapses to one and runs find_route_callers.
+
+    Behavior note (from enabling r.topic resolution): ``callers <topic>`` with no
+    --kind now searches all kinds, and a topic with both a Producer and a server
+    Route is genuinely ambiguous (previously only the Producer matched, which
+    then errored because callers only accepts Symbol/Route roots — jrag.py:2371).
+    ``--kind route`` is the disambiguation that yields a working callers run.
+    (--kind producer resolves to the Producer but callers rejects a Producer
+    root; use the ``producers`` command for the producer view.)
+    """
+    env = _env_for(corpus_root, ladybug_db_path)
+
+    # No --kind: Producer + Route both match -> ambiguous (no traversal).
+    proc_both = _run_jrag(
+        ["callers", "banking.chat.incoming", "--format", "json"], env=env
+    )
+    payload_both = json.loads(proc_both.stdout)
+    assert payload_both["status"] == "ambiguous", (
+        f"expected ambiguous for dual-sided topic without --kind, got {payload_both.get('status')}"
+    )
+    assert len(payload_both.get("candidates", [])) >= 2, (
+        f"expected >=2 candidates (producer + route), got {payload_both.get('candidates')}"
+    )
+
+    # --kind route: resolves to the single server Route and runs find_route_callers.
+    proc_route = _run_jrag(
+        ["callers", "banking.chat.incoming", "--kind", "route", "--format", "json"],
+        env=env,
+    )
+    payload_route = json.loads(proc_route.stdout)
+    assert payload_route["status"] == "ok", (
+        f"expected --kind route to resolve and run, got {payload_route.get('status')}"
+    )
+    assert payload_route.get("root"), "expected a resolved root for --kind route"
+
+
 # ----- Test 16: traversal resolve-ambiguous stops (no auto-pick) -----
 
 

{java_codebase_rag-0.9.0 → java_codebase_rag-0.9.2}/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