#[cfg(feature = "postgres")]
pub(crate) mod postgres;
#[cfg(feature = "sqlite")]
pub(crate) mod sqlite;
use std::time::Duration;
use a2a_protocol_types::task::TaskState;
const TERMINAL_STATES: [TaskState; 4] = [
TaskState::Completed,
TaskState::Failed,
TaskState::Canceled,
TaskState::Rejected,
];
#[must_use]
pub fn terminal_states() -> Vec<TaskState> {
TERMINAL_STATES.to_vec()
}
#[derive(Debug, Clone)]
pub struct RetentionPolicy {
pub terminal_max_age: Duration,
pub batch_size: u32,
pub max_batches: Option<u32>,
}
impl RetentionPolicy {
#[must_use]
pub const fn new(terminal_max_age: Duration) -> Self {
Self {
terminal_max_age,
batch_size: 1_000,
max_batches: None,
}
}
#[must_use]
pub const fn with_batch_size(mut self, batch_size: u32) -> Self {
self.batch_size = batch_size;
self
}
#[must_use]
pub const fn with_max_batches(mut self, max_batches: u32) -> Self {
self.max_batches = Some(max_batches);
self
}
#[cfg(any(feature = "sqlite", feature = "postgres"))]
#[must_use]
pub(crate) const fn effective_batch_size(&self) -> u32 {
if self.batch_size == 0 {
1
} else {
self.batch_size
}
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct PurgeReport {
pub tasks_deleted: u64,
pub journal_orphans_deleted: u64,
pub batches: u32,
pub complete: bool,
}
#[cfg(any(feature = "sqlite", feature = "postgres"))]
pub(crate) fn terminal_state_labels() -> Vec<String> {
terminal_states().iter().map(TaskState::to_string).collect()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn terminal_states_matches_is_terminal_for_every_variant() {
let purgeable = terminal_states();
for state in TaskState::ALL {
assert_eq!(
purgeable.contains(&state),
state.is_terminal(),
"{state} must be purgeable exactly when it is terminal"
);
}
assert_eq!(
purgeable.len(),
TaskState::ALL.iter().filter(|s| s.is_terminal()).count(),
"TERMINAL_STATES has drifted from is_terminal(); a new terminal \
state would otherwise never be purged"
);
}
#[test]
#[cfg(any(feature = "sqlite", feature = "postgres"))]
fn labels_are_the_stored_spellings() {
let labels = terminal_state_labels();
assert!(labels.contains(&"TASK_STATE_COMPLETED".to_string()));
assert!(labels.contains(&"TASK_STATE_REJECTED".to_string()));
assert_eq!(labels.len(), terminal_states().len());
assert_eq!(labels[0], TaskState::Completed.to_string());
}
#[test]
#[cfg(any(feature = "sqlite", feature = "postgres"))]
fn zero_batch_size_cannot_stall_a_sweep() {
let policy = RetentionPolicy::new(Duration::from_secs(1)).with_batch_size(0);
assert_eq!(
policy.effective_batch_size(),
1,
"a zero batch size would delete nothing while looping forever"
);
}
#[test]
fn builders_compose() {
let policy = RetentionPolicy::new(Duration::from_secs(60))
.with_batch_size(50)
.with_max_batches(3);
assert_eq!(policy.terminal_max_age, Duration::from_secs(60));
assert_eq!(policy.batch_size, 50);
assert_eq!(policy.max_batches, Some(3));
}
#[test]
fn report_defaults_to_nothing_done_but_complete() {
let report = PurgeReport::default();
assert_eq!(report.tasks_deleted, 0);
assert!(!report.complete, "default must not claim a completed sweep");
}
}