From fd733cd4e40553c75760404997c46ec85ab396ce Mon Sep 17 00:00:00 2001 From: Stephen Adamson Date: Sun, 16 Nov 2025 00:26:25 +0000 Subject: [PATCH] fix rust --- mcp_codebase.py | 201 +++++++++++++++++++++++++++--------------------- 1 file changed, 112 insertions(+), 89 deletions(-) diff --git a/mcp_codebase.py b/mcp_codebase.py index bc99f59..39bee1d 100644 --- a/mcp_codebase.py +++ b/mcp_codebase.py @@ -2505,7 +2505,7 @@ def parse_rust_file(filepath: Path) -> List[Dict]: @lru_cache(maxsize=1000) def parse_rust_file_cached(filepath_str: str, file_hash: str) -> tuple: - """Enhanced Rust parsing with better metadata extraction.""" + """Fixed Rust parsing that properly handles imports as metadata, not separate entities.""" filepath = Path(filepath_str) chunks = [] @@ -2513,7 +2513,7 @@ def parse_rust_file_cached(filepath_str: str, file_hash: str) -> tuple: rust_helper = Path("./tools/target/release/parse_rust_ast") if not rust_helper.is_file(): - logger.warning(f"Rust parser binary not found: {rust_helper}. Using fallback whole-file chunk for {filepath}") + logger.warning(f"Rust parser binary not found: {rust_helper}. Using fallback.") return create_fallback_chunk(filepath) try: @@ -2530,8 +2530,20 @@ def parse_rust_file_cached(filepath_str: str, file_hash: str) -> tuple: decls = json.loads(proc.stdout) logger.info(f" AST helper found {len(decls)} declarations") - for i, d in enumerate(decls): - # FIX: Ensure valid line numbers + # Collect imports first to attach to actual code entities + file_imports = [] + code_decls = [] + + for d in decls: + if d.get("type_") == "import": + file_imports.extend(d.get("parameters", [])) + else: + code_decls.append(d) + + logger.info(f" Found {len(file_imports)} imports and {len(code_decls)} code declarations") + + # Process actual code entities (functions, structs, etc.) + for i, d in enumerate(code_decls): start_line = max(1, d.get("start_line", 1)) end_line = max(start_line, d.get("end_line", start_line + 1)) @@ -2549,28 +2561,34 @@ def parse_rust_file_cached(filepath_str: str, file_hash: str) -> tuple: "type": d.get("type_", "") or "unknown", "line": start_line, "language": "rust", + "imports": file_imports, # Attach imports to ALL code entities in this file } - # Enhanced metadata extraction + # Enhanced metadata extraction for code entities enhanced_fields = [ "visibility", "is_async", "is_unsafe", "generics", "traits", "fields", "methods", "return_type", "parameters", "docstring", - "imports", "calls" # NEW fields + "calls" # calls from the specific entity ] for field in enhanced_fields: if field in d and d[field] is not None: metadata[field] = d[field] - logger.info(f" Declaration {i+1}: {metadata.get('type', 'unknown')} {metadata.get('name', 'unnamed')}") + logger.info(f" Code {i+1}: {metadata.get('type', 'unknown')} {metadata.get('name', 'unnamed')}") - # Build rich chunk text - chunk_text = build_rust_chunk_text(metadata, chunk_code) + # Build rich chunk text for CODE entities + chunk_text = build_rust_code_chunk_text(metadata, chunk_code) chunks.append({ "text": chunk_text, "metadata": metadata }) + # If we have no code declarations but have imports, create a file-level chunk + if not chunks and file_imports: + logger.info(" Creating file-level chunk for imports-only file") + chunks.append(create_file_level_chunk(filepath, file_imports)) + except Exception as e: logger.warning(f"Rust AST helper failed for {filepath}: {e}") return create_fallback_chunk(filepath) @@ -2582,15 +2600,15 @@ def parse_rust_file_cached(filepath_str: str, file_hash: str) -> tuple: logger.warning(f"❌ No chunks created from Rust AST for {filepath}. Emitting fallback.") return create_fallback_chunk(filepath) -def build_rust_chunk_text(metadata: dict, code: str) -> str: - """Build comprehensive Rust chunk text with all metadata.""" +def build_rust_code_chunk_text(metadata: dict, code: str) -> str: + """Build comprehensive Rust code chunk text.""" parts = [ f"File: {metadata.get('file', '')}", f"Type: {metadata.get('type', 'unknown')}", f"Name: {metadata.get('name', '')}", ] - # Add enhanced metadata + # Add enhanced metadata for CODE entities if metadata.get("visibility"): parts.append(f"Visibility: {metadata['visibility']}") if metadata.get("is_async"): @@ -2607,10 +2625,32 @@ def build_rust_chunk_text(metadata: dict, code: str) -> str: parts.append(f"Imports: {', '.join(metadata['imports'])}") if metadata.get("calls"): parts.append(f"Calls: {', '.join(metadata['calls'])}") + if metadata.get("docstring"): + parts.append(f"Doc: {metadata['docstring']}") parts.append(f"Code:\n{code}") return "\n".join(parts) +def create_file_level_chunk(filepath: Path, imports: list) -> dict: + """Create a chunk for files that only have imports.""" + try: + full_code = filepath.read_text(encoding="utf-8") + except Exception: + full_code = "" + + metadata = { + "file": str(filepath), + "name": filepath.stem, + "type": "module", # Treat import-only files as modules + "line": 1, + "language": "rust", + "imports": imports, + } + + chunk_text = f"File: {filepath}\nType: module\nName: {filepath.stem}\nImports: {', '.join(imports)}\nCode:\n{full_code}" + + return {"text": chunk_text, "metadata": metadata} + def create_fallback_chunk(filepath: Path) -> tuple: """Create a minimal fallback chunk with proper metadata.""" try: @@ -3533,97 +3573,80 @@ def build_indexes(): # ------------------------------------------------- def _add_import_edges_enhanced(graph, file_node_id, importing_node_id, imports, lang): - """Enhanced import edge creation. Returns dict of new nodes to add.""" + """Enhanced import edge creation that resolves Rust paths to actual files.""" new_nodes = {} for imp_path in imports: if not imp_path or not isinstance(imp_path, str): continue - # Normalize path separators based on language - if lang == "python": - normalized_path = imp_path.replace('.', '::') - elif lang in ["javascript", "typescript", "svelte"]: - if imp_path.startswith('.') and len(imp_path) < 3: - continue - normalized_path = imp_path.replace('/', '::').replace('@', '') - elif lang == "rust": - normalized_path = imp_path - elif lang == "go": - normalized_path = imp_path.replace('/', '::') - elif lang == "java": - normalized_path = imp_path.replace('.', '::') - else: - normalized_path = imp_path.replace('.', '::').replace('/', '::') - - parts = [p.strip() for p in normalized_path.split('::') if p.strip()] - if not parts: + # Skip standard library and external crates for now + if any(imp_path.startswith(prefix) for prefix in ['std::', 'core::', 'alloc::', 'bevy::', 'serde::']): continue - # Create package/crate/module hierarchy - package_name = parts[0] - package_node_id = f"package::{lang}::{package_name}" - - if package_node_id not in graph.graph.nodes: - graph.add_node( - package_node_id, - type=package_name, - name=package_name, - file=imp_path, - lang=lang, - is_external=True - ) - new_nodes[package_node_id] = { - "type": package_name, - "name": package_name, - "file": imp_path, - "lang": lang - } - - if not graph.graph.has_edge(file_node_id, package_node_id): - graph.add_edge(file_node_id, package_node_id, "imports") - - # Build hierarchical module structure - current_path = package_name - parent_node = package_node_id - - for i, part in enumerate(parts[1:], 1): - current_path = f"{current_path}::{part}" - - if i == len(parts) - 1: - node_type_name = part - else: - node_type_name = part - - module_node_id = f"symbol::{lang}::{current_path}" - - if module_node_id not in graph.graph.nodes: + # Try to resolve relative paths within the codebase + resolved_path = resolve_rust_import(imp_path, graph.root_path) + if resolved_path: + # Create edge to the actual file + target_file_node_id = f"file::{resolved_path}" + if not graph.graph.has_edge(importing_node_id, target_file_node_id): + graph.add_edge(importing_node_id, target_file_node_id, "imports") + else: + # Create symbolic node for external/unresolved imports + symbolic_id = f"rust::external::{imp_path}" + if symbolic_id not in graph.graph.nodes: graph.add_node( - module_node_id, - type=node_type_name, - name=current_path, - file=imp_path, - lang=lang, - is_symbol=True + symbolic_id, + type="external", + name=imp_path, + file="", + lang="rust", + is_external=True ) - new_nodes[module_node_id] = { - "type": node_type_name, - "name": current_path, - "file": imp_path, - "lang": lang + new_nodes[symbolic_id] = { + "type": "external", + "name": imp_path, + "file": "", + "lang": "rust" } - - if not graph.graph.has_edge(parent_node, module_node_id): - graph.add_edge(parent_node, module_node_id, "contains") - - parent_node = module_node_id - - if parent_node != package_node_id: - if not graph.graph.has_edge(importing_node_id, parent_node): - graph.add_edge(importing_node_id, parent_node, "imports") + if not graph.graph.has_edge(importing_node_id, symbolic_id): + graph.add_edge(importing_node_id, symbolic_id, "imports") return new_nodes +def resolve_rust_import(import_path: str, root_path: str) -> Optional[str]: + """Try to resolve a Rust import path to an actual file path.""" + # Remove any trailing ::* or other modifiers + clean_path = import_path.replace('::*', '').replace('*', '') + + # Convert :: to path separators + potential_path = clean_path.replace('::', '/') + + # Common Rust file extensions to try + extensions = ['.rs', ''] + + # Common source directories in Rust projects + base_dirs = ['src', 'crates', 'workspace'] + + for base_dir in base_dirs: + for ext in extensions: + # Try as direct file + candidate = Path(root_path) / base_dir / f"{potential_path}{ext}" + if candidate.exists(): + return str(candidate) + + # Try with mod.rs in directory + candidate = Path(root_path) / base_dir / potential_path / f"mod{ext}" + if candidate.exists(): + return str(candidate) + + # Try with lib.rs in directory + candidate = Path(root_path) / base_dir / potential_path / f"lib{ext}" + if candidate.exists(): + return str(candidate) + + return None + def _clean_type_string(type_str: str) -> str: """Clean up type strings for matching across all languages.""" if not type_str: