From 19d234048901e30bdb7fad6607721e4524df9c40 Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Wed, 10 Jun 2026 18:35:59 -0600 Subject: [PATCH] feat: Create the built-in memory management tools --- src/config/paths.rs | 10 +- src/config/request_context.rs | 3 + src/function/memory.rs | 398 ++++++++++++++++++++++++++++++++++ src/function/mod.rs | 14 ++ 4 files changed, 423 insertions(+), 2 deletions(-) create mode 100644 src/function/memory.rs diff --git a/src/config/paths.rs b/src/config/paths.rs index 98e010d..6acbacc 100644 --- a/src/config/paths.rs +++ b/src/config/paths.rs @@ -1,5 +1,5 @@ use super::role::Role; -use super::{AGENT_GRAPH_FILE_NAME, AGENTS_DIR_NAME, BASH_PROMPT_UTILS_FILE_NAME, CONFIG_FILE_NAME, ENV_FILE_NAME, FUNCTIONS_BIN_DIR_NAME, FUNCTIONS_DIR_NAME, GLOBAL_TOOLS_DIR_NAME, GLOBAL_TOOLS_UTILS_DIR_NAME, MACROS_DIR_NAME, MCP_FILE_NAME, ModelsOverride, RAGS_DIR_NAME, ROLES_DIR_NAME, SKILLS_DIR_NAME, MEMORY_DIR_NAME, MEMORY_INDEX_FILE_NAME}; +use super::{AGENT_GRAPH_FILE_NAME, AGENTS_DIR_NAME, BASH_PROMPT_UTILS_FILE_NAME, CONFIG_FILE_NAME, ENV_FILE_NAME, FUNCTIONS_BIN_DIR_NAME, FUNCTIONS_DIR_NAME, GLOBAL_TOOLS_DIR_NAME, GLOBAL_TOOLS_UTILS_DIR_NAME, MACROS_DIR_NAME, MCP_FILE_NAME, ModelsOverride, RAGS_DIR_NAME, ROLES_DIR_NAME, SKILLS_DIR_NAME, MEMORY_DIR_NAME, MEMORY_INDEX_FILE_NAME, WORKSPACE_MEMORY_DIR_NAME}; use crate::client::ProviderModels; use crate::utils::{get_env_name, list_file_names, normalize_env_name}; @@ -8,7 +8,7 @@ use log::LevelFilter; use std::collections::HashSet; use std::env; use std::fs::{read_dir, read_to_string}; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; pub fn config_dir() -> PathBuf { if let Ok(v) = env::var(get_env_name("config_dir")) { @@ -198,6 +198,12 @@ pub fn global_memory_index_path() -> PathBuf { global_memory_dir().join(MEMORY_INDEX_FILE_NAME) } +pub fn workspace_memory_dir_for(workspace_root: &Path) -> PathBuf { + workspace_root + .join(WORKSPACE_MEMORY_DIR_NAME) + .join(MEMORY_DIR_NAME) +} + pub fn log_config() -> Result<(LevelFilter, Option)> { let log_level = env::var(get_env_name("log_level")) .ok() diff --git a/src/config/request_context.rs b/src/config/request_context.rs index 06130b4..3f8c7f0 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -2482,6 +2482,9 @@ impl RequestContext { if app.function_calling_support && policy.skills_enabled { functions.append_skill_functions(); } + if self.should_register_memory_tools() { + functions.append_memory_functions(); + } let tool_tracker = self.tool_scope.tool_tracker.clone(); self.tool_scope = ToolScope { diff --git a/src/function/memory.rs b/src/function/memory.rs new file mode 100644 index 0000000..211eabd --- /dev/null +++ b/src/function/memory.rs @@ -0,0 +1,398 @@ +use std::collections::HashSet; +use std::env; +use std::path::PathBuf; + +use anyhow::{Context, Result, anyhow, bail}; +use indexmap::IndexMap; +use serde_json::{Value, json}; + +use super::{FunctionDeclaration, JsonSchema}; +use crate::config::RequestContext; +use crate::config::memory::{MemoryFile, MemoryFrontmatter, MemoryStore, WorkspaceMemory}; +use crate::config::paths; + +pub const MEMORY_FUNCTION_PREFIX: &str = "memory__"; + +const PER_FILE_SOFT_CAP: usize = 2_000; + +pub fn memory_function_declarations() -> Vec { + vec![ + FunctionDeclaration { + name: format!("{MEMORY_FUNCTION_PREFIX}read"), + description: "Read the full content of a specific memory file by its name slug." + .to_string(), + parameters: JsonSchema { + type_value: Some("object".to_string()), + properties: Some(IndexMap::from([( + "name".to_string(), + JsonSchema { + type_value: Some("string".to_string()), + description: Some( + "The `name:` slug of the memory file to read (from MEMORY.md index)" + .into(), + ), + ..Default::default() + }, + )])), + required: Some(vec!["name".to_string()]), + ..Default::default() + }, + agent: false, + }, + FunctionDeclaration { + name: format!("{MEMORY_FUNCTION_PREFIX}write"), + description: + "Create or replace a memory file. Caller must also update MEMORY.md index." + .to_string(), + parameters: JsonSchema { + type_value: Some("object".to_string()), + properties: Some(IndexMap::from([ + ( + "name".to_string(), + JsonSchema { + type_value: Some("string".to_string()), + description: Some( + "Short kebab-case slug for the file (no extension)".into(), + ), + ..Default::default() + }, + ), + ( + "description".to_string(), + JsonSchema { + type_value: Some("string".to_string()), + description: Some("One-line description for the MEMORY.md index".into()), + ..Default::default() + }, + ), + ( + "type".to_string(), + JsonSchema { + type_value: Some("string".to_string()), + description: Some( + "Memory type: user | feedback | project | reference".into(), + ), + ..Default::default() + }, + ), + ( + "content".to_string(), + JsonSchema { + type_value: Some("string".to_string()), + description: Some("The full markdown body of the memory file".into()), + ..Default::default() + }, + ), + ( + "scope".to_string(), + JsonSchema { + type_value: Some("string".to_string()), + description: Some( + "Where to write: 'global' (user-level) or 'workspace' (project-level)" + .into(), + ), + ..Default::default() + }, + ), + ])), + required: Some(vec![ + "name".to_string(), + "description".to_string(), + "content".to_string(), + "scope".to_string(), + ]), + ..Default::default() + }, + agent: false, + }, + FunctionDeclaration { + name: format!("{MEMORY_FUNCTION_PREFIX}list"), + description: "List all known drill files with metadata (size, type, scope).".to_string(), + parameters: JsonSchema { + type_value: Some("object".to_string()), + properties: Some(IndexMap::new()), + ..Default::default() + }, + agent: false, + }, + FunctionDeclaration { + name: format!("{MEMORY_FUNCTION_PREFIX}lint"), + description: "Health-check memory: orphan files, broken [[wikilinks]], oversized files." + .to_string(), + parameters: JsonSchema { + type_value: Some("object".to_string()), + properties: Some(IndexMap::new()), + ..Default::default() + }, + agent: false, + }, + ] +} + +pub fn handle_memory_tool(ctx: &mut RequestContext, cmd_name: &str, args: &Value) -> Result { + if !ctx.should_register_memory_tools() { + bail!("Memory tools are disabled (memory off or function calling unavailable)."); + } + + let action = cmd_name + .strip_prefix(MEMORY_FUNCTION_PREFIX) + .unwrap_or(cmd_name); + let cwd = env::current_dir().context("get cwd")?; + let store = MemoryStore::new(&cwd); + + match action { + "read" => { + let name = args + .get("name") + .and_then(Value::as_str) + .ok_or_else(|| anyhow!("name is required"))?; + let file = find_file(&store, name)? + .ok_or_else(|| anyhow!("memory file '{}' not found", name))?; + + Ok(json!({ + "name": file.frontmatter.name, + "type": file.frontmatter.kind, + "content": file.body, + })) + } + "list" => { + let files = store.list_files()?; + let entries: Vec<_> = files + .iter() + .map(|f| { + json!({ + "name": f.frontmatter.name, + "description": f.frontmatter.description, + "type": f.frontmatter.kind, + "char_len": f.char_len(), + "path": f.path.display().to_string(), + }) + }) + .collect(); + + Ok(json!({ + "files": entries, + "global_index_exists": paths::global_memory_index_path().exists(), + "workspace": store.workspace.as_ref().map(workspace_label), + })) + } + "write" => { + let name = arg_str(args, "name")?; + let description = arg_str(args, "description")?; + let content = arg_str(args, "content")?; + let scope = arg_str(args, "scope")?; + let kind = args.get("type").and_then(Value::as_str).map(String::from); + + let target_dir = match scope.as_str() { + "global" => paths::global_memory_dir(), + "workspace" => workspace_write_dir(&store)?, + other => bail!("unknown scope '{}': use 'global' or 'workspace'", other), + }; + let file = MemoryFile { + path: target_dir.join(format!("{name}.md")), + frontmatter: MemoryFrontmatter { + name: name.clone(), + description: Some(description), + kind, + }, + body: content, + }; + file.save()?; + + Ok(json!({ + "status": "ok", + "path": file.path.display().to_string(), + "reminder": "Update MEMORY.md to keep the index accurate.", + })) + } + "lint" => lint_memory(&store), + _ => bail!("unknown memory action: {action}"), + } +} + +fn arg_str(args: &Value, key: &str) -> Result { + args.get(key) + .and_then(Value::as_str) + .map(String::from) + .ok_or_else(|| anyhow!("{} is required", key)) +} + +fn find_file(store: &MemoryStore, name: &str) -> Result> { + Ok(store + .list_files()? + .into_iter() + .find(|f| f.frontmatter.name == name)) +} + +fn workspace_write_dir(store: &MemoryStore) -> Result { + match &store.workspace { + Some(WorkspaceMemory::Structured { dir, .. }) => Ok(dir.clone()), + Some(WorkspaceMemory::Lite { workspace_root, .. }) => { + Ok(paths::workspace_memory_dir_for(workspace_root)) + } + None => bail!("no workspace memory discoverable; cannot write workspace-scoped memory"), + } +} + +fn workspace_label(w: &WorkspaceMemory) -> Value { + match w { + WorkspaceMemory::Structured { workspace_root, .. } => json!({ + "mode": "structured", + "root": workspace_root.display().to_string(), + }), + WorkspaceMemory::Lite { + workspace_root, + file, + } => json!({ + "mode": "lite", + "root": workspace_root.display().to_string(), + "file": file.display().to_string(), + }), + } +} + +fn lint_memory(store: &MemoryStore) -> Result { + let files = store.list_files()?; + let names: HashSet<&str> = + files.iter().map(|f| f.frontmatter.name.as_str()).collect(); + + let mut oversized = Vec::new(); + let mut broken_links = Vec::new(); + for f in &files { + if f.char_len() > PER_FILE_SOFT_CAP { + oversized.push(json!({"name": &f.frontmatter.name, "chars": f.char_len()})); + } + for link in extract_wikilinks(&f.body) { + if !names.contains(link.as_str()) { + broken_links.push(json!({"from": &f.frontmatter.name, "to": link})); + } + } + } + + let index_content = store + .load_global_index()? + .or_else(|| store.load_workspace_index().ok().flatten()) + .unwrap_or_default(); + let mut orphans = Vec::new(); + for f in &files { + if !index_content.contains(&f.frontmatter.name) { + orphans.push(f.frontmatter.name.clone()); + } + } + + Ok(json!({ + "total_files": files.len(), + "oversized": oversized, + "broken_wikilinks": broken_links, + "orphans": orphans, + })) +} + +fn extract_wikilinks(body: &str) -> Vec { + let mut out = Vec::new(); + let bytes = body.as_bytes(); + let mut i = 0; + while i + 1 < bytes.len() { + if bytes[i] == b'[' && bytes[i + 1] == b'[' { + if let Some(end_rel) = body[i + 2..].find("]]") { + out.push(body[i + 2..i + 2 + end_rel].to_string()); + i = i + 2 + end_rel + 2; + continue; + } + } + i += 1; + } + + out +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::memory::discover_workspace_memory; + use std::fs; + use std::time; + + fn temp_root(label: &str) -> PathBuf { + let unique = time::SystemTime::now() + .duration_since(time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let root = env::temp_dir().join(format!("coyote-function-memory-{label}-{unique}")); + fs::create_dir_all(&root).unwrap(); + root + } + + #[test] + fn extract_wikilinks_finds_all_pairs() { + let body = "see [[alpha]] and [[bravo]] but not [single] or [[unclosed"; + + assert_eq!( + extract_wikilinks(body), + vec!["alpha".to_string(), "bravo".to_string()] + ); + } + + #[test] + fn extract_wikilinks_handles_empty_and_no_links() { + assert!(extract_wikilinks("").is_empty()); + assert!(extract_wikilinks("nothing here").is_empty()); + } + + #[test] + fn lint_flags_orphans_broken_links_and_oversized() { + let root = temp_root("lint"); + let workspace = root.join("ws"); + let structured = workspace.join(".coyote").join("memory"); + fs::create_dir_all(&structured).unwrap(); + + fs::write(structured.join("MEMORY.md"), "- referenced\n").unwrap(); + fs::write( + structured.join("referenced.md"), + "---\nname: referenced\n---\nlinks to [[missing]] and [[also_missing]]\n", + ) + .unwrap(); + fs::write( + structured.join("orphan.md"), + "---\nname: orphan\n---\nnot in the index\n", + ) + .unwrap(); + let huge_body = "x".repeat(PER_FILE_SOFT_CAP + 100); + fs::write( + structured.join("huge.md"), + format!("---\nname: huge\n---\n{huge_body}\n"), + ) + .unwrap(); + + let store = MemoryStore { + global_dir: root.join("nonexistent_global"), + workspace: discover_workspace_memory(&workspace), + }; + + let report = lint_memory(&store).unwrap(); + assert_eq!(report["total_files"], 3); + + let orphans = report["orphans"].as_array().unwrap(); + let orphan_names: Vec<&str> = orphans.iter().filter_map(|v| v.as_str()).collect(); + assert!(orphan_names.contains(&"orphan")); + assert!(orphan_names.contains(&"huge")); + assert!(!orphan_names.contains(&"referenced")); + + let broken = report["broken_wikilinks"].as_array().unwrap(); + let broken_targets: Vec<&str> = broken + .iter() + .filter_map(|v| v["to"].as_str()) + .collect(); + assert!(broken_targets.contains(&"missing")); + assert!(broken_targets.contains(&"also_missing")); + + let oversized = report["oversized"].as_array().unwrap(); + let oversized_names: Vec<&str> = oversized + .iter() + .filter_map(|v| v["name"].as_str()) + .collect(); + assert_eq!(oversized_names, vec!["huge"]); + + let _ = fs::remove_dir_all(&root); + } +} diff --git a/src/function/mod.rs b/src/function/mod.rs index ba66837..703bed8 100644 --- a/src/function/mod.rs +++ b/src/function/mod.rs @@ -1,3 +1,4 @@ +pub(crate) mod memory; pub(crate) mod skill; pub(crate) mod supervisor; pub(crate) mod todo; @@ -22,6 +23,7 @@ use indoc::formatdoc; use rust_embed::Embed; use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; +use memory::MEMORY_FUNCTION_PREFIX; use skill::SKILL_FUNCTION_PREFIX; use std::collections::VecDeque; use std::ffi::OsStr; @@ -355,6 +357,11 @@ impl Functions { self.declarations.extend(todo::todo_function_declarations()); } + pub fn append_memory_functions(&mut self) { + self.declarations + .extend(memory::memory_function_declarations()); + } + pub fn append_skill_functions(&mut self) { self.declarations .extend(skill::skill_function_declarations()); @@ -1046,6 +1053,13 @@ impl ToolCall { json!({"tool_call_error": error_msg}) }) } + _ if cmd_name.starts_with(MEMORY_FUNCTION_PREFIX) => { + memory::handle_memory_tool(ctx, &cmd_name, &json_data).unwrap_or_else(|e| { + let error_msg = format!("Memory tool failed: {e}"); + eprintln!("{}", warning_text(&format!("⚠️ {error_msg} ⚠️"))); + json!({"tool_call_error": error_msg}) + }) + } _ if cmd_name.starts_with(SKILL_FUNCTION_PREFIX) => { skill::handle_skill_tool(ctx, &cmd_name, &json_data) .await