use chrono::{DateTime, Duration, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use crate::agency::{Agent, AgentBuilder, AgentConfig, AgentRole, AgentStatus};
use crate::database::ChatDatabase;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ArchivalPolicy {
pub name: String,
pub enabled: bool,
pub inactive_days: u32,
pub min_messages: u32,
pub max_messages: Option<u32>,
pub providers: Vec<String>,
pub workspace_ids: Vec<String>,
pub exclude_tags: Vec<String>,
pub include_tags: Vec<String>,
pub compress: bool,
pub notify: bool,
}
impl Default for ArchivalPolicy {
fn default() -> Self {
Self {
name: "default".to_string(),
enabled: true,
inactive_days: 30,
min_messages: 5,
max_messages: None,
providers: vec![],
workspace_ids: vec![],
exclude_tags: vec!["pinned".to_string(), "important".to_string()],
include_tags: vec!["archive".to_string()],
compress: true,
notify: true,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ArchivalCandidate {
pub session_id: String,
pub title: String,
pub provider: String,
pub workspace_id: Option<String>,
pub message_count: u32,
pub last_activity: DateTime<Utc>,
pub days_inactive: u32,
pub policy: String,
pub reason: String,
pub priority: u8,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ArchivalDecision {
pub session_id: String,
pub should_archive: bool,
pub confidence: f64,
pub reasoning: String,
pub policies: Vec<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ArchivalResult {
pub archived_count: u32,
pub skipped_count: u32,
pub bytes_saved: u64,
pub errors: Vec<String>,
pub timestamp: DateTime<Utc>,
pub duration_ms: u64,
}
pub struct ArchivalAgentState {
policies: Vec<ArchivalPolicy>,
last_run: Option<DateTime<Utc>>,
stats: ArchivalStats,
pending: Vec<ArchivalCandidate>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ArchivalStats {
pub total_runs: u64,
pub total_archived: u64,
pub total_skipped: u64,
pub total_bytes_saved: u64,
pub avg_confidence: f64,
}
pub struct ArchivalAgent {
config: AgentConfig,
state: Arc<RwLock<ArchivalAgentState>>,
db: Option<Arc<ChatDatabase>>,
running: Arc<RwLock<bool>>,
}
impl ArchivalAgent {
pub fn new() -> Self {
let config = AgentConfig {
name: "archival-agent".to_string(),
description: "Autonomous session archival agent".to_string(),
instruction: ARCHIVAL_SYSTEM_PROMPT.to_string(),
..Default::default()
};
let state = ArchivalAgentState {
policies: vec![ArchivalPolicy::default()],
last_run: None,
stats: ArchivalStats::default(),
pending: vec![],
};
Self {
config,
state: Arc::new(RwLock::new(state)),
db: None,
running: Arc::new(RwLock::new(false)),
}
}
pub fn with_policies(policies: Vec<ArchivalPolicy>) -> Self {
let agent = Self::new();
let mut state = agent.state.blocking_write();
state.policies = policies;
drop(state);
agent
}
pub fn with_database(mut self, db: Arc<ChatDatabase>) -> Self {
self.db = Some(db);
self
}
pub async fn add_policy(&self, policy: ArchivalPolicy) {
let mut state = self.state.write().await;
state.policies.push(policy);
}
pub async fn remove_policy(&self, name: &str) -> bool {
let mut state = self.state.write().await;
let len_before = state.policies.len();
state.policies.retain(|p| p.name != name);
state.policies.len() < len_before
}
pub async fn get_policies(&self) -> Vec<ArchivalPolicy> {
let state = self.state.read().await;
state.policies.clone()
}
pub async fn scan_candidates(&self) -> Vec<ArchivalCandidate> {
let state = self.state.read().await;
let candidates = Vec::new();
let _now = Utc::now();
for policy in &state.policies {
if !policy.enabled {
continue;
}
}
candidates
}
pub async fn evaluate_session(&self, session_id: &str) -> ArchivalDecision {
let state = self.state.read().await;
let mut matched_policies = Vec::new();
let mut reasons = Vec::new();
let mut should_archive = false;
let mut confidence = 0.0;
for policy in &state.policies {
if !policy.enabled {
continue;
}
matched_policies.push(policy.name.clone());
}
if !matched_policies.is_empty() {
should_archive = true;
confidence = 0.85;
reasons.push("Matched archival policies".to_string());
}
ArchivalDecision {
session_id: session_id.to_string(),
should_archive,
confidence,
reasoning: reasons.join("; "),
policies: matched_policies,
}
}
pub async fn archive_session(&self, _session_id: &str) -> Result<bool, String> {
let mut state = self.state.write().await;
state.stats.total_archived += 1;
Ok(true)
}
pub async fn run(&self) -> ArchivalResult {
let start = std::time::Instant::now();
let mut result = ArchivalResult {
archived_count: 0,
skipped_count: 0,
bytes_saved: 0,
errors: vec![],
timestamp: Utc::now(),
duration_ms: 0,
};
{
let mut running = self.running.write().await;
if *running {
result.errors.push("Agent already running".to_string());
return result;
}
*running = true;
}
let candidates = self.scan_candidates().await;
for candidate in candidates {
let decision = self.evaluate_session(&candidate.session_id).await;
if decision.should_archive && decision.confidence >= 0.7 {
match self.archive_session(&candidate.session_id).await {
Ok(true) => {
result.archived_count += 1;
}
Ok(false) => {
result.skipped_count += 1;
}
Err(e) => {
result
.errors
.push(format!("Failed to archive {}: {}", candidate.session_id, e));
}
}
} else {
result.skipped_count += 1;
}
}
{
let mut state = self.state.write().await;
state.last_run = Some(Utc::now());
state.stats.total_runs += 1;
state.stats.total_bytes_saved += result.bytes_saved;
}
{
let mut running = self.running.write().await;
*running = false;
}
result.duration_ms = start.elapsed().as_millis() as u64;
result
}
pub async fn get_stats(&self) -> ArchivalStats {
let state = self.state.read().await;
state.stats.clone()
}
pub async fn get_last_run(&self) -> Option<DateTime<Utc>> {
let state = self.state.read().await;
state.last_run
}
pub async fn is_running(&self) -> bool {
let running = self.running.read().await;
*running
}
pub async fn stop(&self) {
let mut running = self.running.write().await;
*running = false;
}
}
impl Default for ArchivalAgent {
fn default() -> Self {
Self::new()
}
}
const ARCHIVAL_SYSTEM_PROMPT: &str = r#"You are an autonomous session archival agent for Chasm.
Your role is to analyze chat sessions and determine which should be archived based on:
1. Inactivity period (days since last message)
2. Session size and importance
3. Content relevance and quality
4. User-defined policies and tags
When evaluating a session for archival, consider:
- Is the conversation complete or ongoing?
- Does it contain important information that should be preserved?
- Are there pinned or important tags?
- How much space would archiving save?
Provide clear reasoning for your archival decisions.
"#;
pub struct ArchivalScheduler {
agent: Arc<ArchivalAgent>,
interval: Duration,
active: Arc<RwLock<bool>>,
}
impl ArchivalScheduler {
pub fn new(agent: Arc<ArchivalAgent>, interval_hours: u32) -> Self {
Self {
agent,
interval: Duration::hours(interval_hours as i64),
active: Arc::new(RwLock::new(false)),
}
}
pub async fn start(&self) {
let mut active = self.active.write().await;
*active = true;
drop(active);
println!(
"[ArchivalScheduler] Started with interval {:?}. Call run() to execute.",
self.interval
);
}
pub async fn stop(&self) {
let mut active = self.active.write().await;
*active = false;
}
pub async fn is_active(&self) -> bool {
let active = self.active.read().await;
*active
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_archival_agent_creation() {
let agent = ArchivalAgent::new();
let policies = agent.get_policies().await;
assert_eq!(policies.len(), 1);
assert_eq!(policies[0].name, "default");
}
#[tokio::test]
async fn test_add_remove_policy() {
let agent = ArchivalAgent::new();
let custom_policy = ArchivalPolicy {
name: "aggressive".to_string(),
inactive_days: 7,
..Default::default()
};
agent.add_policy(custom_policy).await;
let policies = agent.get_policies().await;
assert_eq!(policies.len(), 2);
agent.remove_policy("aggressive").await;
let policies = agent.get_policies().await;
assert_eq!(policies.len(), 1);
}
#[tokio::test]
async fn test_evaluate_session() {
let agent = ArchivalAgent::new();
let decision = agent.evaluate_session("test-session-123").await;
assert!(!decision.session_id.is_empty());
}
}