graph build fix again

mcp_codebase.py

@@ -2085,6 +2085,9 @@ def build_indexes():
     graph = LocalGraph(root_path=str(CODEBASE_PATH))
     graph.clear()
 
+    # Track all created nodes for second-pass relationships
+    created_nodes = {}
+
     # Create graph nodes/edges from metadata with enhanced relationships
     for m in metadatas:
         fpath = m.get("file")
@@ -2095,12 +2098,14 @@ def build_indexes():
         component_name = m.get("component_name")
 
         # Skip empty/invalid nodes
-        if not name or node_type == "unknown":
+        if not name or node_type == "unknown" or not fpath:
             continue
 
         # Ensure file node exists
         file_node_id = f"file::{fpath}"
-        graph.add_node(file_node_id, type="File", path=fpath, lang=lang)
+        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}
 
         # Create node ID with language
         node_id = f"{lang}::{node_type}::{fpath}::{name}"
@@ -2113,19 +2118,30 @@ def build_indexes():
             "lang": lang  # Always include language
         }
 
+        # Store for second-pass processing
+        created_nodes[node_id] = base_attrs.copy()
+        created_nodes[node_id]["type"] = node_type
+
         # Create appropriate node based on type
         if node_type in ["function", "method", "constructor"]:
             # Function-like entities
             if node_type == "method" and "class" in m:
                 # Method belongs to a class - use compound ID
-                node_id = f"{lang}::{node_type}::{fpath}::{m['class']}.{name}"
+                class_name = m['class']
+                node_id = f"{lang}::{node_type}::{fpath}::{class_name}.{name}"
+
+                # Update created_nodes with new ID
+                created_nodes[node_id] = base_attrs.copy()
+                created_nodes[node_id]["type"] = node_type
+                created_nodes[node_id]["class_name"] = class_name
+
                 graph.add_node(node_id, **base_attrs, type=node_type,
-                            class_name=m.get("class"),
+                            class_name=class_name,
                             return_type=m.get("return_type"),
                             parameters=m.get("parameters", []))
 
                 # Connect method to its class
-                class_node_id = f"{lang}::class::{fpath}::{m['class']}"
+                class_node_id = f"{lang}::class::{fpath}::{class_name}"
                 graph.add_edge(class_node_id, node_id, "contains")
 
                 # Connect return type if available
@@ -2153,6 +2169,7 @@ def build_indexes():
                 graph.add_node(node_id, **base_attrs, type=node_type,
                             return_type=m.get("return_type"),
                             parameters=m.get("parameters", []))
+                # Connect function to file
                 graph.add_edge(file_node_id, node_id, "contains")
 
                 # Connect return type if available
@@ -2181,27 +2198,17 @@ def build_indexes():
             graph.add_node(node_id, **base_attrs, type=node_type,
                         extends=m.get("extends"),
                         implements=m.get("implements", []),
-                        fields=m.get("fields", []),  # Store fields as metadata, not separate nodes
+                        fields=m.get("fields", []),
-                        variants=m.get("variants", []))  # For enums
+                        variants=m.get("variants", []))
 
+            # Connect type to file
             graph.add_edge(file_node_id, node_id, "contains")
 
-            # Inheritance relationships
+            # Store additional attributes for second-pass
-            extends_class = m.get("extends")
+            created_nodes[node_id]["extends"] = m.get("extends")
-            if extends_class:
+            created_nodes[node_id]["implements"] = m.get("implements", [])
-                for parent_id, parent_data in graph.find_nodes(name=extends_class):
+            created_nodes[node_id]["fields"] = m.get("fields", [])
-                    graph.add_edge(node_id, parent_id, "extends")
+            created_nodes[node_id]["variants"] = m.get("variants", [])
-
-            # Interface implementation
-            implements_interfaces = m.get("implements", [])
-            for interface_name in implements_interfaces:
-                for interface_id, interface_data in graph.find_nodes(name=interface_name):
-                    graph.add_edge(node_id, interface_id, "implements")
-
-            # Bevy-specific: Plugin implementation detection
-            if node_type == "struct" and name.endswith("Plugin"):
-                # FIX: Use add_node with attributes instead of modifying existing node
-                graph.add_node(node_id, **base_attrs, type=node_type, is_bevy_plugin=True)
 
         elif node_type == "impl":
             # Implementation blocks
@@ -2213,21 +2220,15 @@ def build_indexes():
                         traits=m.get("traits", []),
                         methods=m.get("methods", []))
 
+            # Connect impl to file
             graph.add_edge(file_node_id, node_id, "contains")
 
-            # Connect to target type
+            # Store for second-pass
-            for target_id, target_data in graph.find_nodes(name=target):
+            created_nodes[node_id] = base_attrs.copy()
-                graph.add_edge(node_id, target_id, "implements_for")
+            created_nodes[node_id]["type"] = node_type
-
+            created_nodes[node_id]["target"] = target
-            # Connect to implemented traits
+            created_nodes[node_id]["traits"] = m.get("traits", [])
-            for trait_name in m.get("traits", []):
+            created_nodes[node_id]["methods"] = m.get("methods", [])
-                for trait_id, trait_data in graph.find_nodes(name=trait_name):
-                    graph.add_edge(node_id, trait_id, "implements")
-
-            # Bevy-specific: Plugin detection
-            if target.endswith("Plugin"):
-                # FIX: Use add_node with attributes instead of modifying existing node
-                graph.add_node(node_id, **base_attrs, type=node_type, is_bevy_plugin_impl=True)
 
         elif node_type == "trait":
             # Traits
@@ -2235,6 +2236,8 @@ def build_indexes():
                         methods=m.get("methods", []))
             graph.add_edge(file_node_id, node_id, "contains")
 
+            created_nodes[node_id]["methods"] = m.get("methods", [])
+
         elif node_type == "module":
             # Modules
             graph.add_node(node_id, **base_attrs, type=node_type)
@@ -2247,104 +2250,30 @@ def build_indexes():
 
         # Handle imports if present
         for imp in m.get("imports", []):
+            # Create import node if it doesn't exist
+            if imp not in graph.graph.nodes:
+                graph.add_node(imp, type="Import", name=imp)
             graph.add_edge(file_node_id, imp, "imports")
 
         # Handle function calls if present
         for call in m.get("calls", []):
+            # Create call target node if it doesn't exist
+            if call not in graph.graph.nodes:
+                graph.add_node(call, type="FunctionCall", name=call)
             graph.add_edge(node_id, call, "calls")
 
-        # Rust-specific relationships - FIX: Use graph.graph.nodes
+    # -------------------------------------------------
-        if lang == "rust":
+    # Step 2: Second-pass relationship building
-            if node_type == "struct":
+    # -------------------------------------------------
-                # Connect struct fields as metadata, not separate nodes
+    logger.info("Building second-pass relationships...")
-                fields = m.get("fields", [])
-                for field in fields:
-                    # Store field info as node attribute instead of creating separate nodes
-                    if "field_types" not in graph.graph.nodes[node_id]:
-                        graph.graph.nodes[node_id]["field_types"] = []
-                    graph.graph.nodes[node_id]["field_types"].append(field)
 
-            elif node_type == "enum":
+    # Build type relationships
-                # Connect enum variants as metadata
+    type_connections = build_type_relationships(graph, created_nodes)
-                variants = m.get("variants", [])
+    logger.info(f"Added {type_connections} type relationships")
-                for variant in variants:
-                    if "variants" not in graph.graph.nodes[node_id]:
-                        graph.graph.nodes[node_id]["variants"] = []
-                    graph.graph.nodes[node_id]["variants"].append(variant)
 
-            elif node_type == "function":
+    # Build Bevy-specific relationships
-                # Connect function parameters to their types
+    bevy_connections = build_bevy_relationships(graph, created_nodes)
-                parameters = m.get("parameters", [])
+    logger.info(f"Added {bevy_connections} Bevy-specific relationships")
-                for param in parameters:
-                    # Extract type from "Type param_name" format
-                    param_parts = param.split()
-                    if len(param_parts) >= 2:
-                        param_type = param_parts[0]
-                        if param_type not in ['self', '&self', '&mut self', 'mut']:
-                            for type_id, type_data in graph.find_nodes(name=param_type):
-                                graph.add_edge(node_id, type_id, "uses_parameter")
-
-        # Bevy-specific relationship detection - FIX: Use graph.graph.nodes
-        if lang == "rust" and node_type == "function":
-            # Detect Bevy system functions
-            function_name = name.lower()
-            if any(keyword in function_name for keyword in ['system', 'plugin', 'build', 'update']):
-                graph.graph.nodes[node_id]["is_bevy_system"] = True
-
-            # Detect query parameters for system dependencies
-            parameters = m.get("parameters", [])
-            for param in parameters:
-                if 'Query' in param:
-                    graph.graph.nodes[node_id]["has_bevy_query"] = True
-                if 'Res<' in param or 'ResMut<' in param:
-                    graph.graph.nodes[node_id]["has_bevy_resource"] = True
-                if 'Commands' in param:
-                    graph.graph.nodes[node_id]["has_bevy_commands"] = True
-
-    # Second pass for Bevy-specific relationships - FIX: Use graph.graph.nodes
-    bevy_connections = 0
-
-    for node_id, node_data in graph.graph.nodes(data=True):
-        if node_data.get('lang') == 'rust':
-            # Connect Plugin impls to their build functions
-            if node_data.get('type') == 'impl' and node_data.get('is_bevy_plugin_impl'):
-                plugin_struct = node_data.get('target')
-                if plugin_struct:
-                    # Find the build method in this impl
-                    for method_id, method_data in graph.find_nodes(type="method", class_name=plugin_struct):
-                        if method_data.get('name') == 'build':
-                            graph.add_edge(node_id, method_id, "defines_plugin")
-                            bevy_connections += 1
-
-            # Connect systems to their plugins
-            elif node_data.get('is_bevy_system'):
-                # Try to find which plugin this system belongs to
-                file_path = node_data.get('file', '')
-                if file_path:
-                    # Look for plugins in the same file
-                    for plugin_id, plugin_data in graph.find_nodes(type="struct", file=file_path):
-                        if plugin_data.get('name', '').endswith('Plugin'):
-                            graph.add_edge(plugin_id, node_id, "contains_system")
-                            bevy_connections += 1
-
-    if bevy_connections > 0:
-        logger.info(f"Added {bevy_connections} Bevy-specific relationships")
-
-    logger.info("Building method parameter type relationships...")
-    param_connections = 0
-    for node_id, node_data in graph.graph.nodes(data=True):
-        if node_data.get('type') in ['method', 'function', 'constructor']:
-            parameters = node_data.get('parameters', [])
-            for param in parameters:
-                # Extract type from parameter string "TypeName paramName"
-                param_type = param.split()[0] if param and ' ' in param else None
-                if param_type and param_type not in ['void', 'int', 'String', 'boolean', 'long', 'double', 'float']:
-                    # Try to find this type in the graph
-                    for type_id, type_data in graph.find_nodes(name=param_type):
-                        graph.add_edge(node_id, type_id, "uses_parameter")
-                        param_connections += 1
-
-    logger.info(f"Added {param_connections} parameter type relationships")
 
     graph.save()
     graph.to_json()
@@ -2419,6 +2348,206 @@ def build_indexes():
     elapsed = index_build_time - start_time
     logger.info(f"Index build complete with contextual weighting applied. Total time: {elapsed:.1f}s")
 
+def build_type_relationships(graph, created_nodes):
+    """Build type relationships between nodes (parameter types, return types, etc.)"""
+    connections = 0
+
+    for node_id, node_data in created_nodes.items():
+        node_type = node_data.get("type")
+
+        # Handle function/method parameter and return types
+        if node_type in ["function", "method", "constructor"]:
+            connections += _connect_function_types(graph, node_id, node_data)
+
+        # Handle class/struct inheritance and implementation
+        elif node_type in ["class", "struct", "interface"]:
+            connections += _connect_class_relationships(graph, node_id, node_data)
+
+        # Handle impl blocks
+        elif node_type == "impl":
+            connections += _connect_impl_relationships(graph, node_id, node_data)
+
+    return connections
+
+def _connect_function_types(graph, node_id, node_data):
+    """Connect function nodes to their parameter and return types"""
+    connections = 0
+
+    # Connect return type
+    return_type = node_data.get("return_type")
+    if return_type and _is_meaningful_type(return_type):
+        return_type_clean = _clean_type_name(return_type)
+        for type_id, type_data in graph.find_nodes(name=return_type_clean):
+            graph.add_edge(node_id, type_id, "returns")
+            connections += 1
+
+    # Connect parameter types
+    parameters = node_data.get("parameters", [])
+    for param in parameters:
+        param_type = _extract_parameter_type(param)
+        if param_type and _is_meaningful_type(param_type):
+            param_type_clean = _clean_type_name(param_type)
+            for type_id, type_data in graph.find_nodes(name=param_type_clean):
+                graph.add_edge(node_id, type_id, "uses_parameter")
+                connections += 1
+
+    return connections
+
+def _connect_class_relationships(graph, node_id, node_data):
+    """Connect class/struct nodes to parent classes and interfaces"""
+    connections = 0
+
+    # Inheritance
+    extends = node_data.get("extends")
+    if extends and _is_meaningful_type(extends):
+        for parent_id, parent_data in graph.find_nodes(name=extends):
+            graph.add_edge(node_id, parent_id, "extends")
+            connections += 1
+
+    # Interface implementation
+    implements = node_data.get("implements", [])
+    for interface in implements:
+        if _is_meaningful_type(interface):
+            for interface_id, interface_data in graph.find_nodes(name=interface):
+                graph.add_edge(node_id, interface_id, "implements")
+                connections += 1
+
+    return connections
+
+def _connect_impl_relationships(graph, node_id, node_data):
+    """Connect impl blocks to their targets and traits"""
+    connections = 0
+
+    target = node_data.get("target")
+    if target and _is_meaningful_type(target):
+        # Connect to target type
+        for target_id, target_data in graph.find_nodes(name=target):
+            graph.add_edge(node_id, target_id, "implements_for")
+            connections += 1
+
+        # Connect to implemented traits
+        traits = node_data.get("traits", [])
+        for trait in traits:
+            if _is_meaningful_type(trait):
+                for trait_id, trait_data in graph.find_nodes(name=trait):
+                    graph.add_edge(node_id, trait_id, "implements")
+                    connections += 1
+
+    return connections
+
+def build_bevy_relationships(graph, created_nodes):
+    """Build Bevy-specific relationships"""
+    connections = 0
+
+    # Find all plugin structs and impls
+    plugin_structs = {}
+    plugin_impls = {}
+
+    for node_id, node_data in created_nodes.items():
+        node_type = node_data.get("type")
+        name = node_data.get("name", "")
+
+        # Collect plugin structs
+        if node_type in ["struct", "class"] and name.endswith("Plugin"):
+            plugin_structs[name] = node_id
+
+        # Collect plugin impls
+        elif node_type == "impl" and node_data.get("target", "").endswith("Plugin"):
+            plugin_impls[node_data["target"]] = node_id
+
+    # Connect plugin impls to their structs
+    for plugin_name, impl_id in plugin_impls.items():
+        if plugin_name in plugin_structs:
+            struct_id = plugin_structs[plugin_name]
+            graph.add_edge(impl_id, struct_id, "implements_plugin")
+            connections += 1
+
+            # Find build method in this impl
+            impl_methods = created_nodes[impl_id].get("methods", [])
+            for method_name in impl_methods:
+                if method_name == "build":
+                    # Find the method node
+                    method_id = f"{created_nodes[impl_id]['lang']}::method::{created_nodes[impl_id]['file']}::{plugin_name}.{method_name}"
+                    if method_id in created_nodes:
+                        graph.add_edge(impl_id, method_id, "defines_plugin")
+                        connections += 1
+
+    # Connect systems to plugins (simplified - same file assumption)
+    for node_id, node_data in created_nodes.items():
+        if node_data.get("type") in ["function", "method"]:
+            # Simple heuristic for Bevy systems
+            function_name = node_data.get("name", "").lower()
+            if any(keyword in function_name for keyword in ['system', 'update', 'setup']):
+                file_path = node_data.get("file")
+                # Look for plugins in the same file
+                for plugin_id, plugin_data in created_nodes.items():
+                    if (plugin_data.get("file") == file_path and
+                        plugin_data.get("type") in ["struct", "class"] and
+                        plugin_data.get("name", "").endswith("Plugin")):
+                        graph.add_edge(plugin_id, node_id, "contains_system")
+                        connections += 1
+
+    return connections
+
+def _extract_parameter_type(param):
+    """Extract type from parameter string, handling multiple formats"""
+    if not param:
+        return None
+
+    # Handle "type: name" format (TypeScript/Python)
+    if ':' in param:
+        parts = param.split(':')
+        if len(parts) >= 2:
+            return parts[1].strip()
+
+    # Handle "Type name" format (Rust/Java/C++)
+    if ' ' in param:
+        parts = param.split()
+        if len(parts) >= 2:
+            return parts[0].strip()
+
+    # Handle complex types with generics
+    if '<' in param and '>' in param:
+        # Extract the base type before generics
+        base_type = param.split('<')[0].strip()
+        return base_type
+
+    return param.strip()
+
+def _clean_type_name(type_name):
+    """Clean type names by removing common modifiers"""
+    if not type_name:
+        return ""
+
+    # Remove common modifiers
+    modifiers = ['&', 'mut ', 'const ', 'static ', 'pub ', 'private ', 'protected ']
+    cleaned = type_name
+    for mod in modifiers:
+        cleaned = cleaned.replace(mod, '')
+
+    # Remove trailing references/pointers
+    cleaned = cleaned.rstrip('*&')
+
+    return cleaned.strip()
+
+def _is_meaningful_type(type_name):
+    """Check if a type name is meaningful (not primitive, void, etc.)"""
+    if not type_name:
+        return False
+
+    meaningless_types = [
+        'void', 'int', 'i32', 'i64', 'f32', 'f64', 'bool', 'str', 'String',
+        'char', 'u8', 'u16', 'u32', 'u64', 'usize', 'isize',
+        'self', '&self', '&mut self', 'mut self', 'Self',
+        'Result', 'Option', 'Vec', 'String'
+    ]
+
+    cleaned = _clean_type_name(type_name)
+    return (cleaned and
+            cleaned not in meaningless_types and
+            not cleaned.startswith('&') and
+            len(cleaned) > 1)
+
 # -----------------------------
 # Build / load indexes with contextual weighting and batch embedding
 # -----------------------------