vv_agent/runtime/backends/distributed/
dispatch.rs

1use serde::{de::Error as _, ser::Error as _, Deserialize, Deserializer, Serialize, Serializer};
2use serde_json::{Map, Value};
3
4use crate::runtime::CancellationToken;
5use crate::types::{AgentResult, AgentStatus};
6
7use super::DistributedRunEnvelope;
8
9pub const DISTRIBUTED_WORKER_RESPONSE_SCHEMA_VERSION: &str =
10    "vv-agent.distributed-worker-response.v1";
11const JSON_SAFE_INTEGER_MAX: u64 = (1_u64 << 53) - 1;
12const INVALID_AGENT_RESULT: &str =
13    "distributed worker response result must be a complete current AgentResult";
14
15#[derive(Debug, Clone, PartialEq)]
16pub enum CycleDispatchResult {
17    Pending,
18    Committed {
19        checkpoint_revision: u64,
20        committed_cycle: u64,
21    },
22    TerminalCandidate {
23        checkpoint_revision: u64,
24        result: AgentResult,
25    },
26    TerminalReplay {
27        checkpoint_revision: u64,
28        result: AgentResult,
29    },
30}
31
32impl Serialize for CycleDispatchResult {
33    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
34    where
35        S: Serializer,
36    {
37        self.validate().map_err(S::Error::custom)?;
38        self.wire_value().serialize(serializer)
39    }
40}
41
42impl<'de> Deserialize<'de> for CycleDispatchResult {
43    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
44    where
45        D: Deserializer<'de>,
46    {
47        let value = Value::deserialize(deserializer)?;
48        Self::from_dict(&value).map_err(D::Error::custom)
49    }
50}
51
52impl CycleDispatchResult {
53    pub const fn pending() -> Self {
54        Self::Pending
55    }
56
57    pub fn committed(cycle_index: u64, checkpoint_revision: u64) -> Result<Self, String> {
58        let result = Self::Committed {
59            checkpoint_revision,
60            committed_cycle: cycle_index,
61        };
62        result.validate()?;
63        Ok(result)
64    }
65
66    pub fn terminal_candidate(
67        result: AgentResult,
68        checkpoint_revision: u64,
69    ) -> Result<Self, String> {
70        let result = Self::TerminalCandidate {
71            checkpoint_revision,
72            result,
73        };
74        result.validate()?;
75        Ok(result)
76    }
77
78    pub fn terminal_replay(result: AgentResult, checkpoint_revision: u64) -> Result<Self, String> {
79        let result = Self::TerminalReplay {
80            checkpoint_revision,
81            result,
82        };
83        result.validate()?;
84        Ok(result)
85    }
86
87    pub const fn kind(&self) -> &'static str {
88        match self {
89            Self::Pending => "pending",
90            Self::Committed { .. } => "committed",
91            Self::TerminalCandidate { .. } => "terminal_candidate",
92            Self::TerminalReplay { .. } => "terminal_replay",
93        }
94    }
95
96    pub fn to_dict(&self) -> Value {
97        self.validate()
98            .expect("CycleDispatchResult must satisfy the distributed worker response contract");
99        self.wire_value()
100    }
101
102    pub fn from_dict(data: &Value) -> Result<Self, String> {
103        let object = data
104            .as_object()
105            .ok_or_else(|| "distributed worker response must be an object".to_string())?;
106        if object.get("schema_version").and_then(Value::as_str)
107            != Some(DISTRIBUTED_WORKER_RESPONSE_SCHEMA_VERSION)
108        {
109            return Err("unsupported distributed worker response schema_version".to_string());
110        }
111        let response_type = object
112            .get("type")
113            .and_then(Value::as_str)
114            .ok_or_else(|| "unsupported distributed worker response type".to_string())?;
115
116        match response_type {
117            "pending" => {
118                require_exact_fields(object, response_type, &["schema_version", "type"])?;
119                Ok(Self::Pending)
120            }
121            "committed" => {
122                require_exact_fields(
123                    object,
124                    response_type,
125                    &[
126                        "schema_version",
127                        "type",
128                        "checkpoint_revision",
129                        "committed_cycle",
130                    ],
131                )?;
132                Ok(Self::Committed {
133                    checkpoint_revision: read_wire_integer(object, "checkpoint_revision")?,
134                    committed_cycle: read_positive_wire_integer(object, "committed_cycle")?,
135                })
136            }
137            "terminal_candidate" | "terminal_replay" => {
138                require_exact_fields(
139                    object,
140                    response_type,
141                    &["schema_version", "type", "checkpoint_revision", "result"],
142                )?;
143                let checkpoint_revision = read_wire_integer(object, "checkpoint_revision")?;
144                let result = parse_complete_agent_result(
145                    object.get("result").expect("exact fields checked above"),
146                )?;
147                if response_type == "terminal_candidate" {
148                    validate_terminal_candidate_result(&result)?;
149                    Ok(Self::TerminalCandidate {
150                        checkpoint_revision,
151                        result,
152                    })
153                } else {
154                    validate_terminal_replay_result(&result)?;
155                    Ok(Self::TerminalReplay {
156                        checkpoint_revision,
157                        result,
158                    })
159                }
160            }
161            _ => Err("unsupported distributed worker response type".to_string()),
162        }
163    }
164
165    fn validate(&self) -> Result<(), String> {
166        match self {
167            Self::Pending => Ok(()),
168            Self::Committed {
169                checkpoint_revision,
170                committed_cycle,
171            } => {
172                validate_wire_integer(*checkpoint_revision, "checkpoint_revision")?;
173                validate_positive_wire_integer(*committed_cycle, "committed_cycle")
174            }
175            Self::TerminalCandidate {
176                checkpoint_revision,
177                result,
178            } => {
179                validate_wire_integer(*checkpoint_revision, "checkpoint_revision")?;
180                validate_terminal_candidate_result(result)
181            }
182            Self::TerminalReplay {
183                checkpoint_revision,
184                result,
185            } => {
186                validate_wire_integer(*checkpoint_revision, "checkpoint_revision")?;
187                validate_terminal_replay_result(result)
188            }
189        }
190    }
191
192    fn wire_value(&self) -> Value {
193        let mut payload = Map::from_iter([
194            (
195                "schema_version".to_string(),
196                Value::String(DISTRIBUTED_WORKER_RESPONSE_SCHEMA_VERSION.to_string()),
197            ),
198            ("type".to_string(), Value::String(self.kind().to_string())),
199        ]);
200        match self {
201            Self::Pending => {}
202            Self::Committed {
203                checkpoint_revision,
204                committed_cycle,
205            } => {
206                payload.insert(
207                    "checkpoint_revision".to_string(),
208                    Value::from(*checkpoint_revision),
209                );
210                payload.insert("committed_cycle".to_string(), Value::from(*committed_cycle));
211            }
212            Self::TerminalCandidate {
213                checkpoint_revision,
214                result,
215            }
216            | Self::TerminalReplay {
217                checkpoint_revision,
218                result,
219            } => {
220                payload.insert(
221                    "checkpoint_revision".to_string(),
222                    Value::from(*checkpoint_revision),
223                );
224                payload.insert("result".to_string(), result.to_dict());
225            }
226        }
227        Value::Object(payload)
228    }
229}
230
231fn require_exact_fields(
232    object: &Map<String, Value>,
233    response_type: &str,
234    expected: &[&str],
235) -> Result<(), String> {
236    if object.len() != expected.len() || expected.iter().any(|field| !object.contains_key(*field)) {
237        return Err(format!(
238            "distributed worker response fields do not match type {response_type}"
239        ));
240    }
241    Ok(())
242}
243
244fn read_wire_integer(object: &Map<String, Value>, field: &str) -> Result<u64, String> {
245    let value = object
246        .get(field)
247        .and_then(Value::as_u64)
248        .ok_or_else(|| format!("{field} must be a JSON-safe unsigned integer"))?;
249    validate_wire_integer(value, field)?;
250    Ok(value)
251}
252
253fn validate_wire_integer(value: u64, field: &str) -> Result<(), String> {
254    if value > JSON_SAFE_INTEGER_MAX {
255        return Err(format!("{field} must be a JSON-safe unsigned integer"));
256    }
257    Ok(())
258}
259
260fn read_positive_wire_integer(object: &Map<String, Value>, field: &str) -> Result<u64, String> {
261    let value = read_wire_integer(object, field)?;
262    validate_positive_wire_integer(value, field)?;
263    Ok(value)
264}
265
266fn validate_positive_wire_integer(value: u64, field: &str) -> Result<(), String> {
267    validate_wire_integer(value, field)?;
268    if value == 0 {
269        return Err(format!("{field} must be a positive JSON-safe integer"));
270    }
271    Ok(())
272}
273
274fn parse_complete_agent_result(value: &Value) -> Result<AgentResult, String> {
275    let result = AgentResult::from_dict(value).map_err(|_| invalid_agent_result())?;
276    if result.to_dict() != *value {
277        return Err(invalid_agent_result());
278    }
279    Ok(result)
280}
281
282fn validate_terminal_candidate_result(result: &AgentResult) -> Result<(), String> {
283    if matches!(result.status, AgentStatus::Pending | AgentStatus::Running) {
284        return Err(invalid_agent_result());
285    }
286    Ok(())
287}
288
289fn validate_terminal_replay_result(result: &AgentResult) -> Result<(), String> {
290    if matches!(
291        result.status,
292        AgentStatus::Pending | AgentStatus::Running | AgentStatus::ReconciliationRequired
293    ) {
294        return Err(invalid_agent_result());
295    }
296    Ok(())
297}
298
299fn invalid_agent_result() -> String {
300    INVALID_AGENT_RESULT.to_string()
301}
302
303pub trait CycleDispatcher: Send + Sync {
304    fn dispatch_envelope(
305        &self,
306        envelope: &DistributedRunEnvelope,
307    ) -> Result<CycleDispatchResult, String>;
308
309    fn dispatch_envelope_with_cancellation(
310        &self,
311        envelope: &DistributedRunEnvelope,
312        cancellation_token: Option<&CancellationToken>,
313    ) -> Result<CycleDispatchResult, String> {
314        check_cancellation(cancellation_token)?;
315        let result = self.dispatch_envelope(envelope)?;
316        if matches!(
317            &result,
318            CycleDispatchResult::TerminalCandidate { .. }
319                | CycleDispatchResult::TerminalReplay { .. }
320        ) {
321            return Ok(result);
322        }
323        check_cancellation(cancellation_token)?;
324        Ok(result)
325    }
326}
327
328fn check_cancellation(cancellation_token: Option<&CancellationToken>) -> Result<(), String> {
329    cancellation_token
330        .map(CancellationToken::check)
331        .transpose()
332        .map(|_| ())
333        .map_err(|reason| format!("distributed dispatch cancelled: {reason}"))
334}
vv_agent/runtime/backends/distributed/dispatch.rs

vv_agent/runtime/backends/distributed/
dispatch.rs
