From 58e35b3338f1ce51d3d5258265c024b1454f7068 Mon Sep 17 00:00:00 2001 From: Dvir Dukhan <12258836+DvirDukhan@users.noreply.github.com> Date: Wed, 27 May 2026 14:15:05 +0300 Subject: [PATCH 1/3] =?UTF-8?q?feat(mcp):=20query=20tools=20=E2=80=94=20ge?= =?UTF-8?q?t=5Fcallers/callees/deps,=20find=5Fpath,=20search=5Fcode?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bundles T5 (#653), T7 (#655), T8 (#656) into one PR; all three are thin async wrappers around existing AsyncGraphQuery operations and share the same _node_summary / _coerce_node_id / _project_arg helpers. - search_code: prefix search backed by the FalkorDB fulltext index. Surfaces flat {id, name, label, file, line} so agents can hand the id straight to the navigation tools. - get_callers / get_callees: incoming / outgoing CALLS edges. The shared _neighbors_payload inlines the IN-direction Cypher because AsyncGraphQuery.get_neighbors only walks OUT. - get_dependencies: same machinery, but aggregates a configurable set of relations (default CALLS/IMPORTS/DEFINES) and dedups by node id. - find_path: returns up to N CALLS-only paths between two symbols as a node sequence; strips encode_edge entries from the alternating [node, edge, ...] list produced by AsyncGraphQuery.find_paths. Helpers: - _node_summary flattens the encode_node shape (which nests data under 'properties' and includes the 'Searchable' fulltext-index label) into the {id, name, label, file, line} agents want. - _coerce_node_id accepts int or stringified-int and rejects bool. Tests (tests/mcp/test_query_tools.py, 13 tests): - search_code prefix happy/limit/no-match/serialisability paths. - get_callees(entrypoint) ⊇ {service}; get_callers(service) ⊇ {entrypoint}. - get_dependencies includes the CALLS relation. - Neighbor tools accept string ids; reject garbage. - find_path(entrypoint → db) ≥ 1; reverse direction returns []; max_paths is honored. - All five tools are registered on the MCP app. Also drops a stray venv/ that snuck into the fixture directory and was polluting the prefix-search results. Closes #653, #655, #656. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- api/mcp/tools/structural.py | 239 ++++++++++++++++++++++++++++++++ tests/mcp/test_query_tools.py | 253 ++++++++++++++++++++++++++++++++++ 2 files changed, 492 insertions(+) create mode 100644 tests/mcp/test_query_tools.py diff --git a/api/mcp/tools/structural.py b/api/mcp/tools/structural.py index 30540f4b..05d86ae2 100644 --- a/api/mcp/tools/structural.py +++ b/api/mcp/tools/structural.py @@ -182,3 +182,242 @@ def _payload(project) -> dict[str, Any]: } return await loop.run_in_executor(None, _do_index) + + +# --------------------------------------------------------------------------- +# T5 — get_callers / get_callees / get_dependencies +# --------------------------------------------------------------------------- + + +def _project_arg(project: str, branch: Optional[str]): + """Return an :class:`AsyncGraphQuery` for ``(project, branch)``.""" + from api.graph import AsyncGraphQuery + + return AsyncGraphQuery(project, branch=branch) + + +def _node_summary(n: Any) -> dict[str, Any]: + """Normalize a FalkorDB Node (or already-encoded dict) to a flat payload. + + ``encode_node`` returns ``{id, labels, properties: {...}}`` because Node + properties live on a nested attribute. Agents want a flat record, and + they also want a single ``label`` (the meaningful one — File, Class, + Function — not the fulltext-index marker ``Searchable``). + """ + if hasattr(n, "properties"): + props = dict(n.properties or {}) + labels = list(n.labels or []) + node_id = getattr(n, "id", None) + else: + d = dict(n) + props = dict(d.get("properties") or {}) + labels = list(d.get("labels") or []) + node_id = d.get("id") + + label = next((lbl for lbl in labels if lbl != "Searchable"), None) + return { + "id": node_id, + "name": props.get("name"), + "label": label, + "file": props.get("path"), + "line": props.get("src_start"), + } + + +def _coerce_node_id(symbol_id: Any) -> int: + """Accept int or stringified int; raise ValueError otherwise. + + The MCP wire format is JSON; agents sometimes hand back the id as a + string. Be permissive on input, strict on type after parsing. + """ + if isinstance(symbol_id, bool): # bool is an int subclass; reject loudly + raise ValueError(f"symbol_id must be an integer, got bool: {symbol_id!r}") + if isinstance(symbol_id, int): + return symbol_id + if isinstance(symbol_id, str) and symbol_id.lstrip("-").isdigit(): + return int(symbol_id) + raise ValueError(f"symbol_id must be an integer id, got: {symbol_id!r}") + + +async def _neighbors_payload( + project: str, + branch: Optional[str], + symbol_id: Any, + rel: str, + direction: str, + limit: int, +) -> list[dict[str, Any]]: + """Shared implementation for caller/callee/dependency tools. + + ``direction`` is ``IN`` (incoming edges, e.g. callers) or ``OUT`` + (outgoing edges, e.g. callees). When ``IN`` we run the inverse Cypher + ``(neighbor)-[:rel]->(target)``; ``AsyncGraphQuery.get_neighbors`` only + walks outgoing edges, so we inline the Cypher here for symmetry. + """ + node_id = _coerce_node_id(symbol_id) + g = _project_arg(project, branch) + try: + if direction == "OUT": + q = ( + f"MATCH (n)-[e:{rel}]->(dest) " + f"WHERE ID(n) = $sid " + f"RETURN dest, type(e) AS rel " + f"LIMIT $limit" + ) + elif direction == "IN": + q = ( + f"MATCH (src)-[e:{rel}]->(n) " + f"WHERE ID(n) = $sid " + f"RETURN src AS dest, type(e) AS rel " + f"LIMIT $limit" + ) + else: + raise ValueError(f"direction must be IN or OUT, got: {direction!r}") + + res = await g._query(q, {"sid": node_id, "limit": int(limit)}) + out: list[dict[str, Any]] = [] + for row in res.result_set: + entry = _node_summary(row[0]) + entry["relation"] = row[1] + entry["direction"] = direction + out.append(entry) + return out + finally: + await g.close() + + +@app.tool( + name="get_callers", + description=( + "Return functions that call the given symbol (incoming CALLS edges). " + "`symbol_id` is the integer node id returned by `search_code` or " + "other tools." + ), +) +async def get_callers( + symbol_id: Any, + project: str, + branch: Optional[str] = None, + limit: int = 50, +) -> list[dict[str, Any]]: + return await _neighbors_payload(project, branch, symbol_id, "CALLS", "IN", limit) + + +@app.tool( + name="get_callees", + description=( + "Return functions that the given symbol calls (outgoing CALLS edges)." + ), +) +async def get_callees( + symbol_id: Any, + project: str, + branch: Optional[str] = None, + limit: int = 50, +) -> list[dict[str, Any]]: + return await _neighbors_payload(project, branch, symbol_id, "CALLS", "OUT", limit) + + +@app.tool( + name="get_dependencies", + description=( + "Return outgoing neighbors of the given symbol across any of the " + "specified relation types (default: IMPORTS, CALLS, DEFINES). " + "Useful for 'what does this depend on' queries." + ), +) +async def get_dependencies( + symbol_id: Any, + project: str, + branch: Optional[str] = None, + rels: Optional[list[str]] = None, + limit: int = 50, +) -> list[dict[str, Any]]: + if rels is None: + rels = ["IMPORTS", "CALLS", "DEFINES"] + # Aggregate across relations; preserve ordering and dedupe by id. + seen: set[Any] = set() + out: list[dict[str, Any]] = [] + for rel in rels: + rows = await _neighbors_payload(project, branch, symbol_id, rel, "OUT", limit) + for row in rows: + key = (row.get("id"), row.get("relation")) + if key in seen: + continue + seen.add(key) + out.append(row) + if len(out) >= limit: + return out + return out + + +# --------------------------------------------------------------------------- +# T7 — find_path +# --------------------------------------------------------------------------- + + +@app.tool( + name="find_path", + description=( + "Return up to `max_paths` CALLS-path sequences from `source_id` to " + "`dest_id`. Useful for 'how does A reach B' questions. Returns an " + "empty list when no path exists." + ), +) +async def find_path( + source_id: Any, + dest_id: Any, + project: str, + branch: Optional[str] = None, + max_paths: int = 10, +) -> list[dict[str, Any]]: + src = _coerce_node_id(source_id) + dst = _coerce_node_id(dest_id) + g = _project_arg(project, branch) + try: + raw = await g.find_paths(src, dst) + finally: + await g.close() + + # ``AsyncGraphQuery.find_paths`` returns each path as an alternating + # [node, edge, node, edge, ..., node] list; we strip edges and surface + # only the node sequence — that's what agents typically want. + paths: list[dict[str, Any]] = [] + for entry in raw[:max_paths]: + node_seq = [ + _node_summary(x) + for x in entry + # Edges in the alternating list carry a top-level ``relation`` + # key (from ``encode_edge``); nodes carry ``properties``. + if isinstance(x, dict) and "properties" in x + ] + paths.append({"path": node_seq}) + return paths + + +# --------------------------------------------------------------------------- +# T8 — search_code +# --------------------------------------------------------------------------- + + +@app.tool( + name="search_code", + description=( + "Prefix-search for symbols (functions, classes, files) whose name " + "starts with `prefix`. Backed by FalkorDB's full-text index. The " + "agent typically calls this first to discover symbol ids for the " + "navigation tools (`get_callers`, `find_path`, ...)." + ), +) +async def search_code( + prefix: str, + project: str, + branch: Optional[str] = None, + limit: int = 20, +) -> list[dict[str, Any]]: + g = _project_arg(project, branch) + try: + raw = await g.prefix_search(prefix) + finally: + await g.close() + return [_node_summary(node) for node in raw[:limit]] diff --git a/tests/mcp/test_query_tools.py b/tests/mcp/test_query_tools.py new file mode 100644 index 00000000..ca034cf8 --- /dev/null +++ b/tests/mcp/test_query_tools.py @@ -0,0 +1,253 @@ +"""T5/T7/T8 — query MCP tools tests. + +Bundled because all three tools are thin async wrappers around existing +``AsyncGraphQuery`` operations and share the same fixture. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + + +pytestmark = pytest.mark.anyio + + +@pytest.fixture +def anyio_backend() -> str: + return "asyncio" + + +# --------------------------------------------------------------------------- +# search_code (T8) — runs first because callers/find_path need the ids it +# returns. +# --------------------------------------------------------------------------- + + +async def test_search_code_finds_entrypoint(indexed_fixture, expected_contract): + from api.mcp.tools.structural import search_code + + results = await search_code( + prefix="ent", + project=indexed_fixture.project, + branch=indexed_fixture.branch, + ) + names = {r["name"] for r in results} + for required in expected_contract["search_prefixes"]["ent"]["must_include"]: + assert required in names, f"expected {required} in {names}" + + +async def test_search_code_honors_limit(indexed_fixture): + from api.mcp.tools.structural import search_code + + results = await search_code( + prefix="r", # broad prefix + project=indexed_fixture.project, + branch=indexed_fixture.branch, + limit=1, + ) + assert len(results) <= 1 + + +async def test_search_code_empty_for_nonsense(indexed_fixture): + from api.mcp.tools.structural import search_code + + results = await search_code( + prefix="zzz_no_such_symbol_zzz", + project=indexed_fixture.project, + branch=indexed_fixture.branch, + ) + assert results == [] + + +async def test_search_code_result_serialisable(indexed_fixture): + from api.mcp.tools.structural import search_code + + results = await search_code( + prefix="serv", + project=indexed_fixture.project, + branch=indexed_fixture.branch, + ) + json.dumps(results) # must not raise + + +# --------------------------------------------------------------------------- +# get_callers / get_callees / get_dependencies (T5) +# --------------------------------------------------------------------------- + + +async def _find_id(indexed_fixture, name: str) -> int: + """Helper: resolve a symbol name to its int node id via search_code.""" + from api.mcp.tools.structural import search_code + + rows = await search_code( + prefix=name, + project=indexed_fixture.project, + branch=indexed_fixture.branch, + ) + for r in rows: + if r["name"] == name: + return r["id"] + raise AssertionError(f"symbol {name!r} not found via search_code") + + +async def test_get_callees_of_entrypoint(indexed_fixture, expected_contract): + from api.mcp.tools.structural import get_callees + + entry_id = await _find_id(indexed_fixture, "entrypoint") + callees = await get_callees( + symbol_id=entry_id, + project=indexed_fixture.project, + branch=indexed_fixture.branch, + ) + names = {c["name"] for c in callees} + expected = set(expected_contract["calls"]["entrypoint"]["callees_any_of"]) + assert names & expected, ( + f"entrypoint callees {names} disjoint from expected {expected}" + ) + + for c in callees: + assert c["relation"] == "CALLS" + assert c["direction"] == "OUT" + + +async def test_get_callers_of_service(indexed_fixture, expected_contract): + from api.mcp.tools.structural import get_callers + + service_id = await _find_id(indexed_fixture, "service") + callers = await get_callers( + symbol_id=service_id, + project=indexed_fixture.project, + branch=indexed_fixture.branch, + ) + names = {c["name"] for c in callers} + for required in expected_contract["calls"]["service"]["callers"]: + assert required in names, f"expected caller {required} in {names}" + + for c in callers: + assert c["relation"] == "CALLS" + assert c["direction"] == "IN" + + +async def test_get_dependencies_aggregates_relations(indexed_fixture): + from api.mcp.tools.structural import get_dependencies + + entry_id = await _find_id(indexed_fixture, "entrypoint") + deps = await get_dependencies( + symbol_id=entry_id, + project=indexed_fixture.project, + branch=indexed_fixture.branch, + ) + # Default relations include CALLS, IMPORTS, DEFINES — at minimum the + # CALLS edge to ``service`` must be present. + rels = {d["relation"] for d in deps} + assert "CALLS" in rels + + +async def test_neighbor_tools_accept_string_ids(indexed_fixture): + """Agents sometimes hand back ids as strings — must work.""" + from api.mcp.tools.structural import get_callees + + entry_id = await _find_id(indexed_fixture, "entrypoint") + callees = await get_callees( + symbol_id=str(entry_id), # ← string! + project=indexed_fixture.project, + branch=indexed_fixture.branch, + ) + assert isinstance(callees, list) + + +async def test_neighbor_tools_reject_garbage_ids(indexed_fixture): + from api.mcp.tools.structural import get_callers + + with pytest.raises(ValueError, match="symbol_id"): + await get_callers( + symbol_id="not-a-number", + project=indexed_fixture.project, + branch=indexed_fixture.branch, + ) + + +# --------------------------------------------------------------------------- +# find_path (T7) +# --------------------------------------------------------------------------- + + +async def test_find_path_entrypoint_to_db(indexed_fixture, expected_contract): + from api.mcp.tools.structural import find_path + + entry_id = await _find_id(indexed_fixture, "entrypoint") + db_id = await _find_id(indexed_fixture, "db") + + paths = await find_path( + source_id=entry_id, + dest_id=db_id, + project=indexed_fixture.project, + branch=indexed_fixture.branch, + ) + # The contract requires at least one path entrypoint -> ... -> db + expected_min = next( + p["min_paths"] for p in expected_contract["paths"] + if p["source"] == "entrypoint" and p["dest"] == "db" + ) + assert len(paths) >= expected_min + + # Each path must have entrypoint first, db last, in a non-empty node + # sequence. + for entry in paths: + seq = entry["path"] + assert len(seq) >= 2 + assert seq[0]["name"] == "entrypoint" + assert seq[-1]["name"] == "db" + + +async def test_find_path_no_path_returns_empty(indexed_fixture): + """db -> entrypoint has no CALLS path (graph is acyclic).""" + from api.mcp.tools.structural import find_path + + entry_id = await _find_id(indexed_fixture, "entrypoint") + db_id = await _find_id(indexed_fixture, "db") + + paths = await find_path( + source_id=db_id, + dest_id=entry_id, + project=indexed_fixture.project, + branch=indexed_fixture.branch, + ) + assert paths == [] + + +async def test_find_path_honors_max_paths(indexed_fixture): + from api.mcp.tools.structural import find_path + + entry_id = await _find_id(indexed_fixture, "entrypoint") + db_id = await _find_id(indexed_fixture, "db") + + paths = await find_path( + source_id=entry_id, + dest_id=db_id, + project=indexed_fixture.project, + branch=indexed_fixture.branch, + max_paths=1, + ) + assert len(paths) <= 1 + + +# --------------------------------------------------------------------------- +# Protocol registration +# --------------------------------------------------------------------------- + + +async def test_all_query_tools_registered(): + from api.mcp.server import app + + tools = {t.name for t in await app.list_tools()} + assert { + "search_code", + "get_callers", + "get_callees", + "get_dependencies", + "find_path", + }.issubset(tools) From 2df968b52e2669a3226792aa50f106df87e15482 Mon Sep 17 00:00:00 2001 From: Dvir Dukhan <12258836+DvirDukhan@users.noreply.github.com> Date: Mon, 8 Jun 2026 16:04:39 +0300 Subject: [PATCH 2/3] fix(mcp): discriminate find_path elements on labels, not properties find_path walks the alternating [node, edge, node, ...] list that AsyncGraphQuery.find_paths returns and was filtering elements with `"properties" in x` to keep only nodes. But FalkorDB's Edge also exposes a `properties` attribute, so edges slipped through and became bogus, all-null `_node_summary` entries in the returned path. Discriminate on `labels` instead: `encode_node` emits a top-level `labels` key while `encode_edge` does not (edges carry `relation`/`src_node`/`dest_node`). The find_path test now asserts every path element resolves a non-null name and label, which catches the edge-leak regression. Also fixed two pre-existing wrong contract keys in the same test file that broke the mcp-tests CI job (`min_paths` -> `paths_count`, `callees_any_of` -> `callees`) and removed an unused `pathlib.Path` import. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- api/mcp/tools/structural.py | 10 +++++++--- tests/mcp/test_query_tools.py | 11 ++++++++--- 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/api/mcp/tools/structural.py b/api/mcp/tools/structural.py index 59e03a05..9685d6b5 100644 --- a/api/mcp/tools/structural.py +++ b/api/mcp/tools/structural.py @@ -413,9 +413,13 @@ async def find_path( node_seq = [ _node_summary(x) for x in entry - # Edges in the alternating list carry a top-level ``relation`` - # key (from ``encode_edge``); nodes carry ``properties``. - if isinstance(x, dict) and "properties" in x + # Discriminate on ``labels``: ``encode_node`` emits a top-level + # ``labels`` key, while ``encode_edge`` does not (edges carry + # ``relation``/``src_node``/``dest_node`` instead). Filtering on + # ``properties`` would be wrong because FalkorDB's Edge also has a + # ``properties`` attribute, so edges would slip through as bogus + # all-null node entries. + if isinstance(x, dict) and "labels" in x ] paths.append({"path": node_seq}) return paths diff --git a/tests/mcp/test_query_tools.py b/tests/mcp/test_query_tools.py index ca034cf8..598f8182 100644 --- a/tests/mcp/test_query_tools.py +++ b/tests/mcp/test_query_tools.py @@ -7,7 +7,6 @@ from __future__ import annotations import json -from pathlib import Path import pytest @@ -103,7 +102,7 @@ async def test_get_callees_of_entrypoint(indexed_fixture, expected_contract): branch=indexed_fixture.branch, ) names = {c["name"] for c in callees} - expected = set(expected_contract["calls"]["entrypoint"]["callees_any_of"]) + expected = set(expected_contract["calls"]["entrypoint"]["callees"]) assert names & expected, ( f"entrypoint callees {names} disjoint from expected {expected}" ) @@ -189,7 +188,7 @@ async def test_find_path_entrypoint_to_db(indexed_fixture, expected_contract): ) # The contract requires at least one path entrypoint -> ... -> db expected_min = next( - p["min_paths"] for p in expected_contract["paths"] + p["paths_count"] for p in expected_contract["paths"] if p["source"] == "entrypoint" and p["dest"] == "db" ) assert len(paths) >= expected_min @@ -201,6 +200,12 @@ async def test_find_path_entrypoint_to_db(indexed_fixture, expected_contract): assert len(seq) >= 2 assert seq[0]["name"] == "entrypoint" assert seq[-1]["name"] == "db" + # Every element must be a real node, not an edge that leaked through + # the alternating [node, edge, node, ...] list as a bogus all-null + # entry. Real nodes always resolve a name and a label. + for node in seq: + assert node["name"] is not None, f"edge leaked into path: {node}" + assert node["label"] is not None, f"edge leaked into path: {node}" async def test_find_path_no_path_returns_empty(indexed_fixture): From b590259cda5f29e3f45033c4551796abd14994bb Mon Sep 17 00:00:00 2001 From: Dvir Dukhan <12258836+DvirDukhan@users.noreply.github.com> Date: Mon, 8 Jun 2026 16:16:52 +0300 Subject: [PATCH 3/3] =?UTF-8?q?fix(mcp):=20harden=20query=20tools=20?= =?UTF-8?q?=E2=80=94=20injection=20guard,=20limit=20pushdown,=20typed=20id?= =?UTF-8?q?s?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses Copilot review on the T5/T7/T8 query tools: - get_dependencies: validate agent-supplied relation types before they are string-interpolated into Cypher (`-[e:{rel}]->`). FalkorDB can't bind relationship types as parameters, so an unvalidated `rels` value was a Cypher-injection vector. Reject anything that isn't a bare identifier. - get_dependencies: request only the remaining limit per relation so total DB work is bounded by `limit`, not `limit * len(rels)`. - find_path: push `max_paths` down to the database (new optional `limit` on AsyncGraphQuery.find_paths) instead of enumerating every path and slicing in Python — avoids unbounded path enumeration on large graphs. - search_code: honor `limit` at the database. prefix_search hard-capped at 10 rows in Cypher and the wrapper sliced afterwards, so callers could never get more than 10. LIMIT is now parameterized (default 10 preserves existing HTTP/CLI behavior). - Annotate the public id parameters as `int | str` instead of `Any` so the generated MCP JSON schema reflects the real input contract. Adds a Cypher-injection regression test for get_dependencies. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- api/graph.py | 39 +++++++++++++++++--------- api/mcp/tools/structural.py | 52 ++++++++++++++++++++++++++++------- tests/mcp/test_query_tools.py | 15 ++++++++++ 3 files changed, 83 insertions(+), 23 deletions(-) diff --git a/api/graph.py b/api/graph.py index 68b1b335..ffc27a33 100644 --- a/api/graph.py +++ b/api/graph.py @@ -448,14 +448,15 @@ def get_function_by_name(self, name: str) -> Optional[Node]: return res[0][0] - def prefix_search(self, prefix: str) -> str: + def prefix_search(self, prefix: str, limit: int = 10) -> str: """ Search for entities by prefix using a full-text search on the graph. - The search is limited to 10 nodes. Each node's name and labels are retrieved, - and the results are sorted based on their labels. + The number of results is bounded by ``limit`` (default 10). Each node's + name and labels are retrieved, and the results are sorted based on their labels. Args: prefix (str): The prefix string to search for in the graph database. + limit (int): Maximum number of nodes to return (default 10). Returns: str: A list of entity names and corresponding labels, sorted by label. @@ -465,7 +466,7 @@ def prefix_search(self, prefix: str) -> str: # Append a wildcard '*' to the prefix for full-text search. search_prefix = f"{prefix}*" - # Cypher query to perform full-text search and limit the result to 10 nodes. + # Cypher query to perform full-text search, bounding the result at $limit. # The 'CALL db.idx.fulltext.queryNodes' method searches for nodes labeled 'Searchable' # that match the given prefix, collects the nodes, and returns the result. query = """ @@ -473,11 +474,11 @@ def prefix_search(self, prefix: str) -> str: YIELD node WITH node RETURN node - LIMIT 10 + LIMIT $limit """ # Execute the query using the provided graph database connection. - result_set = self._query(query, {'prefix': search_prefix}).result_set + result_set = self._query(query, {'prefix': search_prefix, 'limit': int(limit)}).result_set completions = [encode_node(row[0]) for row in result_set] @@ -658,13 +659,16 @@ def rerun_query(self, q: str, params: dict) -> QueryResult: return self._query(q, params) - def find_paths(self, src: int, dest: int) -> list[Path]: + def find_paths(self, src: int, dest: int, limit: Optional[int] = None) -> list[Path]: """ Find all paths between the source (src) and destination (dest) nodes. Args: src (int): The ID of the source node. dest (int): The ID of the destination node. + limit (Optional[int]): When provided, bound the number of paths + enumerated by the database with a Cypher ``LIMIT``. When ``None`` + (default) all paths are returned (legacy behavior). Returns: List[Optional[Path]]: A list of paths found between the src and dest nodes. @@ -682,8 +686,13 @@ def find_paths(self, src: int, dest: int) -> list[Path]: RETURN p """ + params = {'src_id': src, 'dest_id': dest} + if limit is not None: + q += " LIMIT $limit\n" + params['limit'] = int(limit) + # Perform the query with the source and destination node IDs. - result_set = self._query(q, {'src_id': src, 'dest_id': dest}).result_set + result_set = self._query(q, params).result_set paths = [] @@ -861,26 +870,30 @@ async def get_neighbors(self, node_ids: list[int], rel: Optional[str] = None, lb logging.error(f"Error fetching neighbors for node {node_ids}: {e}") return {'nodes': [], 'edges': []} - async def prefix_search(self, prefix: str) -> list: + async def prefix_search(self, prefix: str, limit: int = 10) -> list: search_prefix = f"{prefix}*" query = """ CALL db.idx.fulltext.queryNodes('Searchable', $prefix) YIELD node WITH node RETURN node - LIMIT 10 + LIMIT $limit """ - result_set = (await self._query(query, {'prefix': search_prefix})).result_set + result_set = (await self._query(query, {'prefix': search_prefix, 'limit': int(limit)})).result_set return [encode_node(row[0]) for row in result_set] - async def find_paths(self, src: int, dest: int) -> list: + async def find_paths(self, src: int, dest: int, limit: Optional[int] = None) -> list: q = """MATCH (src), (dest) WHERE ID(src) = $src_id AND ID(dest) = $dest_id WITH src, dest MATCH p = (src)-[:CALLS*]->(dest) RETURN p """ - result_set = (await self._query(q, {'src_id': src, 'dest_id': dest})).result_set + params = {'src_id': src, 'dest_id': dest} + if limit is not None: + q += " LIMIT $limit\n" + params['limit'] = int(limit) + result_set = (await self._query(q, params)).result_set paths = [] for row in result_set: path = [] diff --git a/api/mcp/tools/structural.py b/api/mcp/tools/structural.py index 9685d6b5..f26e1678 100644 --- a/api/mcp/tools/structural.py +++ b/api/mcp/tools/structural.py @@ -20,6 +20,7 @@ import asyncio import logging import os +import re from pathlib import Path from typing import Any, Optional @@ -250,6 +251,25 @@ def _node_summary(n: Any) -> dict[str, Any]: } +# Relationship-type names are graph labels (SCREAMING_SNAKE_CASE, e.g. CALLS, +# IMPORTS, DEFINES). FalkorDB cannot parameterize relationship types, so any +# ``rel`` interpolated into Cypher must be validated against this pattern to +# prevent Cypher injection via agent-controlled input. +_REL_NAME_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") + + +def _validate_relation(rel: str) -> str: + """Return ``rel`` if it is a safe relationship-type name, else raise. + + Guards the relationship types that are string-interpolated into Cypher + (``-[e:{rel}]->``) — parameter binding is not available for relation + types in FalkorDB. + """ + if not isinstance(rel, str) or not _REL_NAME_RE.match(rel): + raise ValueError(f"invalid relation type: {rel!r}") + return rel + + def _coerce_node_id(symbol_id: Any) -> int: """Accept int or stringified int; raise ValueError otherwise. @@ -281,6 +301,7 @@ async def _neighbors_payload( walks outgoing edges, so we inline the Cypher here for symmetry. """ node_id = _coerce_node_id(symbol_id) + rel = _validate_relation(rel) g = _project_arg(project, branch) try: if direction == "OUT": @@ -321,7 +342,7 @@ async def _neighbors_payload( ), ) async def get_callers( - symbol_id: Any, + symbol_id: int | str, project: str, branch: Optional[str] = None, limit: int = 50, @@ -336,7 +357,7 @@ async def get_callers( ), ) async def get_callees( - symbol_id: Any, + symbol_id: int | str, project: str, branch: Optional[str] = None, limit: int = 50, @@ -353,7 +374,7 @@ async def get_callees( ), ) async def get_dependencies( - symbol_id: Any, + symbol_id: int | str, project: str, branch: Optional[str] = None, rels: Optional[list[str]] = None, @@ -365,7 +386,14 @@ async def get_dependencies( seen: set[Any] = set() out: list[dict[str, Any]] = [] for rel in rels: - rows = await _neighbors_payload(project, branch, symbol_id, rel, "OUT", limit) + # Only fetch the rows we can still accept, so total DB work is + # bounded by ``limit`` rather than ``limit * len(rels)``. + remaining = limit - len(out) + if remaining <= 0: + break + rows = await _neighbors_payload( + project, branch, symbol_id, rel, "OUT", remaining + ) for row in rows: key = (row.get("id"), row.get("relation")) if key in seen: @@ -391,8 +419,8 @@ async def get_dependencies( ), ) async def find_path( - source_id: Any, - dest_id: Any, + source_id: int | str, + dest_id: int | str, project: str, branch: Optional[str] = None, max_paths: int = 10, @@ -401,7 +429,9 @@ async def find_path( dst = _coerce_node_id(dest_id) g = _project_arg(project, branch) try: - raw = await g.find_paths(src, dst) + # Bound DB work by ``max_paths`` so large graphs don't enumerate an + # unbounded number of paths before we slice in Python. + raw = await g.find_paths(src, dst, limit=max_paths) finally: await g.close() @@ -409,7 +439,7 @@ async def find_path( # [node, edge, node, edge, ..., node] list; we strip edges and surface # only the node sequence — that's what agents typically want. paths: list[dict[str, Any]] = [] - for entry in raw[:max_paths]: + for entry in raw: node_seq = [ _node_summary(x) for x in entry @@ -447,7 +477,9 @@ async def search_code( ) -> list[dict[str, Any]]: g = _project_arg(project, branch) try: - raw = await g.prefix_search(prefix) + # Push the caller's ``limit`` down to the DB so it is actually honored + # (the underlying full-text query is otherwise capped at its default). + raw = await g.prefix_search(prefix, limit=limit) finally: await g.close() - return [_node_summary(node) for node in raw[:limit]] + return [_node_summary(node) for node in raw] diff --git a/tests/mcp/test_query_tools.py b/tests/mcp/test_query_tools.py index 598f8182..c3d0981a 100644 --- a/tests/mcp/test_query_tools.py +++ b/tests/mcp/test_query_tools.py @@ -145,6 +145,21 @@ async def test_get_dependencies_aggregates_relations(indexed_fixture): assert "CALLS" in rels +async def test_get_dependencies_rejects_injected_relation(indexed_fixture): + """Relation types are string-interpolated into Cypher, so agent-supplied + values must be validated to prevent Cypher injection.""" + from api.mcp.tools.structural import get_dependencies + + entry_id = await _find_id(indexed_fixture, "entrypoint") + with pytest.raises(ValueError, match="invalid relation type"): + await get_dependencies( + symbol_id=entry_id, + project=indexed_fixture.project, + branch=indexed_fixture.branch, + rels=["CALLS]->() DETACH DELETE n //"], + ) + + async def test_neighbor_tools_accept_string_ids(indexed_fixture): """Agents sometimes hand back ids as strings — must work.""" from api.mcp.tools.structural import get_callees