diff --git a/mcp_codebase.py b/mcp_codebase.py index 40bcd0a..053891c 100644 --- a/mcp_codebase.py +++ b/mcp_codebase.py @@ -2101,11 +2101,24 @@ def build_indexes(): if not name or node_type == "unknown" or not fpath: continue - # Ensure file node exists - file_node_id = f"file::{fpath}" + # ---- File node --------------------------------------------- + # Keep the same ID (`file::`) but store the full set of + # attributes so that downstream code can rely on `name`, `lang`, + # and `file` (even though `name` will just be the file path). if file_node_id not in graph.graph.nodes: - graph.add_node(file_node_id, type="File", path=fpath, lang=lang) - created_nodes[file_node_id] = {"type": "File", "path": fpath} + graph.add_node( + file_node_id, + type="File", + name=fpath, + path=fpath, + lang=lang, + ) + created_nodes[file_node_id] = { + "type": "File", + "name": fpath, + "path": fpath, + "lang": lang, + } # Create node ID with language node_id = f"{lang}::{node_type}::{fpath}::{name}" @@ -2260,8 +2273,29 @@ def build_indexes(): if not candidate_id: # If nothing exists, fall back to a generic import node. candidate_id = f"import::{imp}" + # ---- Import node -------------------------------------------- + # For imports we want a node that mirrors the target module if + # it exists; otherwise we still create a stub with language & + # file attributes for consistency. if candidate_id not in graph.graph.nodes: - graph.add_node(candidate_id, type="Import", name=imp) + # First try to find a module node with that name + found = False + for n_id, n_data in created_nodes.items(): + if n_data.get("name") == imp and n_data.get("type") == "module": + candidate_id = n_id + found = True + break + if not found: + # Generic import stub + candidate_id = f"import::{imp}" + graph.add_node( + candidate_id, + type="Import", + name=imp, + lang="unknown", + file="", + ) + graph.add_edge(file_node_id, candidate_id, "imports") graph.add_edge(file_node_id, candidate_id, "imports") @@ -2277,8 +2311,32 @@ def build_indexes(): if not target_id: # Fallback to a generic call node. target_id = f"call::{call}" + # ---- Call node --------------------------------------------- + # If we already know the target node (by name), reuse it. Otherwise + # create a generic node that still carries language & file context + # so it can be displayed or filtered later. if target_id not in graph.graph.nodes: - graph.add_node(target_id, type="FunctionCall", name=call) + # Attempt a best‑effort lookup for the actual function node + found = False + for n_id, n_data in created_nodes.items(): + if ( + n_data.get("name") == call + and n_data.get("type") in {"function", "method", "constructor"} + ): + target_id = n_id + found = True + break + if not found: + # Fallback to a generic FunctionCall node + target_id = f"call::{call}" + graph.add_node( + target_id, + type="FunctionCall", + name=call, + lang="unknown", + file="", + ) + graph.add_edge(node_id, target_id, "calls") graph.add_edge(node_id, target_id, "calls")