vv_agent/runtime/backends/distributed/
dispatch.rs

1use serde::{de::Error as _, Deserialize, Deserializer, Serialize, Serializer};
2use serde_json::Value;
3
4use crate::runtime::backends::RuntimeRecipe;
5use crate::runtime::CancellationToken;
6use crate::types::{AgentResult, AgentTask};
7
8use super::DistributedRunEnvelope;
9
10#[derive(Debug, Clone, PartialEq)]
11pub struct CycleDispatchResult {
12    pub finished: bool,
13    pub result: Option<AgentResult>,
14    pub checkpoint_revision: Option<u64>,
15    pub committed_cycle: Option<u64>,
16    pub terminal_candidate: bool,
17    pub terminal_replay: bool,
18}
19
20impl Serialize for CycleDispatchResult {
21    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
22    where
23        S: Serializer,
24    {
25        self.to_dict().serialize(serializer)
26    }
27}
28
29impl<'de> Deserialize<'de> for CycleDispatchResult {
30    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
31    where
32        D: Deserializer<'de>,
33    {
34        let value = Value::deserialize(deserializer)?;
35        Self::from_dict(&value).map_err(D::Error::custom)
36    }
37}
38
39impl CycleDispatchResult {
40    pub fn unfinished() -> Self {
41        Self {
42            finished: false,
43            result: None,
44            checkpoint_revision: None,
45            committed_cycle: None,
46            terminal_candidate: false,
47            terminal_replay: false,
48        }
49    }
50
51    pub fn committed(cycle_index: u64, checkpoint_revision: u64) -> Self {
52        Self {
53            finished: false,
54            result: None,
55            checkpoint_revision: Some(checkpoint_revision),
56            committed_cycle: Some(cycle_index),
57            terminal_candidate: false,
58            terminal_replay: false,
59        }
60    }
61
62    pub fn finished(result: AgentResult) -> Self {
63        Self::finished_at_revision(result, None)
64    }
65
66    pub fn finished_at_revision(result: AgentResult, checkpoint_revision: Option<u64>) -> Self {
67        Self {
68            finished: true,
69            result: Some(result),
70            checkpoint_revision,
71            committed_cycle: None,
72            terminal_candidate: false,
73            terminal_replay: false,
74        }
75    }
76
77    pub fn terminal_candidate(result: AgentResult, checkpoint_revision: u64) -> Self {
78        Self {
79            finished: true,
80            result: Some(result),
81            checkpoint_revision: Some(checkpoint_revision),
82            committed_cycle: None,
83            terminal_candidate: true,
84            terminal_replay: false,
85        }
86    }
87
88    pub fn terminal_replay(result: AgentResult, checkpoint_revision: u64) -> Self {
89        Self {
90            finished: true,
91            result: Some(result),
92            checkpoint_revision: Some(checkpoint_revision),
93            committed_cycle: None,
94            terminal_candidate: false,
95            terminal_replay: true,
96        }
97    }
98
99    pub fn to_dict(&self) -> Value {
100        let mut payload =
101            serde_json::Map::from_iter([("finished".to_string(), Value::Bool(self.finished))]);
102        if let Some(result) = &self.result {
103            payload.insert("result".to_string(), result.to_dict());
104        }
105        if let Some(revision) = self.checkpoint_revision {
106            payload.insert("checkpoint_revision".to_string(), Value::from(revision));
107        }
108        if let Some(cycle_index) = self.committed_cycle {
109            payload.insert("committed_cycle".to_string(), Value::from(cycle_index));
110        }
111        if self.terminal_candidate {
112            payload.insert("terminal_candidate".to_string(), Value::Bool(true));
113        }
114        if self.terminal_replay {
115            payload.insert("terminal_replay".to_string(), Value::Bool(true));
116        }
117        Value::Object(payload)
118    }
119
120    pub fn from_dict(data: &Value) -> Result<Self, String> {
121        let object = data
122            .as_object()
123            .ok_or_else(|| "CycleDispatchResult payload must be an object".to_string())?;
124        let finished = object
125            .get("finished")
126            .and_then(Value::as_bool)
127            .ok_or_else(|| "CycleDispatchResult finished must be a boolean".to_string())?;
128        let result = object
129            .get("result")
130            .filter(|value| !value.is_null())
131            .map(AgentResult::from_dict)
132            .transpose()?;
133        if finished != result.is_some() {
134            return Err(
135                "CycleDispatchResult result must be present exactly when finished is true"
136                    .to_string(),
137            );
138        }
139        let checkpoint_revision = match object.get("checkpoint_revision") {
140            None | Some(Value::Null) => None,
141            Some(value) => Some(value.as_u64().ok_or_else(|| {
142                "CycleDispatchResult checkpoint_revision must be an unsigned integer".to_string()
143            })?),
144        };
145        let committed_cycle = match object.get("committed_cycle") {
146            None | Some(Value::Null) => None,
147            Some(value) => Some(value.as_u64().ok_or_else(|| {
148                "CycleDispatchResult committed_cycle must be an unsigned integer".to_string()
149            })?),
150        };
151        let terminal_candidate = optional_bool(object, "terminal_candidate")?;
152        let terminal_replay = optional_bool(object, "terminal_replay")?;
153        let result = Self {
154            finished,
155            result,
156            checkpoint_revision,
157            committed_cycle,
158            terminal_candidate,
159            terminal_replay,
160        };
161        result.validate()?;
162        Ok(result)
163    }
164
165    fn validate(&self) -> Result<(), String> {
166        if self.terminal_candidate && self.terminal_replay {
167            return Err(
168                "CycleDispatchResult terminal_candidate and terminal_replay are mutually exclusive"
169                    .to_string(),
170            );
171        }
172        if (self.terminal_candidate || self.terminal_replay)
173            && (!self.finished || self.result.is_none() || self.checkpoint_revision.is_none())
174        {
175            return Err(
176                "CycleDispatchResult terminal disposition requires a finished result and checkpoint_revision"
177                    .to_string(),
178            );
179        }
180        if self.committed_cycle.is_some()
181            && (self.finished
182                || self.result.is_some()
183                || self.terminal_candidate
184                || self.terminal_replay)
185        {
186            return Err(
187                "CycleDispatchResult committed_cycle is only valid for unfinished progress"
188                    .to_string(),
189            );
190        }
191        Ok(())
192    }
193}
194
195fn optional_bool(object: &serde_json::Map<String, Value>, field: &str) -> Result<bool, String> {
196    match object.get(field) {
197        None => Ok(false),
198        Some(Value::Bool(value)) => Ok(*value),
199        Some(_) => Err(format!("CycleDispatchResult {field} must be a boolean")),
200    }
201}
202
203pub trait CycleDispatcher: Send + Sync {
204    fn dispatch_cycle(
205        &self,
206        task: &AgentTask,
207        recipe: &RuntimeRecipe,
208        cycle_name: &str,
209        cycle_index: u32,
210    ) -> Result<CycleDispatchResult, String>;
211
212    fn dispatch_envelope(
213        &self,
214        envelope: &DistributedRunEnvelope,
215    ) -> Result<CycleDispatchResult, String> {
216        self.dispatch_cycle(
217            &envelope.task,
218            &envelope.recipe,
219            &envelope.cycle_name,
220            envelope.cycle_index,
221        )
222    }
223
224    fn dispatch_envelope_with_cancellation(
225        &self,
226        envelope: &DistributedRunEnvelope,
227        cancellation_token: Option<&CancellationToken>,
228    ) -> Result<CycleDispatchResult, String> {
229        check_cancellation(cancellation_token)?;
230        let result = self.dispatch_envelope(envelope)?;
231        if result.finished {
232            return Ok(result);
233        }
234        check_cancellation(cancellation_token)?;
235        Ok(result)
236    }
237}
238
239fn check_cancellation(cancellation_token: Option<&CancellationToken>) -> Result<(), String> {
240    cancellation_token
241        .map(CancellationToken::check)
242        .transpose()
243        .map(|_| ())
244        .map_err(|reason| format!("distributed dispatch cancelled: {reason}"))
245}
vv_agent/runtime/backends/distributed/dispatch.rs

vv_agent/runtime/backends/distributed/
dispatch.rs
