use async_trait::async_trait;
use serde_json::json;
use super::{Tool, ToolContext, ToolOutput};
use crate::error::Result;
use crate::memory::MemoryStore;
use crate::storage;
use crate::trust::RiskLabel;
const DEFAULT_MEMORY_LIMIT: usize = 2200;
const DEFAULT_USER_LIMIT: usize = 1400;
pub struct MemoryTool;
#[async_trait]
impl Tool for MemoryTool {
fn name(&self) -> &str {
"memory"
}
fn label(&self) -> &str {
"Memory"
}
fn description(&self) -> &str {
"Manage persistent memory across sessions. Use to save environment facts, \
user preferences, and lessons learned. Target 'memory' for agent notes, \
'user' for user profile."
}
fn parameters(&self) -> serde_json::Value {
json!({
"type": "object",
"required": ["action", "target"],
"properties": {
"action": {
"type": "string",
"enum": ["add", "replace", "remove"],
"description": "Action to perform"
},
"target": {
"type": "string",
"enum": ["memory", "user"],
"description": "Which store: 'memory' for agent notes, 'user' for user profile"
},
"content": {
"type": "string",
"description": "Content to add or replacement text"
},
"old_text": {
"type": "string",
"description": "Unique substring identifying the entry to replace or remove"
}
}
})
}
fn is_readonly(&self) -> bool {
false
}
async fn execute(
&self,
_call_id: &str,
params: serde_json::Value,
ctx: ToolContext,
) -> Result<ToolOutput> {
let action = params["action"].as_str().unwrap_or("");
let target = params["target"].as_str().unwrap_or("");
if action.is_empty() {
return Ok(ToolOutput::error("Missing required parameter: action"));
}
if target.is_empty() {
return Ok(ToolOutput::error("Missing required parameter: target"));
}
let (path, char_limit) = match target {
"memory" => (storage::global_memory_path(), DEFAULT_MEMORY_LIMIT),
"user" => (storage::global_user_path(), DEFAULT_USER_LIMIT),
other => {
return Ok(ToolOutput::error(format!(
"Unknown target \"{other}\". Use \"memory\" or \"user\"."
)));
}
};
let mut store = match MemoryStore::load(&path, char_limit) {
Ok(s) => s,
Err(e) => return Ok(ToolOutput::error(format!("Failed to load memory: {e}"))),
};
let result = match action {
"add" => {
let content = params["content"].as_str().unwrap_or("");
if content.is_empty() {
return Ok(ToolOutput::error(
"Missing required parameter: content (for 'add' action)",
));
}
if ctx
.supporting_provenance
.iter()
.any(|provenance| provenance.is_low_trust())
{
return Ok(ToolOutput::error(
"Low-trust context cannot be written to durable memory without explicit user adoption.",
));
}
if ctx.supporting_provenance.iter().any(|provenance| {
provenance.risk.iter().any(|risk| {
matches!(
risk,
RiskLabel::PossiblePromptInjection | RiskLabel::SecretAdjacent
)
})
}) {
return Ok(ToolOutput::error(
"Risk-labeled context cannot be written to durable memory without review.",
));
}
store.add(content)?
}
"replace" => {
let old_text = params["old_text"].as_str().unwrap_or("");
let content = params["content"].as_str().unwrap_or("");
if old_text.is_empty() {
return Ok(ToolOutput::error(
"Missing required parameter: old_text (for 'replace' action)",
));
}
if content.is_empty() {
return Ok(ToolOutput::error(
"Missing required parameter: content (for 'replace' action)",
));
}
if ctx
.supporting_provenance
.iter()
.any(|provenance| provenance.is_low_trust())
{
return Ok(ToolOutput::error(
"Low-trust context cannot replace durable memory without explicit user adoption.",
));
}
store.replace(old_text, content)?
}
"remove" => {
let old_text = params["old_text"].as_str().unwrap_or("");
if old_text.is_empty() {
return Ok(ToolOutput::error(
"Missing required parameter: old_text (for 'remove' action)",
));
}
store.remove(old_text)?
}
other => {
return Ok(ToolOutput::error(format!(
"Unknown action \"{other}\". Use \"add\", \"replace\", or \"remove\"."
)));
}
};
let json_text = serde_json::to_string_pretty(&result.to_json()).unwrap_or_default();
if result.success {
Ok(ToolOutput::text(json_text))
} else {
Ok(ToolOutput::error(json_text))
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::tools::ToolContext;
use std::sync::Arc;
fn test_ctx() -> ToolContext {
let (tx, _rx) = tokio::sync::mpsc::channel(16);
let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::channel(16);
let dir = std::env::temp_dir();
ToolContext {
cwd: dir,
cancelled: Arc::new(std::sync::atomic::AtomicBool::new(false)),
update_tx: tx,
command_tx: cmd_tx,
ui: Arc::new(crate::ui::NullInterface),
file_cache: Arc::new(crate::tools::FileCache::new()),
checkpoint_state: Arc::new(crate::tools::CheckpointState::new()),
file_tracker: Arc::new(std::sync::Mutex::new(crate::tools::FileTracker::new())),
anchor_store: Arc::new(crate::tools::AnchorStore::new()),
lua_tool_loader: None,
mode: crate::config::AgentMode::Full,
read_max_lines: 500,
turn_mana_review: Arc::new(std::sync::Mutex::new(
crate::mana_review::TurnManaReviewAccumulator::default(),
)),
config: Arc::new(crate::config::Config::default()),
run_policy: Default::default(),
supporting_provenance: Vec::new(),
}
}
#[tokio::test]
async fn memory_tool_validates_params() {
let tool = MemoryTool;
let r = tool
.execute("c1", json!({"target": "memory"}), test_ctx())
.await
.unwrap();
assert!(r.is_error);
let r = tool
.execute("c2", json!({"action": "add"}), test_ctx())
.await
.unwrap();
assert!(r.is_error);
let r = tool
.execute(
"c3",
json!({"action": "add", "target": "memory"}),
test_ctx(),
)
.await
.unwrap();
assert!(r.is_error);
}
#[tokio::test]
async fn memory_tool_blocks_low_trust_durable_writes() {
let tool = MemoryTool;
let mut ctx = test_ctx();
ctx.supporting_provenance
.push(crate::trust::Provenance::external_web(
"https://example.com",
));
let result = tool
.execute(
"c4",
json!({"action": "add", "target": "memory", "content": "remember this"}),
ctx,
)
.await
.unwrap();
assert!(result.is_error);
assert!(result.content.iter().any(|block| {
matches!(block, imp_llm::ContentBlock::Text { text } if text.contains("Low-trust context"))
}));
}
}