fix: apply pre-flatten passes to imported subgraphs

core/src/ten_rust/src/graph/mod.rs

@@ -526,28 +526,51 @@ impl Graph {
         Ok(Some(new_graph))
     }
 
-    /// Convenience method for flattening a graph instance without preserving
-    /// exposed info. This is the main public API for flattening graphs.
+    /// Applies the passes that must run before subgraph flattening:
     ///
-    /// Returns `Ok(None)` if the graph doesn't need flattening. Returns
-    /// `Ok(Some(flattened_graph))` if the graph was successfully flattened.
-    pub async fn flatten_graph(&self, current_base_dir: Option<&str>) -> Result<Option<Graph>> {
+    /// 1. Expand `names` arrays into individual `name` items.
+    /// 2. Replace selector references with the nodes they match.
+    /// 3. Convert reversed connections (`source`) to forward connections
+    ///    (`dest`).
+    ///
+    /// `flatten_subgraphs` relies on all three having already run - it expects
+    /// `name` to be `Some`, `names` to be `None`, and every node to be an
+    /// extension node. Any graph handed to it, including one imported through
+    /// `import_uri`, has to go through here first.
+    ///
+    /// Returns `Ok(None)` if none of the passes changed anything.
+    pub fn apply_pre_flatten_passes(&self) -> Result<Option<Graph>> {
         let mut processing_graph = self;
 
-        // Step 1: Expand names arrays to individual name items
         let expanded_names_graph = processing_graph.expand_names_to_individual_items()?;
         processing_graph = expanded_names_graph.as_ref().unwrap_or(processing_graph);
 
-        // Step 2: Match nodes according to selector rules and replace them in
-        // connections
         let flattened_selector_graph = processing_graph.flatten_selectors()?;
         processing_graph = flattened_selector_graph.as_ref().unwrap_or(processing_graph);
 
-        // Step 3: Convert reversed connections to forward connections if needed
         let reversed_graph =
             processing_graph.convert_reversed_connections_to_forward_connections()?;
         processing_graph = reversed_graph.as_ref().unwrap_or(processing_graph);
 
+        if std::ptr::eq(processing_graph, self) {
+            return Ok(None);
+        }
+
+        Ok(Some(processing_graph.clone()))
+    }
+
+    /// Convenience method for flattening a graph instance without preserving
+    /// exposed info. This is the main public API for flattening graphs.
+    ///
+    /// Returns `Ok(None)` if the graph doesn't need flattening. Returns
+    /// `Ok(Some(flattened_graph))` if the graph was successfully flattened.
+    pub async fn flatten_graph(&self, current_base_dir: Option<&str>) -> Result<Option<Graph>> {
+        let mut processing_graph = self;
+
+        // Steps 1-3: names expansion, selector resolution, reversed connections.
+        let pre_flattened_graph = processing_graph.apply_pre_flatten_passes()?;
+        processing_graph = pre_flattened_graph.as_ref().unwrap_or(processing_graph);
+
         // Step 4: Flatten subgraphs
         let flattened = Self::flatten_subgraphs(processing_graph, current_base_dir, false)
             .await

core/src/ten_rust/src/graph/subgraph/flatten.rs

@@ -534,6 +534,15 @@ impl Graph {
         let mut new_base_dir: Option<String> = None;
         let subgraph = load_graph_from_uri(import_uri, current_base_dir, &mut new_base_dir).await?;
 
+        // `load_graph_from_uri` only deserializes the file, so an imported
+        // subgraph has not been through the pre-flatten passes yet. Run them
+        // here, otherwise `names`, selectors and reversed connections written
+        // inside the subgraph reach the flattening logic unlowered.
+        let pre_flattened_subgraph = subgraph.apply_pre_flatten_passes().map_err(|e| {
+            anyhow::anyhow!("Failed to process subgraph '{}': {}", subgraph_node.name, e)
+        })?;
+        let subgraph = pre_flattened_subgraph.unwrap_or(subgraph);
+
         Self::process_loaded_subgraph(
             subgraph_node,
             &subgraph,

core/src/ten_rust/tests/test_case/graph/subgraph.rs

@@ -1339,4 +1339,171 @@ mod tests {
         assert_eq!(expanded_property.name, "config_b");
         assert!(expanded_property.subgraph.is_none());
     }
+
+    /// Builds a main graph whose only node imports `subgraph_path`.
+    fn main_graph_importing(subgraph_path: &str) -> Graph {
+        Graph {
+            nodes: vec![GraphNode::new_subgraph_node(
+                "sg".to_string(),
+                None,
+                GraphResource { import_uri: format!("file://{subgraph_path}") },
+            )],
+            connections: None,
+            exposed_messages: None,
+            exposed_properties: None,
+        }
+    }
+
+    /// A reversed connection (`source`) written inside an imported subgraph
+    /// must be converted to a forward connection, with the subgraph name
+    /// prefix applied, just like one written in the main graph.
+    #[tokio::test]
+    async fn test_flatten_subgraph_with_reversed_connection() {
+        let temp_dir = tempdir().unwrap();
+        let subgraph_file_path = temp_dir.path().join("subgraph_reversed.json");
+
+        // Inside the subgraph: ext_a --hello--> ext_b, written in reverse form.
+        fs::write(
+            &subgraph_file_path,
+            r#"{
+              "nodes": [
+                {"type": "extension", "name": "ext_a", "addon": "addon_a"},
+                {"type": "extension", "name": "ext_b", "addon": "addon_b"}
+              ],
+              "connections": [
+                {
+                  "extension": "ext_b",
+                  "cmd": [{"name": "hello", "source": [{"extension": "ext_a"}]}]
+                }
+              ]
+            }"#,
+        )
+        .unwrap();
+
+        let main_graph = main_graph_importing(subgraph_file_path.to_str().unwrap());
+        let flattened = main_graph.flatten_graph(None).await.unwrap().unwrap();
+
+        let connections = flattened.connections.as_ref().unwrap();
+
+        // No reversed connection may survive flattening.
+        assert!(
+            connections
+                .iter()
+                .flat_map(|conn| conn.cmd.iter().flatten())
+                .all(|flow| flow.source.is_empty()),
+            "reversed connection inside the subgraph was not converted: {connections:?}"
+        );
+
+        // It must have become sg_ext_a --hello--> sg_ext_b.
+        let forward = connections
+            .iter()
+            .find(|conn| conn.loc.extension.as_deref() == Some("sg_ext_a"))
+            .expect("no forward connection originating from sg_ext_a");
+
+        let flow = forward
+            .cmd
+            .as_ref()
+            .unwrap()
+            .iter()
+            .find(|flow| flow.name.as_deref() == Some("hello"))
+            .expect("no 'hello' cmd flow");
+
+        assert_eq!(flow.dest.len(), 1);
+        assert_eq!(flow.dest[0].loc.extension.as_deref(), Some("sg_ext_b"));
+    }
+
+    /// A `names` array written inside an imported subgraph must be expanded
+    /// into individual `name` items.
+    #[tokio::test]
+    async fn test_flatten_subgraph_with_names_array() {
+        let temp_dir = tempdir().unwrap();
+        let subgraph_file_path = temp_dir.path().join("subgraph_names.json");
+
+        fs::write(
+            &subgraph_file_path,
+            r#"{
+              "nodes": [
+                {"type": "extension", "name": "ext_a", "addon": "addon_a"},
+                {"type": "extension", "name": "ext_b", "addon": "addon_b"}
+              ],
+              "connections": [
+                {
+                  "extension": "ext_a",
+                  "cmd": [{"names": ["hello", "world"], "dest": [{"extension": "ext_b"}]}]
+                }
+              ]
+            }"#,
+        )
+        .unwrap();
+
+        let main_graph = main_graph_importing(subgraph_file_path.to_str().unwrap());
+        let flattened = main_graph.flatten_graph(None).await.unwrap().unwrap();
+
+        let connections = flattened.connections.as_ref().unwrap();
+        let flows: Vec<_> = connections.iter().flat_map(|conn| conn.cmd.iter().flatten()).collect();
+
+        assert!(
+            flows.iter().all(|flow| flow.names.is_none()),
+            "'names' inside the subgraph was not expanded: {flows:?}"
+        );
+
+        let mut names: Vec<&str> = flows.iter().filter_map(|flow| flow.name.as_deref()).collect();
+        names.sort_unstable();
+        assert_eq!(names, vec!["hello", "world"]);
+
+        for flow in &flows {
+            assert_eq!(flow.dest.len(), 1);
+            assert_eq!(flow.dest[0].loc.extension.as_deref(), Some("sg_ext_b"));
+        }
+    }
+
+    /// A selector node written inside an imported subgraph must be resolved to
+    /// the nodes it matches and removed, instead of reaching the flattening
+    /// logic and panicking as a non-extension node.
+    #[tokio::test]
+    async fn test_flatten_subgraph_with_selector() {
+        let temp_dir = tempdir().unwrap();
+        let subgraph_file_path = temp_dir.path().join("subgraph_selector.json");
+
+        fs::write(
+            &subgraph_file_path,
+            r#"{
+              "nodes": [
+                {"type": "extension", "name": "ext_a", "addon": "addon_a"},
+                {"type": "extension", "name": "ext_b", "addon": "addon_b"},
+                {
+                  "type": "selector",
+                  "name": "targets",
+                  "filter": {"field": "name", "operator": "exact", "value": "ext_b"}
+                }
+              ],
+              "connections": [
+                {
+                  "extension": "ext_a",
+                  "cmd": [{"name": "hello", "dest": [{"selector": "targets"}]}]
+                }
+              ]
+            }"#,
+        )
+        .unwrap();
+
+        let main_graph = main_graph_importing(subgraph_file_path.to_str().unwrap());
+        let flattened = main_graph.flatten_graph(None).await.unwrap().unwrap();
+
+        // The selector node must not survive into the flattened graph.
+        assert!(flattened.nodes.iter().all(|node| node.get_type() == GraphNodeType::Extension));
+        assert_eq!(flattened.nodes.len(), 2);
+
+        let connections = flattened.connections.as_ref().unwrap();
+        let flow = connections
+            .iter()
+            .find(|conn| conn.loc.extension.as_deref() == Some("sg_ext_a"))
+            .and_then(|conn| conn.cmd.as_ref())
+            .and_then(|flows| flows.iter().find(|flow| flow.name.as_deref() == Some("hello")))
+            .expect("no 'hello' cmd flow from sg_ext_a");
+
+        assert_eq!(flow.dest.len(), 1);
+        assert_eq!(flow.dest[0].loc.extension.as_deref(), Some("sg_ext_b"));
+        assert!(flow.dest[0].loc.selector.is_none());
+    }
 }