feat: Create the built-in memory management tools

2026-06-10 18:35:59 -06:00
parent 4ece3d3df1
commit 19d2340489
4 changed files with 423 additions and 2 deletions
@@ -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<FunctionDeclaration> {
+    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<Value> {
+    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<String> {
+    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<Option<MemoryFile>> {
+    Ok(store
+        .list_files()?
+        .into_iter()
+        .find(|f| f.frontmatter.name == name))
+}
+
+fn workspace_write_dir(store: &MemoryStore) -> Result<PathBuf> {
+    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<Value> {
+    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<String> {
+    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);
+    }
+}
@@ -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