mocra_core/common/model/
module_dag.rs

1//! Module DAG definition types — the graph a [`ModuleTrait`](crate::common::interface::ModuleTrait)
2//! describes via `dag_definition()` / `add_step()`.
3//!
4//! These live in `common` (not `engine`) so the `ModuleTrait` contract does not depend on the
5//! engine — the engine's `ModuleDagProcessor` consumes them. Node-execution policy / placement
6//! types come from the standalone `mocra-dag` crate.
7
8use std::collections::HashMap;
9use std::sync::Arc;
10
11use indexmap::IndexMap;
12use mocra_dag::{DagNodeExecutionPolicy, NodePlacement};
13
14use crate::common::interface::ModuleNodeTrait;
15
16#[derive(Clone)]
17pub struct ModuleDagNodeDef {
18    pub node_id: String,
19    pub node: Arc<dyn ModuleNodeTrait>,
20    pub placement_override: Option<NodePlacement>,
21    pub policy_override: Option<DagNodeExecutionPolicy>,
22    pub tags: Vec<String>,
23}
24impl ModuleDagNodeDef {
25    pub fn new(node: Arc<dyn ModuleNodeTrait>) -> Self {
26        let key = node.stable_node_key();
27        let node_id = if key.is_empty() {
28            uuid::Uuid::now_v7().to_string()
29        } else {
30            key.to_string()
31        };
32        Self {
33            node_id,
34            node,
35            placement_override: None,
36            policy_override: None,
37            tags: Vec::new(),
38        }
39    }
40
41    /// Overrides the auto-generated node ID with an explicit stable identifier.
42    ///
43    /// Use this when you need multiple instances of the same node type in one DAG, or
44    /// when you cannot implement `stable_node_key()` on the node struct directly.
45    /// The ID must be unique within the DAG.
46    pub fn with_id(mut self, id: impl Into<String>) -> Self {
47        self.node_id = id.into();
48        self
49    }
50}
51
52#[derive(Debug, Clone)]
53pub struct ModuleDagEdgeDef {
54    pub from: String,
55    pub to: String,
56}
57impl ModuleDagEdgeDef {
58    pub fn new(from: &ModuleDagNodeDef, to: &ModuleDagNodeDef) -> Self {
59        Self {
60            from: from.node_id.clone(),
61            to: to.node_id.clone(),
62        }
63    }
64}
65
66#[derive(Clone, Default)]
67pub struct ModuleDagDefinition {
68    pub nodes: Vec<ModuleDagNodeDef>,
69    pub edges: Vec<ModuleDagEdgeDef>,
70    pub entry_nodes: Vec<String>,
71    pub default_policy: Option<DagNodeExecutionPolicy>,
72    pub metadata: HashMap<String, String>,
73}
74
75impl ModuleDagDefinition {
76    /// Returns a fluent builder.
77    ///
78    /// ```ignore
79    /// let dag = ModuleDagDefinition::builder()
80    ///     .edge(&login_node, &cate_list_node)
81    ///     .edge(&cate_list_node, &brand_rank_downloader)
82    ///     .edge(&cate_list_node, &goods_cate_downloader)
83    ///     .edge(&brand_rank_downloader, &download_url_node)
84    ///     .edge(&goods_cate_downloader, &download_url_node)
85    ///     .edge(&download_url_node, &download_file_node)
86    ///     .build();
87    /// ```
88    ///
89    /// Nodes are collected automatically from edges in first-seen order.
90    /// Entry nodes (those with no incoming edges) are derived automatically.
91    /// Isolated nodes (no edges at all) can be added with `.node()`.
92    pub fn builder() -> ModuleDagBuilder {
93        ModuleDagBuilder::new()
94    }
95
96    pub fn from_linear_steps(steps: Vec<Arc<dyn ModuleNodeTrait>>) -> Self {
97        let mut nodes = Vec::with_capacity(steps.len());
98        let mut edges = Vec::with_capacity(steps.len().saturating_sub(1));
99
100        for step in steps {
101            nodes.push(ModuleDagNodeDef::new(step));
102        }
103        for idx in 1..nodes.len() {
104            edges.push(ModuleDagEdgeDef::new(&nodes[idx - 1], &nodes[idx]));
105        }
106
107        let entry_nodes = nodes
108            .first()
109            .map(|n| vec![n.node_id.clone()])
110            .unwrap_or_default();
111
112        Self {
113            nodes,
114            edges,
115            entry_nodes,
116            default_policy: None,
117            metadata: HashMap::new(),
118        }
119    }
120}
121
122// ── Builder ──────────────────────────────────────────────────────────────────
123
124/// Fluent builder for `ModuleDagDefinition`.
125///
126/// Collects nodes automatically from `.edge()` calls. Entry nodes and the
127/// `nodes` list are both derived automatically when `.build()` is called, so
128/// callers only need to describe the edges between node definitions.
129pub struct ModuleDagBuilder {
130    /// Nodes in insertion order, keyed by node_id to deduplicate.
131    nodes: IndexMap<String, ModuleDagNodeDef>,
132    edges: Vec<ModuleDagEdgeDef>,
133    default_policy: Option<DagNodeExecutionPolicy>,
134}
135
136impl ModuleDagBuilder {
137    pub fn new() -> Self {
138        Self {
139            nodes: IndexMap::new(),
140            edges: Vec::new(),
141            default_policy: None,
142        }
143    }
144
145    /// Registers a directed edge `from → to`.
146    ///
147    /// Both nodes are inserted into the node registry on first encounter
148    /// (subsequent calls with the same `node_id` are silently ignored so
149    /// the same `ModuleDagNodeDef` reference is safe to reuse across calls).
150    pub fn edge(mut self, from: &ModuleDagNodeDef, to: &ModuleDagNodeDef) -> Self {
151        self.nodes
152            .entry(from.node_id.clone())
153            .or_insert_with(|| from.clone());
154        self.nodes
155            .entry(to.node_id.clone())
156            .or_insert_with(|| to.clone());
157        self.edges.push(ModuleDagEdgeDef {
158            from: from.node_id.clone(),
159            to: to.node_id.clone(),
160        });
161        self
162    }
163
164    /// Registers an isolated node (no edges). Useful for single-node DAGs.
165    pub fn node(mut self, node: &ModuleDagNodeDef) -> Self {
166        self.nodes
167            .entry(node.node_id.clone())
168            .or_insert_with(|| node.clone());
169        self
170    }
171
172    /// Sets a default execution policy applied to every node.
173    pub fn default_policy(mut self, policy: DagNodeExecutionPolicy) -> Self {
174        self.default_policy = Some(policy);
175        self
176    }
177
178    /// Consumes the builder and produces a `ModuleDagDefinition`.
179    ///
180    /// Entry nodes are derived as any node that never appears as a `to` in an edge.
181    pub fn build(self) -> ModuleDagDefinition {
182        let to_set: std::collections::HashSet<&str> =
183            self.edges.iter().map(|e| e.to.as_str()).collect();
184
185        let entry_nodes: Vec<String> = self
186            .nodes
187            .keys()
188            .filter(|id| !to_set.contains(id.as_str()))
189            .cloned()
190            .collect();
191
192        ModuleDagDefinition {
193            nodes: self.nodes.into_values().collect(),
194            edges: self.edges,
195            entry_nodes,
196            default_policy: self.default_policy,
197            metadata: HashMap::new(),
198        }
199    }
200}
201
202impl Default for ModuleDagBuilder {
203    fn default() -> Self {
204        Self::new()
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    use std::sync::Arc;
211
212    use async_trait::async_trait;
213    use serde_json::Map;
214
215    use super::ModuleDagDefinition;
216    use crate::common::interface::{ModuleNodeTrait, SyncBoxStream, ToSyncBoxStream};
217    use crate::common::model::login_info::LoginInfo;
218    use crate::common::model::message::TaskOutputEvent;
219    use crate::common::model::{ModuleConfig, Request, Response};
220    use crate::errors::Result;
221
222    struct DummyNode;
223
224    #[async_trait]
225    impl ModuleNodeTrait for DummyNode {
226        async fn generate(
227            &self,
228            _config: Arc<ModuleConfig>,
229            _params: Map<String, serde_json::Value>,
230            _login_info: Option<LoginInfo>,
231        ) -> Result<SyncBoxStream<'static, Request>> {
232            Ok(Vec::<Request>::new().to_stream())
233        }
234
235        async fn parser(
236            &self,
237            _response: Response,
238            _config: Option<Arc<ModuleConfig>>,
239        ) -> Result<TaskOutputEvent> {
240            Ok(TaskOutputEvent::default())
241        }
242    }
243
244    fn dummy_node() -> Arc<dyn ModuleNodeTrait> {
245        Arc::new(DummyNode)
246    }
247
248    #[test]
249    fn from_linear_steps_builds_expected_shape() {
250        let steps = vec![dummy_node(), dummy_node(), dummy_node()];
251        let definition = ModuleDagDefinition::from_linear_steps(steps);
252
253        assert_eq!(definition.nodes.len(), 3);
254        assert_eq!(definition.edges.len(), 2);
255        assert_eq!(
256            definition.entry_nodes,
257            vec![definition.nodes[0].node_id.clone()]
258        );
259        assert_eq!(definition.edges[0].from, definition.nodes[0].node_id);
260        assert_eq!(definition.edges[0].to, definition.nodes[1].node_id);
261        assert_eq!(definition.edges[1].from, definition.nodes[1].node_id);
262        assert_eq!(definition.edges[1].to, definition.nodes[2].node_id);
263    }
264}
mocra_core/common/model/module_dag.rs

mocra_core/common/model/
module_dag.rs
