Skip to main content

jellyflow_runtime/runtime/connection/
connect.rs

1use serde::{Deserialize, Serialize};
2
3use crate::rules::{
4    ConnectPlan, Diagnostic, plan_connect_typed_with_mode_and_policy,
5    plan_connect_with_mode_and_policy,
6};
7use crate::runtime::store::{DispatchError, DispatchOutcome, NodeGraphStore};
8use jellyflow_core::core::{EdgeId, PortId};
9use jellyflow_core::interaction::NodeGraphConnectionMode;
10use jellyflow_core::ops::{GraphOp, GraphTransaction};
11use jellyflow_core::types::DefaultTypeCompatibility;
12
13/// Default transaction label used for committed connect updates.
14pub const CONNECT_EDGE_TRANSACTION_LABEL: &str = "connect edge";
15
16/// Rules-driven request for connecting two existing ports.
17#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
18pub struct ConnectEdgeRequest {
19    pub from: PortId,
20    pub to: PortId,
21    #[serde(default, skip_serializing_if = "Option::is_none")]
22    pub edge: Option<EdgeId>,
23    #[serde(default)]
24    pub mode: NodeGraphConnectionMode,
25}
26
27impl ConnectEdgeRequest {
28    pub fn new(from: PortId, to: PortId, mode: NodeGraphConnectionMode) -> Self {
29        Self {
30            from,
31            to,
32            edge: None,
33            mode,
34        }
35    }
36
37    pub fn with_edge_id(mut self, edge: EdgeId) -> Self {
38        self.edge = Some(edge);
39        self
40    }
41}
42
43/// Error returned when a connect request could not be committed.
44#[derive(Debug, thiserror::Error)]
45pub enum ConnectEdgeError {
46    #[error("connect edge was rejected")]
47    Rejected { diagnostics: Vec<Diagnostic> },
48    #[error(transparent)]
49    Dispatch(#[from] DispatchError),
50}
51
52impl ConnectEdgeError {
53    pub fn diagnostics(&self) -> Option<&[Diagnostic]> {
54        match self {
55            Self::Rejected { diagnostics } => Some(diagnostics),
56            Self::Dispatch(_) => None,
57        }
58    }
59}
60
61pub fn connect_edge_transaction(plan: &ConnectPlan) -> Option<GraphTransaction> {
62    connect_edge_transaction_with_optional_edge_id(plan, None)
63}
64
65pub fn connect_edge_transaction_with_edge_id(
66    plan: &ConnectPlan,
67    edge: EdgeId,
68) -> Option<GraphTransaction> {
69    connect_edge_transaction_with_optional_edge_id(plan, Some(edge))
70}
71
72fn connect_edge_transaction_with_optional_edge_id(
73    plan: &ConnectPlan,
74    edge: Option<EdgeId>,
75) -> Option<GraphTransaction> {
76    if !plan.is_accept() || plan.ops().is_empty() {
77        return None;
78    }
79
80    Some(
81        GraphTransaction::from_ops(connect_edge_ops(plan.ops().iter().cloned(), edge))
82            .with_label(CONNECT_EDGE_TRANSACTION_LABEL),
83    )
84}
85
86fn connect_edge_transaction_from_request(
87    plan: ConnectPlan,
88    request: ConnectEdgeRequest,
89) -> Option<GraphTransaction> {
90    if !plan.is_accept() || plan.ops().is_empty() {
91        return None;
92    }
93
94    Some(
95        GraphTransaction::from_ops(connect_edge_ops(plan.into_ops(), request.edge))
96            .with_label(CONNECT_EDGE_TRANSACTION_LABEL),
97    )
98}
99
100fn connect_edge_ops(
101    ops: impl IntoIterator<Item = GraphOp>,
102    edge: Option<EdgeId>,
103) -> impl Iterator<Item = GraphOp> {
104    ops.into_iter().map(move |op| match (edge, op) {
105        (Some(id), GraphOp::AddEdge { edge, .. }) => GraphOp::AddEdge { id, edge },
106        (_, op) => op,
107    })
108}
109
110impl NodeGraphStore {
111    /// Plans connecting two existing ports against the resolved interaction policy.
112    pub fn plan_connect_edge(&self, request: ConnectEdgeRequest) -> ConnectPlan {
113        let interaction = self.resolved_interaction_state();
114        let has_typed_endpoint = |port: PortId| {
115            self.graph()
116                .ports()
117                .get(&port)
118                .is_some_and(|port| port.ty.is_some())
119        };
120        if has_typed_endpoint(request.from) || has_typed_endpoint(request.to) {
121            let mut compat = DefaultTypeCompatibility;
122            return plan_connect_typed_with_mode_and_policy(
123                self.graph(),
124                request.from,
125                request.to,
126                request.mode,
127                &interaction,
128                |graph, port| graph.ports().get(&port).and_then(|port| port.ty.clone()),
129                &mut compat,
130            );
131        }
132
133        plan_connect_with_mode_and_policy(
134            self.graph(),
135            request.from,
136            request.to,
137            request.mode,
138            &interaction,
139        )
140    }
141
142    /// Commits a connect request through normal store dispatch.
143    pub fn apply_connect_edge(
144        &mut self,
145        request: ConnectEdgeRequest,
146    ) -> Result<Option<DispatchOutcome>, ConnectEdgeError> {
147        let plan = self.plan_connect_edge(request);
148        if plan.is_reject() {
149            return Err(ConnectEdgeError::Rejected {
150                diagnostics: plan.diagnostics,
151            });
152        }
153
154        let Some(transaction) = connect_edge_transaction_from_request(plan, request) else {
155            return Ok(None);
156        };
157
158        self.dispatch_transaction(&transaction)
159            .map(Some)
160            .map_err(ConnectEdgeError::from)
161    }
162}