use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PendingAction {
pub action_id: String,
pub tool_name: String,
pub arguments: serde_json::Value,
pub issued_at: u64,
pub expected_outcome: ExpectedOutcome,
pub status: PendingActionStatus,
}
impl PendingAction {
pub fn new(
action_id: String,
tool_name: String,
arguments: serde_json::Value,
expected_outcome: ExpectedOutcome,
) -> Self {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0);
Self {
action_id,
tool_name,
arguments,
issued_at: now,
expected_outcome,
status: PendingActionStatus::InFlight,
}
}
pub fn is_in_flight(&self) -> bool {
matches!(self.status, PendingActionStatus::InFlight)
}
pub fn is_stale(&self, timeout_secs: u64) -> bool {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0);
matches!(self.status, PendingActionStatus::InFlight) && now.saturating_sub(self.issued_at) > timeout_secs
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PendingActionStatus {
InFlight,
Completed {
completed_at: u64,
success: bool,
},
Failed {
completed_at: u64,
error: String,
},
TimedOut,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ExpectedOutcome {
FileModification {
paths: Vec<String>,
},
ReadOperation,
CommandExecution,
SearchOperation,
Generic,
}
impl ExpectedOutcome {
pub fn has_side_effects(&self) -> bool {
matches!(self, ExpectedOutcome::FileModification { .. } | ExpectedOutcome::CommandExecution)
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct PendingActions {
actions: Vec<PendingAction>,
max_retained: usize,
}
impl PendingActions {
pub fn new(max_retained: usize) -> Self {
Self {
actions: Vec::with_capacity(max_retained.saturating_add(16)),
max_retained,
}
}
pub fn register(&mut self, action: PendingAction) {
self.actions.push(action);
}
pub fn resolve(&mut self, action_id: &str, success: bool, error: Option<String>) -> bool {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0);
if let Some(action) = self
.actions
.iter_mut()
.find(|a| a.action_id == action_id && matches!(a.status, PendingActionStatus::InFlight))
{
action.status = match error {
Some(err) => PendingActionStatus::Failed { completed_at: now, error: err },
None => PendingActionStatus::Completed { completed_at: now, success },
};
self.prune();
true
} else {
false
}
}
pub fn mark_timed_out(&mut self, action_id: &str) -> bool {
if let Some(action) = self
.actions
.iter_mut()
.find(|a| a.action_id == action_id && matches!(a.status, PendingActionStatus::InFlight))
{
action.status = PendingActionStatus::TimedOut;
self.prune();
true
} else {
false
}
}
pub fn in_flight(&self) -> impl Iterator<Item = &PendingAction> {
self.actions.iter().filter(|a| a.is_in_flight())
}
pub fn stale(&self, timeout_secs: u64) -> impl Iterator<Item = &PendingAction> {
self.actions.iter().filter(move |a| a.is_stale(timeout_secs))
}
pub fn in_flight_count(&self) -> usize {
self.in_flight().count()
}
pub fn contains(&self, action_id: &str) -> bool {
self.actions.iter().any(|a| a.action_id == action_id)
}
fn prune(&mut self) {
if self.max_retained == 0 {
return;
}
let terminal_indices: Vec<usize> = self
.actions
.iter()
.enumerate()
.filter(|(_, a)| !a.is_in_flight())
.map(|(i, _)| i)
.collect();
if terminal_indices.len() > self.max_retained {
let to_remove = terminal_indices.len() - self.max_retained;
for idx in terminal_indices.into_iter().take(to_remove).rev() {
self.actions.remove(idx);
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn make_action(action_id: &str) -> PendingAction {
PendingAction::new(
action_id.to_string(),
"test_tool".to_string(),
serde_json::json!({}),
ExpectedOutcome::Generic,
)
}
#[test]
fn test_register_and_resolve() {
let mut store = PendingActions::new(10);
store.register(make_action("call_1"));
assert_eq!(store.in_flight_count(), 1);
assert!(store.contains("call_1"));
assert!(store.resolve("call_1", true, None));
assert_eq!(store.in_flight_count(), 0);
assert!(store.contains("call_1"));
}
#[test]
fn test_resolve_with_error() {
let mut store = PendingActions::new(10);
store.register(make_action("call_1"));
assert!(store.resolve("call_1", false, Some("timeout".to_string())));
let actions: Vec<_> = store.in_flight().collect();
assert!(actions.is_empty());
}
#[test]
fn test_resolve_unknown_action() {
let mut store = PendingActions::new(10);
assert!(!store.resolve("unknown", true, None));
}
#[test]
fn test_stale_detection() {
let mut store = PendingActions::new(10);
store.register(PendingAction {
action_id: "stale_call".to_string(),
tool_name: "slow_tool".to_string(),
arguments: serde_json::json!({}),
issued_at: 0, expected_outcome: ExpectedOutcome::Generic,
status: PendingActionStatus::InFlight,
});
let stale: Vec<_> = store.stale(1).collect(); assert_eq!(stale.len(), 1);
}
#[test]
fn test_in_flight_count() {
let mut store = PendingActions::new(10);
store.register(make_action("call_1"));
store.register(make_action("call_2"));
assert_eq!(store.in_flight_count(), 2);
store.resolve("call_1", true, None);
assert_eq!(store.in_flight_count(), 1);
}
#[test]
fn test_expected_outcome_side_effects() {
assert!(ExpectedOutcome::FileModification { paths: vec!["a.txt".to_string()] }.has_side_effects());
assert!(ExpectedOutcome::CommandExecution.has_side_effects());
assert!(!ExpectedOutcome::ReadOperation.has_side_effects());
assert!(!ExpectedOutcome::SearchOperation.has_side_effects());
assert!(!ExpectedOutcome::Generic.has_side_effects());
}
#[test]
fn test_pending_action_serde_roundtrip() {
let action = make_action("call_roundtrip");
let json = serde_json::to_string(&action).unwrap();
let deserialized: PendingAction = serde_json::from_str(&json).unwrap();
assert_eq!(deserialized.action_id, "call_roundtrip");
assert!(deserialized.is_in_flight());
}
#[test]
fn test_prune_retains_in_flight() {
let mut store = PendingActions::new(1); store.register(make_action("keep"));
store.register(make_action("prune_me"));
store.resolve("prune_me", true, None);
store.resolve("keep", true, None);
assert_eq!(store.in_flight_count(), 0);
assert!(store.contains("prune_me"));
assert!(!store.contains("keep"));
}
}