1use serde::{Deserialize, Serialize};
15
16#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct PendingAction {
20 pub action_id: String,
22 pub tool_name: String,
24 pub arguments: serde_json::Value,
26 pub issued_at: u64,
28 pub expected_outcome: ExpectedOutcome,
30 pub status: PendingActionStatus,
32}
33
34impl PendingAction {
35 pub fn new(
37 action_id: String,
38 tool_name: String,
39 arguments: serde_json::Value,
40 expected_outcome: ExpectedOutcome,
41 ) -> Self {
42 let now = std::time::SystemTime::now()
43 .duration_since(std::time::UNIX_EPOCH)
44 .map(|d| d.as_secs())
45 .unwrap_or(0);
46 Self {
47 action_id,
48 tool_name,
49 arguments,
50 issued_at: now,
51 expected_outcome,
52 status: PendingActionStatus::InFlight,
53 }
54 }
55
56 pub fn is_in_flight(&self) -> bool {
58 matches!(self.status, PendingActionStatus::InFlight)
59 }
60
61 pub fn is_stale(&self, timeout_secs: u64) -> bool {
63 let now = std::time::SystemTime::now()
64 .duration_since(std::time::UNIX_EPOCH)
65 .map(|d| d.as_secs())
66 .unwrap_or(0);
67 matches!(self.status, PendingActionStatus::InFlight) && now.saturating_sub(self.issued_at) > timeout_secs
68 }
69}
70
71#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
73#[serde(rename_all = "snake_case")]
74pub enum PendingActionStatus {
75 InFlight,
77 Completed {
79 completed_at: u64,
81 success: bool,
83 },
84 Failed {
86 completed_at: u64,
88 error: String,
90 },
91 TimedOut,
93}
94
95#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
101#[serde(rename_all = "snake_case")]
102pub enum ExpectedOutcome {
103 FileModification {
105 paths: Vec<String>,
107 },
108 ReadOperation,
110 CommandExecution,
112 SearchOperation,
114 Generic,
116}
117
118impl ExpectedOutcome {
119 pub fn has_side_effects(&self) -> bool {
121 matches!(self, ExpectedOutcome::FileModification { .. } | ExpectedOutcome::CommandExecution)
122 }
123}
124
125#[derive(Debug, Clone, Default, Serialize, Deserialize)]
128pub struct PendingActions {
129 actions: Vec<PendingAction>,
131 max_retained: usize,
134}
135
136impl PendingActions {
137 pub fn new(max_retained: usize) -> Self {
139 Self {
140 actions: Vec::with_capacity(max_retained.saturating_add(16)),
141 max_retained,
142 }
143 }
144
145 pub fn register(&mut self, action: PendingAction) {
147 self.actions.push(action);
148 }
149
150 pub fn resolve(&mut self, action_id: &str, success: bool, error: Option<String>) -> bool {
153 let now = std::time::SystemTime::now()
154 .duration_since(std::time::UNIX_EPOCH)
155 .map(|d| d.as_secs())
156 .unwrap_or(0);
157 if let Some(action) = self
158 .actions
159 .iter_mut()
160 .find(|a| a.action_id == action_id && matches!(a.status, PendingActionStatus::InFlight))
161 {
162 action.status = match error {
163 Some(err) => PendingActionStatus::Failed { completed_at: now, error: err },
164 None => PendingActionStatus::Completed { completed_at: now, success },
165 };
166 self.prune();
167 true
168 } else {
169 false
170 }
171 }
172
173 pub fn mark_timed_out(&mut self, action_id: &str) -> bool {
175 if let Some(action) = self
176 .actions
177 .iter_mut()
178 .find(|a| a.action_id == action_id && matches!(a.status, PendingActionStatus::InFlight))
179 {
180 action.status = PendingActionStatus::TimedOut;
181 self.prune();
182 true
183 } else {
184 false
185 }
186 }
187
188 pub fn in_flight(&self) -> impl Iterator<Item = &PendingAction> {
190 self.actions.iter().filter(|a| a.is_in_flight())
191 }
192
193 pub fn stale(&self, timeout_secs: u64) -> impl Iterator<Item = &PendingAction> {
195 self.actions.iter().filter(move |a| a.is_stale(timeout_secs))
196 }
197
198 pub fn in_flight_count(&self) -> usize {
200 self.in_flight().count()
201 }
202
203 pub fn contains(&self, action_id: &str) -> bool {
206 self.actions.iter().any(|a| a.action_id == action_id)
207 }
208
209 fn prune(&mut self) {
211 if self.max_retained == 0 {
212 return;
213 }
214 let terminal_indices: Vec<usize> = self
216 .actions
217 .iter()
218 .enumerate()
219 .filter(|(_, a)| !a.is_in_flight())
220 .map(|(i, _)| i)
221 .collect();
222 if terminal_indices.len() > self.max_retained {
223 let to_remove = terminal_indices.len() - self.max_retained;
224 for idx in terminal_indices.into_iter().take(to_remove).rev() {
226 self.actions.remove(idx);
227 }
228 }
229 }
230}
231
232#[cfg(test)]
233mod tests {
234 use super::*;
235
236 fn make_action(action_id: &str) -> PendingAction {
237 PendingAction::new(
238 action_id.to_string(),
239 "test_tool".to_string(),
240 serde_json::json!({}),
241 ExpectedOutcome::Generic,
242 )
243 }
244
245 #[test]
246 fn test_register_and_resolve() {
247 let mut store = PendingActions::new(10);
248 store.register(make_action("call_1"));
249 assert_eq!(store.in_flight_count(), 1);
250 assert!(store.contains("call_1"));
251
252 assert!(store.resolve("call_1", true, None));
253 assert_eq!(store.in_flight_count(), 0);
254 assert!(store.contains("call_1"));
255 }
256
257 #[test]
258 fn test_resolve_with_error() {
259 let mut store = PendingActions::new(10);
260 store.register(make_action("call_1"));
261
262 assert!(store.resolve("call_1", false, Some("timeout".to_string())));
263 let actions: Vec<_> = store.in_flight().collect();
264 assert!(actions.is_empty());
265 }
266
267 #[test]
268 fn test_resolve_unknown_action() {
269 let mut store = PendingActions::new(10);
270 assert!(!store.resolve("unknown", true, None));
271 }
272
273 #[test]
274 fn test_stale_detection() {
275 let mut store = PendingActions::new(10);
276 store.register(PendingAction {
277 action_id: "stale_call".to_string(),
278 tool_name: "slow_tool".to_string(),
279 arguments: serde_json::json!({}),
280 issued_at: 0, expected_outcome: ExpectedOutcome::Generic,
282 status: PendingActionStatus::InFlight,
283 });
284
285 let stale: Vec<_> = store.stale(1).collect(); assert_eq!(stale.len(), 1);
287 }
288
289 #[test]
290 fn test_in_flight_count() {
291 let mut store = PendingActions::new(10);
292 store.register(make_action("call_1"));
293 store.register(make_action("call_2"));
294 assert_eq!(store.in_flight_count(), 2);
295
296 store.resolve("call_1", true, None);
297 assert_eq!(store.in_flight_count(), 1);
298 }
299
300 #[test]
301 fn test_expected_outcome_side_effects() {
302 assert!(ExpectedOutcome::FileModification { paths: vec!["a.txt".to_string()] }.has_side_effects());
303 assert!(ExpectedOutcome::CommandExecution.has_side_effects());
304 assert!(!ExpectedOutcome::ReadOperation.has_side_effects());
305 assert!(!ExpectedOutcome::SearchOperation.has_side_effects());
306 assert!(!ExpectedOutcome::Generic.has_side_effects());
307 }
308
309 #[test]
310 fn test_pending_action_serde_roundtrip() {
311 let action = make_action("call_roundtrip");
312 let json = serde_json::to_string(&action).unwrap();
313 let deserialized: PendingAction = serde_json::from_str(&json).unwrap();
314 assert_eq!(deserialized.action_id, "call_roundtrip");
315 assert!(deserialized.is_in_flight());
316 }
317
318 #[test]
319 fn test_prune_retains_in_flight() {
320 let mut store = PendingActions::new(1); store.register(make_action("keep"));
322 store.register(make_action("prune_me"));
323 store.resolve("prune_me", true, None);
324 store.resolve("keep", true, None);
325
326 assert_eq!(store.in_flight_count(), 0);
328 assert!(store.contains("prune_me"));
330 assert!(!store.contains("keep"));
332 }
333}