use crate::config::{CleanupStrategy, RetentionConfig};
use crate::io::inbox::inbox_update;
use crate::schema::InboxMessage;
use anyhow::{Context, Result};
use chrono::{DateTime, Duration, Utc};
use std::fs;
use std::path::{Path, PathBuf};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RetentionResult {
pub kept: usize,
pub removed: usize,
pub archived: usize,
}
impl RetentionResult {
pub fn new(kept: usize, removed: usize, archived: usize) -> Self {
Self {
kept,
removed,
archived,
}
}
}
pub fn apply_retention(
inbox_path: &Path,
team: &str,
agent: &str,
policy: &RetentionConfig,
dry_run: bool,
) -> Result<RetentionResult> {
if !inbox_path.exists() {
return Ok(RetentionResult::new(0, 0, 0));
}
let content = fs::read_to_string(inbox_path)
.with_context(|| format!("Failed to read inbox at {}", inbox_path.display()))?;
let messages: Vec<InboxMessage> = serde_json::from_str(&content)
.with_context(|| format!("Failed to parse inbox at {}", inbox_path.display()))?;
if policy.max_age.is_none() && policy.max_count.is_none() {
return Ok(RetentionResult::new(messages.len(), 0, 0));
}
let now = Utc::now();
let max_age_duration = if let Some(ref age_str) = policy.max_age {
Some(parse_duration(age_str)?)
} else {
None
};
let mut to_keep = Vec::new();
let mut to_remove = Vec::new();
for message in messages {
let should_remove = should_remove_message(
&message,
&max_age_duration,
now,
policy.max_count,
to_keep.len(),
);
if should_remove {
to_remove.push(message);
} else {
to_keep.push(message);
}
}
if to_remove.is_empty() {
return Ok(RetentionResult::new(to_keep.len(), 0, 0));
}
if dry_run {
let archived = if policy.strategy == CleanupStrategy::Archive {
to_remove.len()
} else {
0
};
return Ok(RetentionResult::new(
to_keep.len(),
to_remove.len(),
archived,
));
}
let archived = if policy.strategy == CleanupStrategy::Archive {
let archive_dir = determine_archive_dir(policy)?;
archive_messages(&to_remove, team, agent, &archive_dir)?;
to_remove.len()
} else {
0
};
inbox_update(inbox_path, team, agent, |messages| {
messages.clear();
messages.extend(to_keep.clone());
})?;
Ok(RetentionResult::new(
to_keep.len(),
to_remove.len(),
archived,
))
}
fn should_remove_message(
message: &InboxMessage,
max_age_duration: &Option<Duration>,
now: DateTime<Utc>,
max_count: Option<usize>,
current_kept_count: usize,
) -> bool {
if let Some(max_age) = max_age_duration
&& is_expired_by_age(message, max_age, now)
{
return true;
}
if let Some(max_count) = max_count
&& current_kept_count >= max_count
{
return true;
}
false
}
fn is_expired_by_age(message: &InboxMessage, max_age: &Duration, now: DateTime<Utc>) -> bool {
if let Ok(msg_time) = DateTime::parse_from_rfc3339(&message.timestamp) {
let msg_time_utc = msg_time.with_timezone(&Utc);
let age = now.signed_duration_since(msg_time_utc);
age > *max_age
} else {
true
}
}
pub fn parse_duration(s: &str) -> Result<Duration> {
let s = s.trim();
if s.is_empty() {
anyhow::bail!("Empty duration string");
}
let (num_part, unit) = match s.find(|c: char| !c.is_ascii_digit()) {
Some(idx) => (&s[..idx], &s[idx..]),
None => anyhow::bail!("Duration must have a unit (h or d): {s}"),
};
let num: i64 = num_part
.parse()
.with_context(|| format!("Invalid number in duration: {s}"))?;
match unit {
"h" => Ok(Duration::hours(num)),
"d" => Ok(Duration::days(num)),
_ => anyhow::bail!("Unknown duration unit '{unit}'. Use 'h' for hours or 'd' for days"),
}
}
fn determine_archive_dir(policy: &RetentionConfig) -> Result<PathBuf> {
if let Some(ref dir_str) = policy.archive_dir {
Ok(PathBuf::from(dir_str))
} else {
let home = crate::home::get_home_dir()?;
Ok(home.join(".config/atm/archive"))
}
}
fn archive_messages(
messages: &[InboxMessage],
team: &str,
agent: &str,
archive_dir: &Path,
) -> Result<()> {
let team_agent_dir = archive_dir.join(team).join(agent);
fs::create_dir_all(&team_agent_dir).with_context(|| {
format!(
"Failed to create archive directory: {}",
team_agent_dir.display()
)
})?;
let timestamp = Utc::now().format("%Y%m%d-%H%M%S");
let archive_file = team_agent_dir.join(format!("archive-{timestamp}.json"));
let json = serde_json::to_string_pretty(messages)
.context("Failed to serialize messages for archiving")?;
fs::write(&archive_file, json)
.with_context(|| format!("Failed to write archive file: {}", archive_file.display()))?;
Ok(())
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CleanReportResult {
pub deleted_count: usize,
pub skipped_count: usize,
}
impl CleanReportResult {
pub fn new(deleted_count: usize, skipped_count: usize) -> Self {
Self {
deleted_count,
skipped_count,
}
}
}
pub fn clean_report_files(report_dir: &Path, max_age: &Duration) -> Result<CleanReportResult> {
if !report_dir.exists() {
return Ok(CleanReportResult::new(0, 0));
}
let now = Utc::now();
let mut deleted = 0;
let mut skipped = 0;
let entries = fs::read_dir(report_dir)
.with_context(|| format!("Failed to read report directory: {}", report_dir.display()))?;
for entry in entries {
let entry = entry.context("Failed to read directory entry")?;
let path = entry.path();
if let Some(ext) = path.extension() {
let ext_str = ext.to_string_lossy();
if ext_str != "json" && ext_str != "md" {
continue;
}
} else {
continue;
}
let metadata = match fs::metadata(&path) {
Ok(m) => m,
Err(e) => {
tracing::warn!("Failed to get metadata for {}: {}", path.display(), e);
continue;
}
};
let modified = metadata
.modified()
.context("Failed to get file modification time")?;
let modified_datetime: DateTime<Utc> = modified.into();
let age = now.signed_duration_since(modified_datetime);
if age > *max_age {
match fs::remove_file(&path) {
Ok(()) => {
deleted += 1;
tracing::debug!("Deleted old report file: {}", path.display());
}
Err(e) => {
tracing::warn!("Failed to delete report file {}: {}", path.display(), e);
}
}
} else {
skipped += 1;
}
}
Ok(CleanReportResult::new(deleted, skipped))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse_duration_hours() {
let duration = parse_duration("24h").unwrap();
assert_eq!(duration, Duration::hours(24));
}
#[test]
fn test_parse_duration_days() {
let duration = parse_duration("7d").unwrap();
assert_eq!(duration, Duration::days(7));
}
#[test]
fn test_parse_duration_large_values() {
let duration = parse_duration("168h").unwrap();
assert_eq!(duration, Duration::hours(168));
let duration = parse_duration("30d").unwrap();
assert_eq!(duration, Duration::days(30));
}
#[test]
fn test_parse_duration_invalid() {
assert!(parse_duration("").is_err());
assert!(parse_duration("24").is_err());
assert!(parse_duration("24m").is_err());
assert!(parse_duration("abc").is_err());
}
#[test]
fn test_is_expired_by_age() {
let now = Utc::now();
let max_age = Duration::days(7);
let old_message = InboxMessage {
from: "test".to_string(),
text: "old message".to_string(),
timestamp: (now - Duration::days(10)).to_rfc3339(),
read: false,
summary: None,
message_id: None,
unknown_fields: std::collections::HashMap::new(),
};
assert!(is_expired_by_age(&old_message, &max_age, now));
let recent_message = InboxMessage {
from: "test".to_string(),
text: "recent message".to_string(),
timestamp: (now - Duration::days(3)).to_rfc3339(),
read: false,
summary: None,
message_id: None,
unknown_fields: std::collections::HashMap::new(),
};
assert!(!is_expired_by_age(&recent_message, &max_age, now));
}
#[test]
fn test_retention_result() {
let result = RetentionResult::new(10, 5, 5);
assert_eq!(result.kept, 10);
assert_eq!(result.removed, 5);
assert_eq!(result.archived, 5);
}
#[test]
fn test_clean_report_files_deletes_old_files() {
use std::fs::File;
use std::io::Write;
use tempfile::TempDir;
let temp_dir = TempDir::new().unwrap();
let report_dir = temp_dir.path().join("reports");
fs::create_dir_all(&report_dir).unwrap();
let old_json = report_dir.join("old.json");
let old_md = report_dir.join("old.md");
File::create(&old_json).unwrap().write_all(b"{}").unwrap();
File::create(&old_md)
.unwrap()
.write_all(b"# Report")
.unwrap();
std::thread::sleep(std::time::Duration::from_millis(500));
let max_age = Duration::milliseconds(200);
let result = super::clean_report_files(&report_dir, &max_age).unwrap();
assert_eq!(result.deleted_count, 2);
assert_eq!(result.skipped_count, 0);
assert!(!old_json.exists());
assert!(!old_md.exists());
}
#[test]
fn test_clean_report_files_skips_recent_files() {
use std::fs::File;
use std::io::Write;
use tempfile::TempDir;
let temp_dir = TempDir::new().unwrap();
let report_dir = temp_dir.path().join("reports");
fs::create_dir_all(&report_dir).unwrap();
let recent_json = report_dir.join("recent.json");
let recent_md = report_dir.join("recent.md");
File::create(&recent_json)
.unwrap()
.write_all(b"{}")
.unwrap();
File::create(&recent_md)
.unwrap()
.write_all(b"# Report")
.unwrap();
let max_age = Duration::hours(1);
let result = super::clean_report_files(&report_dir, &max_age).unwrap();
assert_eq!(result.deleted_count, 0);
assert_eq!(result.skipped_count, 2);
assert!(recent_json.exists());
assert!(recent_md.exists());
}
#[test]
fn test_clean_report_files_empty_directory() {
use tempfile::TempDir;
let temp_dir = TempDir::new().unwrap();
let report_dir = temp_dir.path().join("reports");
fs::create_dir_all(&report_dir).unwrap();
let max_age = Duration::hours(1);
let result = super::clean_report_files(&report_dir, &max_age).unwrap();
assert_eq!(result.deleted_count, 0);
assert_eq!(result.skipped_count, 0);
}
#[test]
fn test_clean_report_files_nonexistent_directory() {
use tempfile::TempDir;
let temp_dir = TempDir::new().unwrap();
let report_dir = temp_dir.path().join("nonexistent");
let max_age = Duration::hours(1);
let result = super::clean_report_files(&report_dir, &max_age).unwrap();
assert_eq!(result.deleted_count, 0);
assert_eq!(result.skipped_count, 0);
}
#[test]
fn test_clean_report_files_only_targets_json_and_md() {
use std::fs::File;
use std::io::Write;
use tempfile::TempDir;
let temp_dir = TempDir::new().unwrap();
let report_dir = temp_dir.path().join("reports");
fs::create_dir_all(&report_dir).unwrap();
let json_file = report_dir.join("report.json");
let md_file = report_dir.join("report.md");
let txt_file = report_dir.join("report.txt");
let log_file = report_dir.join("report.log");
File::create(&json_file).unwrap().write_all(b"{}").unwrap();
File::create(&md_file)
.unwrap()
.write_all(b"# Report")
.unwrap();
File::create(&txt_file).unwrap().write_all(b"text").unwrap();
File::create(&log_file).unwrap().write_all(b"log").unwrap();
std::thread::sleep(std::time::Duration::from_millis(500));
let max_age = Duration::milliseconds(200);
let result = super::clean_report_files(&report_dir, &max_age).unwrap();
assert_eq!(result.deleted_count, 2);
assert!(!json_file.exists());
assert!(!md_file.exists());
assert!(txt_file.exists()); assert!(log_file.exists()); }
#[test]
fn test_clean_report_result() {
let result = CleanReportResult::new(5, 10);
assert_eq!(result.deleted_count, 5);
assert_eq!(result.skipped_count, 10);
}
#[test]
fn test_parse_inbox_filename_simple() {
let filename = "agent.json";
let base = &filename[..filename.len() - 5];
assert_eq!(base, "agent");
}
#[test]
fn test_parse_inbox_filename_with_hostname() {
let filename = "agent.hostname.json";
let base = &filename[..filename.len() - 5];
assert_eq!(base, "agent.hostname");
}
#[test]
fn test_parse_inbox_filename_dotted_agent() {
let filename = "dotted.agent.name.json";
let base = &filename[..filename.len() - 5];
assert_eq!(base, "dotted.agent.name");
}
#[test]
fn test_parse_inbox_filename_dotted_agent_with_hostname() {
let filename = "dotted.agent.hostname.json";
let base = &filename[..filename.len() - 5];
assert_eq!(base, "dotted.agent.hostname");
}
#[test]
fn test_retention_only_scans_inboxes_subdirectory() {
use std::fs::File;
use std::io::Write;
use tempfile::TempDir;
let temp_dir = TempDir::new().unwrap();
let team_dir = temp_dir.path().join("test-team");
let inboxes_dir = team_dir.join("inboxes");
fs::create_dir_all(&inboxes_dir).unwrap();
let inbox_path = inboxes_dir.join("agent.json");
let old_timestamp = (Utc::now() - Duration::days(10)).to_rfc3339();
let inbox_data = serde_json::json!([
{
"from": "test",
"text": "old message",
"timestamp": old_timestamp,
"read": false
}
]);
File::create(&inbox_path)
.unwrap()
.write_all(serde_json::to_string(&inbox_data).unwrap().as_bytes())
.unwrap();
let config_path = team_dir.join("config.json");
File::create(&config_path)
.unwrap()
.write_all(b"{\"name\":\"test-team\"}")
.unwrap();
let policy = RetentionConfig {
max_age: Some("7d".to_string()),
max_count: None,
strategy: CleanupStrategy::Delete,
archive_dir: None,
enabled: true,
interval_secs: 300,
};
let result = apply_retention(&inbox_path, "test-team", "agent", &policy, false).unwrap();
assert_eq!(result.kept, 0);
assert_eq!(result.removed, 1);
assert!(config_path.exists());
let config_content = fs::read_to_string(&config_path).unwrap();
assert_eq!(config_content, "{\"name\":\"test-team\"}");
}
#[test]
fn test_report_dir_lookup_resolves_ci_monitor_key() {
use tempfile::TempDir;
let temp_dir = TempDir::new().unwrap();
let report_dir = temp_dir.path().join("ci-reports");
fs::create_dir_all(&report_dir).unwrap();
let report_path = report_dir.join("report.json");
std::fs::File::create(&report_path).unwrap();
let max_age = Duration::hours(24);
let result = clean_report_files(&report_dir, &max_age).unwrap();
assert_eq!(result.deleted_count, 0);
assert_eq!(result.skipped_count, 1);
}
}