Skip to main content

sz_rust_workflow/engine/
plugin_node.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2024-2026 SZ-Rust Team
3//
4use std::sync::Arc;
5use std::time::Duration;
6
7use serde_json::Value;
8
9use crate::definition::NodeConfig;
10use crate::error::{WorkflowError, WorkflowErrorCode, WorkflowResult};
11use crate::integration::SensitiveFieldRegistry;
12use crate::scheduling::fault_strategy::{FaultDecision, FaultStrategyHandler, PluginNodeOutcome};
13
14/// 插件节点执行器,对齐 design 2.2.2.8。
15pub struct PluginNodeExecutor {
16    capability_registry: Arc<sz_rust_capability::CapabilityRegistry>,
17    fault_handler: Arc<dyn FaultStrategyHandler>,
18    sensitive_registry: Arc<SensitiveFieldRegistry>,
19    timeout: Duration,
20}
21
22impl PluginNodeExecutor {
23    pub fn new(
24        capability_registry: Arc<sz_rust_capability::CapabilityRegistry>,
25        fault_handler: Arc<dyn FaultStrategyHandler>,
26        sensitive_registry: Arc<SensitiveFieldRegistry>,
27        timeout: Duration,
28    ) -> Self {
29        Self {
30            capability_registry,
31            fault_handler,
32            sensitive_registry,
33            timeout,
34        }
35    }
36
37    /// 执行插件节点。
38    pub async fn execute(
39        &self,
40        node: &NodeConfig,
41        context: &mut Value,
42    ) -> WorkflowResult<PluginNodeOutcome> {
43        let (capability_name, version_range, args_mapping, fault_strategy, _output_schema, next) =
44            match node {
45                NodeConfig::Plugin {
46                    capability_name,
47                    capability_version_range,
48                    args_mapping,
49                    fault_strategy,
50                    output_schema,
51                    next,
52                } => (
53                    capability_name,
54                    capability_version_range,
55                    args_mapping,
56                    *fault_strategy,
57                    output_schema,
58                    next,
59                ),
60                _ => {
61                    return Err(WorkflowError::new(
62                        WorkflowErrorCode::CapabilityNotFound,
63                        "非插件节点",
64                    ))
65                }
66            };
67
68        let _ = next;
69        let args = self.resolve_args(args_mapping, context);
70
71        let capability = match self.capability_registry.get(capability_name) {
72            Some(c) => c,
73            None => {
74                let wf_error = WorkflowError::with_field(
75                    WorkflowErrorCode::CapabilityNotFound,
76                    "能力不存在",
77                    "capability",
78                    capability_name,
79                );
80                let decision = self.fault_handler.decide(fault_strategy, &wf_error, 0);
81                return match decision {
82                    FaultDecision::Terminate => Ok(PluginNodeOutcome::InstanceTerminated),
83                    FaultDecision::Skip => Ok(PluginNodeOutcome::Skipped),
84                    FaultDecision::Retry { .. } => Ok(PluginNodeOutcome::InstanceTerminated),
85                };
86            }
87        };
88
89        let req = semver::VersionReq::parse(version_range).map_err(|e| {
90            WorkflowError::with_field(
91                WorkflowErrorCode::CapabilityNotFound,
92                format!("非法版本范围:{e}"),
93                "range",
94                version_range,
95            )
96        })?;
97        let cap_version =
98            semver::Version::parse(capability.version()).unwrap_or(semver::Version::new(0, 0, 0));
99        if !req.matches(&cap_version) {
100            return Err(WorkflowError::with_field(
101                WorkflowErrorCode::CapabilityNotFound,
102                format!(
103                    "能力版本 {} 不满足范围 {}",
104                    capability.version(),
105                    version_range
106                ),
107                "capability",
108                capability_name,
109            ));
110        }
111
112        let mut attempt = 0u32;
113        loop {
114            let call_result = tokio::time::timeout(
115                self.timeout,
116                self.capability_registry.call(capability_name, args.clone()),
117            )
118            .await;
119
120            match call_result {
121                Ok(Ok(result)) => {
122                    self.sensitive_registry
123                        .merge_capability_result(context, result);
124                    return Ok(PluginNodeOutcome::Completed);
125                }
126                Ok(Err(e)) => {
127                    let wf_error = WorkflowError::new(
128                        WorkflowErrorCode::CapabilityNotFound,
129                        format!("能力调用失败:{e}"),
130                    );
131                    let decision = self
132                        .fault_handler
133                        .decide(fault_strategy, &wf_error, attempt);
134                    match decision {
135                        FaultDecision::Terminate => {
136                            return Ok(PluginNodeOutcome::InstanceTerminated)
137                        }
138                        FaultDecision::Skip => return Ok(PluginNodeOutcome::Skipped),
139                        FaultDecision::Retry { backoff, .. } => {
140                            attempt += 1;
141                            tokio::time::sleep(backoff).await;
142                            continue;
143                        }
144                    }
145                }
146                Err(_) => {
147                    let wf_error = WorkflowError::with_field(
148                        WorkflowErrorCode::CapabilityTimeout,
149                        "能力调用超时",
150                        "capability",
151                        capability_name,
152                    );
153                    let decision = self
154                        .fault_handler
155                        .decide(fault_strategy, &wf_error, attempt);
156                    match decision {
157                        FaultDecision::Terminate => {
158                            return Ok(PluginNodeOutcome::InstanceTerminated)
159                        }
160                        FaultDecision::Skip => return Ok(PluginNodeOutcome::Skipped),
161                        FaultDecision::Retry { backoff, .. } => {
162                            attempt += 1;
163                            tokio::time::sleep(backoff).await;
164                            continue;
165                        }
166                    }
167                }
168            }
169        }
170    }
171
172    fn resolve_args(&self, mapping: &Value, context: &Value) -> Value {
173        if mapping.is_null() {
174            return context.clone();
175        }
176        if let Value::Object(map) = mapping {
177            let mut result = serde_json::Map::new();
178            for (k, v) in map {
179                if let Some(s) = v.as_str() {
180                    if let Some(path) = s.strip_prefix("$.") {
181                        if let Some(val) = lookup(context, path) {
182                            result.insert(k.clone(), val.clone());
183                            continue;
184                        }
185                    }
186                }
187                result.insert(k.clone(), v.clone());
188            }
189            return Value::Object(result);
190        }
191        mapping.clone()
192    }
193}
194
195fn lookup<'a>(ctx: &'a Value, path: &str) -> Option<&'a Value> {
196    let mut current = ctx;
197    for part in path.split('.') {
198        if let Value::Object(obj) = current {
199            current = obj.get(part)?;
200        } else {
201            return None;
202        }
203    }
204    Some(current)
205}
206
207#[cfg(test)]
208mod tests {
209    use super::*;
210    use crate::definition::FaultStrategy;
211    use crate::scheduling::fault_strategy::DefaultFaultStrategyHandler;
212    use async_trait::async_trait;
213    use sz_rust_capability::{CapResult, Capability, CapabilitySource};
214
215    struct EchoCap;
216    #[async_trait]
217    impl Capability for EchoCap {
218        fn name(&self) -> &'static str {
219            "test.echo"
220        }
221        fn description(&self) -> &'static str {
222            "回显"
223        }
224        fn schema(&self) -> Value {
225            Value::Object(serde_json::Map::new())
226        }
227        fn tags(&self) -> &[&'static str] {
228            &["test"]
229        }
230        fn source(&self) -> CapabilitySource {
231            CapabilitySource::Skill
232        }
233        async fn call(&self, args: Value) -> CapResult<Value> {
234            Ok(args)
235        }
236    }
237
238    #[tokio::test]
239    async fn execute_success() {
240        let registry = Arc::new(sz_rust_capability::CapabilityRegistry::new());
241        registry.register(Arc::new(EchoCap));
242        let executor = PluginNodeExecutor::new(
243            registry,
244            Arc::new(DefaultFaultStrategyHandler::default()),
245            Arc::new(SensitiveFieldRegistry::new()),
246            Duration::from_secs(5),
247        );
248        let node = NodeConfig::Plugin {
249            capability_name: "test.echo".into(),
250            capability_version_range: "*".into(),
251            args_mapping: serde_json::json!({"key": "value"}),
252            fault_strategy: FaultStrategy::Fail,
253            output_schema: None,
254            next: "end".into(),
255        };
256        let mut ctx = serde_json::json!({});
257        let result = executor.execute(&node, &mut ctx).await.unwrap();
258        assert_eq!(result, PluginNodeOutcome::Completed);
259        assert_eq!(ctx["key"], "value");
260    }
261
262    #[tokio::test]
263    async fn execute_capability_not_found() {
264        let registry = Arc::new(sz_rust_capability::CapabilityRegistry::new());
265        let executor = PluginNodeExecutor::new(
266            registry,
267            Arc::new(DefaultFaultStrategyHandler::default()),
268            Arc::new(SensitiveFieldRegistry::new()),
269            Duration::from_secs(5),
270        );
271        let node = NodeConfig::Plugin {
272            capability_name: "nonexistent.cap".into(),
273            capability_version_range: "*".into(),
274            args_mapping: Value::Null,
275            fault_strategy: FaultStrategy::Fail,
276            output_schema: None,
277            next: "end".into(),
278        };
279        let mut ctx = serde_json::json!({});
280        let result = executor.execute(&node, &mut ctx).await.unwrap();
281        assert_eq!(result, PluginNodeOutcome::InstanceTerminated);
282    }
283}