nexus_sdk/scheduler/
input_plan.rs

1//! Preflight and materialization of task inputs with explicit remote selections.
2
3use {
4    super::input_json::{hint_remote_fields, is_canonical_nexus_data, nexus_data_from_json_value},
5    crate::{
6        execution_limits::MAX_RESOLVED_DATA_BYTES,
7        scheduler::TaskInputs,
8        types::NexusData,
9        walrus::WalrusUploadData,
10    },
11    anyhow::{anyhow, ensure, Result as AnyResult},
12    serde_json::Value,
13    std::collections::{BTreeMap, HashSet},
14};
15
16/// Local preflight commitments and the exact bytes awaiting explicit upload.
17#[derive(Debug)]
18pub struct TaskInputPlan {
19    prepared: TaskInputs,
20    uploads: BTreeMap<String, WalrusUploadData>,
21}
22
23impl TaskInputPlan {
24    /// Parses all local inputs and remote selectors without making a request.
25    /// Arrays denote Many ports. Canonical references can be supplied directly.
26    pub fn new(input: &Value, remote: &[String]) -> AnyResult<Self> {
27        let vertices = input
28            .as_object()
29            .ok_or_else(|| anyhow!("Input JSON must be an object with vertex names as keys"))?;
30        let mut selected = HashSet::new();
31        for handle in remote {
32            ensure!(
33                selected.insert(handle.clone()),
34                "Remote input selector '{handle}' is duplicated"
35            );
36        }
37        let mut prepared = TaskInputs::new();
38        let mut uploads = BTreeMap::new();
39        let mut flattened = serde_json::Map::new();
40        for (vertex, data) in vertices {
41            let mut total_bytes = 0usize;
42            let ports = data.as_object().ok_or_else(|| {
43                anyhow!(
44                    "Input JSON for vertex '{vertex}' must be an object with port names as keys"
45                )
46            })?;
47            let mut result = BTreeMap::new();
48            for (port, value) in ports {
49                let handle = format!("{vertex}.{port}");
50                ensure!(
51                    !flattened.contains_key(&handle),
52                    "ambiguous input selector '{handle}'"
53                );
54                flattened.insert(handle.clone(), value.clone());
55                let value = if selected.contains(&handle) {
56                    ensure!(!is_canonical_nexus_data(value), "remote selections accept JSON data; reuse existing canonical references directly");
57                    let upload = WalrusUploadData::from_json_document(
58                        serde_json::to_vec(value)?,
59                        value.is_array(),
60                    )?;
61                    total_bytes = total_bytes
62                        .checked_add(upload.byte_len())
63                        .ok_or_else(|| anyhow!("input size overflow"))?;
64                    let placeholder = upload.preflight_reference()?;
65                    uploads.insert(handle, upload);
66                    placeholder
67                } else {
68                    nexus_data_from_json_value(value.clone()).map_err(|error| anyhow!("input '{handle}': {error}; select this port for remote storage when its data is large"))?
69                };
70                result.insert(port.clone(), value);
71            }
72            ensure!(
73                total_bytes <= MAX_RESOLVED_DATA_BYTES,
74                "remote inputs for vertex '{vertex}' exceed the execution byte limit"
75            );
76            prepared.insert(vertex.clone(), result);
77        }
78        for handle in &selected {
79            ensure!(
80                flattened.contains_key(handle),
81                "Remote input selector '{handle}' does not identify an input field"
82            );
83        }
84        let hints = hint_remote_fields(&Value::Object(flattened), &selected)?;
85        ensure!(
86            hints.is_empty(),
87            "Inputs exceed transaction size limits; select these remote ports: {}",
88            hints.join(",")
89        );
90        Ok(Self { prepared, uploads })
91    }
92
93    /// Returns commitments for authoritative DAG preflight. Pending uploads use placeholders.
94    pub fn preflight_inputs(&self) -> TaskInputs {
95        self.prepared.clone()
96    }
97
98    /// Whether the plan contains payloads awaiting upload.
99    pub fn has_uploads(&self) -> bool {
100        !self.uploads.is_empty()
101    }
102
103    /// Whether the plan includes remote commitments that already exist.
104    pub fn has_references(&self) -> bool {
105        self.prepared.iter().any(|(vertex, ports)| {
106            ports.iter().any(|(port, data)| {
107                data.has_walrus() && !self.uploads.contains_key(&format!("{vertex}.{port}"))
108            })
109        })
110    }
111
112    /// Resolves existing commitments before a new paid upload is allowed to start.
113    pub async fn verify_references(&self, reader: &crate::walrus::WalrusReader) -> AnyResult<()> {
114        for (vertex, ports) in &self.prepared {
115            let existing = ports
116                .iter()
117                .filter(|(port, _)| !self.uploads.contains_key(&format!("{vertex}.{port}")))
118                .map(|(port, value)| (port.clone(), value.clone()))
119                .collect();
120            let resolved = reader.resolve_ports(existing).await?;
121            let mut remaining = MAX_RESOLVED_DATA_BYTES;
122            for value in resolved.values().flatten() {
123                let size = match value {
124                    crate::types::NexusValue::InlineData { bytes } => bytes.len(),
125                    _ => 32,
126                };
127                remaining = remaining
128                    .checked_sub(size)
129                    .ok_or_else(|| anyhow!("resolved inputs exceed the execution byte limit"))?;
130            }
131            for port in ports.keys() {
132                if let Some(upload) = self.uploads.get(&format!("{vertex}.{port}")) {
133                    remaining = remaining.checked_sub(upload.byte_len()).ok_or_else(|| {
134                        anyhow!(
135                            "resolved inputs and pending uploads exceed the execution byte limit"
136                        )
137                    })?;
138                }
139            }
140        }
141        Ok(())
142    }
143
144    /// Uploads after preflight through the caller's storage policy and checks every commitment.
145    pub async fn materialize_with<F, Fut>(mut self, mut upload: F) -> AnyResult<TaskInputs>
146    where
147        F: FnMut(String, WalrusUploadData) -> Fut,
148        Fut: std::future::Future<Output = AnyResult<NexusData>>,
149    {
150        for (handle, data) in self.uploads {
151            let result = upload(handle.clone(), data.clone()).await?;
152            data.verify_reference(&result)?;
153            for (vertex, ports) in &mut self.prepared {
154                for (port, value) in ports {
155                    if format!("{vertex}.{port}") == handle {
156                        ensure!(
157                            result.is_well_formed()
158                                && result.has_walrus()
159                                && result.is_many() == value.is_many()
160                                && result.values()?.len() == value.values()?.len(),
161                            "upload changed the preflight port shape"
162                        );
163                        *value = result.clone();
164                    }
165                }
166            }
167        }
168        Ok(self.prepared)
169    }
170}
171
172#[cfg(test)]
173mod tests {
174    use {super::*, serde_json::json};
175
176    #[test]
177    fn each_vertex_has_an_independent_execution_budget() {
178        let data = "x".repeat(MAX_RESOLVED_DATA_BYTES / 2);
179        let separate = TaskInputPlan::new(
180            &json!({"first": {"input": data}, "second": {"input": data}}),
181            &["first.input".into(), "second.input".into()],
182        )
183        .expect("separate invocations do not share a data budget");
184        assert_eq!(separate.uploads.len(), 2);
185
186        let error = TaskInputPlan::new(
187            &json!({"one": {"first": data, "second": data}}),
188            &["one.first".into(), "one.second".into()],
189        )
190        .unwrap_err();
191        assert!(error.to_string().contains("vertex 'one'"));
192    }
193
194    #[tokio::test]
195    async fn prepares_large_remote_values_and_materializes_only_selected_ports() {
196        let value = "x".repeat(100_000);
197        let plan = TaskInputPlan::new(
198            &json!({"v":{"large":value, "inline":"small"}}),
199            &["v.large".into()],
200        )
201        .unwrap();
202        assert!(plan.preflight_inputs()["v"]["large"].has_walrus());
203        let result = plan
204            .materialize_with(|handle, bytes| async move {
205                assert_eq!(handle, "v.large");
206                assert_eq!(
207                    bytes.values(),
208                    &[serde_json::to_vec(&"x".repeat(100_000)).unwrap()]
209                );
210                {
211                    use sha2::{Digest as _, Sha256};
212                    NexusData::walrus_data(
213                        b"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
214                        Sha256::digest(&bytes.values()[0]).to_vec(),
215                    )
216                }
217            })
218            .await
219            .unwrap();
220        assert!(result["v"]["large"].has_walrus());
221        assert!(!result["v"]["inline"].has_walrus());
222    }
223
224    #[test]
225    fn validates_all_local_decisions_before_uploading() {
226        assert!(TaskInputPlan::new(&json!("invalid"), &[]).is_err());
227        assert!(
228            TaskInputPlan::new(&json!({"a":{"port":"x"},"z":false}), &["a.port".into()]).is_err()
229        );
230        assert!(TaskInputPlan::new(
231            &json!({"a":{"port":"x"}}),
232            &["a.port".into(), "a.port".into()]
233        )
234        .is_err());
235        assert!(TaskInputPlan::new(&json!({"a":{"port":"x"}}), &["a.missing".into()]).is_err());
236        assert!(TaskInputPlan::new(&json!({"a":{"port":"x".repeat(100_000)}}), &[]).is_err());
237        let ports = (0..6)
238            .map(|i| (format!("p{i}"), json!(vec!["x"; 128])))
239            .collect::<serde_json::Map<_, _>>();
240        assert!(TaskInputPlan::new(
241            &json!({"a":ports}),
242            &(0..6).map(|i| format!("a.p{i}")).collect::<Vec<_>>()
243        )
244        .is_err());
245    }
246
247    #[tokio::test]
248    async fn verifies_existing_references_before_materializing_new_uploads() {
249        use sha2::{Digest as _, Sha256};
250        let mut server = mockito::Server::new_async().await;
251        let bytes = b"\"existing data\"";
252        let reference = NexusData::walrus_data(
253            b"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
254            Sha256::digest(bytes).to_vec(),
255        )
256        .unwrap();
257        let get = server
258            .mock(
259                "GET",
260                "/v1/blobs/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
261            )
262            .with_body(bytes)
263            .expect(1)
264            .create_async()
265            .await;
266        let plan = TaskInputPlan::new(
267            &json!({"v": {
268                "existing": reference.to_json_value().unwrap(),
269                "new": "new data",
270                "object": NexusData::object(crate::sui::types::Address::TWO).to_json_value().unwrap(),
271            }}),
272            &["v.new".into()],
273        )
274        .unwrap();
275        assert!(plan.has_uploads());
276        assert!(plan.has_references());
277        plan.verify_references(&crate::walrus::WalrusReader::new(&server.url(), 100).unwrap())
278            .await
279            .unwrap();
280        get.assert_async().await;
281        let error = plan
282            .materialize_with(|handle, _| async move {
283                assert_eq!(handle, "v.new");
284                NexusData::walrus_data(b"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA", [0; 32])
285            })
286            .await
287            .unwrap_err();
288        assert!(error.to_string().contains("different content"));
289    }
290
291    #[tokio::test]
292    async fn existing_bytes_and_pending_uploads_must_fit_one_execution() {
293        use sha2::{Digest as _, Sha256};
294        let mut server = mockito::Server::new_async().await;
295        let bytes = vec![b'a'; MAX_RESOLVED_DATA_BYTES - 10];
296        let reference = NexusData::walrus_data(
297            b"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
298            Sha256::digest(&bytes).to_vec(),
299        )
300        .unwrap();
301        let get = server
302            .mock(
303                "GET",
304                "/v1/blobs/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
305            )
306            .with_body(bytes)
307            .expect(1)
308            .create_async()
309            .await;
310        let plan = TaskInputPlan::new(
311            &json!({"v": {"existing": reference.to_json_value().unwrap(), "new": "data exceeds the remainder"}}),
312            &["v.new".into()],
313        )
314        .unwrap();
315        let reader =
316            crate::walrus::WalrusReader::new(&server.url(), MAX_RESOLVED_DATA_BYTES).unwrap();
317        let error = plan.verify_references(&reader).await.unwrap_err();
318        assert!(error.to_string().contains("pending uploads exceed"));
319        get.assert_async().await;
320    }
321}
nexus_sdk/scheduler/input_plan.rs

nexus_sdk/scheduler/
input_plan.rs
