telos_agent/tools/builtin/memory/
maintenance.rs

1use async_trait::async_trait;
2use serde_json::{Value, json};
3use std::sync::{Arc, Mutex};
4
5use crate::error::AgentError;
6use crate::knowledge::memory::{MemoryMaintenancePolicy, MemoryStore};
7use crate::tools::api::{PermissionDecision, Tool, ToolContext, ToolDefinition, ToolOutput};
8
9pub struct MemoryMaintenanceTool {
10    store: Arc<Mutex<MemoryStore>>,
11}
12
13impl MemoryMaintenanceTool {
14    pub fn new(store: Arc<Mutex<MemoryStore>>) -> Self {
15        Self { store }
16    }
17}
18
19#[async_trait]
20impl Tool for MemoryMaintenanceTool {
21    fn definition(&self) -> ToolDefinition {
22        ToolDefinition {
23            name: "MemoryMaintenance".into(),
24            description: "Preview or apply conservative memory cleanup by archiving stale low-value entries. Defaults to dry-run.".into(),
25            input_schema: json!({"type":"object","properties":{
26                "apply":{"type":"boolean","description":"When true, archive the reported candidates. Defaults to false."},
27                "archive_deprecated":{"type":"boolean","description":"Archive deprecated memories. Defaults to true."},
28                "max_auto_learned_commands":{"type":"integer","minimum":0,"description":"Keep at most this many active auto-learned command memories. Defaults to 20. Set null to disable this rule."},
29                "max_active_entries":{"type":"integer","minimum":0,"description":"Optional cap for active memories overall. Omit or null to disable."}
30            }}),
31        }
32    }
33
34    async fn check_permission(
35        &self,
36        _: &Value,
37        _: &ToolContext,
38    ) -> Result<PermissionDecision, AgentError> {
39        Ok(PermissionDecision::Allow)
40    }
41
42    async fn invoke(&self, args: Value, _: ToolContext) -> Result<ToolOutput, AgentError> {
43        let apply = args.get("apply").and_then(|v| v.as_bool()).unwrap_or(false);
44        let archive_deprecated =
45            args.get("archive_deprecated").and_then(|v| v.as_bool()).unwrap_or(true);
46        let max_auto_learned_commands = if args.get("max_auto_learned_commands").is_some() {
47            optional_usize_arg(&args, "max_auto_learned_commands")?
48        } else {
49            Some(20)
50        };
51        let max_active_entries = optional_usize_arg(&args, "max_active_entries")?;
52        let policy = MemoryMaintenancePolicy {
53            max_auto_learned_commands,
54            max_active_entries,
55            archive_deprecated,
56        };
57        let store = self.store.clone();
58        tokio::task::spawn_blocking(move || {
59            let mut store = store.lock().map_err(|e| AgentError::ToolExecution {
60                tool: "MemoryMaintenance".into(),
61                message: format!("memory store poisoned: {e}"),
62            })?;
63            let report = if apply {
64                store.apply_maintenance(&policy).map_err(|e| AgentError::ToolExecution {
65                    tool: "MemoryMaintenance".into(),
66                    message: e.to_string(),
67                })?
68            } else {
69                store.maintenance_report(&policy)
70            };
71            let content = serde_json::to_value(report).map_err(|e| AgentError::ToolExecution {
72                tool: "MemoryMaintenance".into(),
73                message: format!("failed to serialize maintenance report: {e}"),
74            })?;
75            Ok(ToolOutput::json(content))
76        })
77        .await
78        .map_err(|e| AgentError::ToolExecution {
79            tool: "MemoryMaintenance".into(),
80            message: format!("memory task panicked: {e}"),
81        })?
82    }
83}
84
85fn optional_usize_arg(args: &Value, name: &str) -> Result<Option<usize>, AgentError> {
86    let Some(value) = args.get(name) else {
87        return Ok(None);
88    };
89    if value.is_null() {
90        return Ok(None);
91    }
92    let Some(n) = value.as_u64() else {
93        return Err(AgentError::Validation(format!("`{name}` must be a non-negative integer")));
94    };
95    usize::try_from(n)
96        .map(Some)
97        .map_err(|_| AgentError::Validation(format!("`{name}` is too large")))
98}
telos_agent/tools/builtin/memory/maintenance.rs

telos_agent/tools/builtin/memory/
maintenance.rs
