From c04f2855e4dde90e1761f382781fb32d49dd601a Mon Sep 17 00:00:00 2001
From: Edvard <ecschoye@stud.ntnu.no>
Date: Tue, 17 Feb 2026 18:40:01 -0500
Subject: [PATCH] feat(tools): expose custom memory categories in memory_store
 tool

The MemoryCategory::Custom variant already exists in the memory backend
but the memory_store tool only accepted core/daily/conversation. Now any
string is accepted as a category, passing through to Custom(name) for
non-builtin values.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 src/tools/memory_store.rs | 25 +++++++++++++++++++++----
 1 file changed, 21 insertions(+), 4 deletions(-)

diff --git a/src/tools/memory_store.rs b/src/tools/memory_store.rs
index d2aad40..8d28714 100644
--- a/src/tools/memory_store.rs
+++ b/src/tools/memory_store.rs
@@ -22,7 +22,7 @@ impl Tool for MemoryStoreTool {
     }
 
     fn description(&self) -> &str {
-        "Store a fact, preference, or note in long-term memory. Use category 'core' for permanent facts, 'daily' for session notes, 'conversation' for chat context."
+        "Store a fact, preference, or note in long-term memory. Use category 'core' for permanent facts, 'daily' for session notes, 'conversation' for chat context, or a custom category name."
     }
 
     fn parameters_schema(&self) -> serde_json::Value {
@@ -39,8 +39,7 @@ impl Tool for MemoryStoreTool {
                 },
                 "category": {
                     "type": "string",
-                    "enum": ["core", "daily", "conversation"],
-                    "description": "Memory category: core (permanent), daily (session), conversation (chat)"
+                    "description": "Memory category: 'core' (permanent), 'daily' (session), 'conversation' (chat), or a custom category name. Defaults to 'core'."
                 }
             },
             "required": ["key", "content"]
@@ -59,9 +58,10 @@ impl Tool for MemoryStoreTool {
             .ok_or_else(|| anyhow::anyhow!("Missing 'content' parameter"))?;
 
         let category = match args.get("category").and_then(|v| v.as_str()) {
+            Some("core") | None => MemoryCategory::Core,
             Some("daily") => MemoryCategory::Daily,
             Some("conversation") => MemoryCategory::Conversation,
-            _ => MemoryCategory::Core,
+            Some(other) => MemoryCategory::Custom(other.to_string()),
         };
 
         match self.memory.store(key, content, category, None).await {
@@ -128,6 +128,23 @@ mod tests {
         assert!(result.success);
     }
 
+    #[tokio::test]
+    async fn store_with_custom_category() {
+        let (_tmp, mem) = test_mem();
+        let tool = MemoryStoreTool::new(mem.clone());
+        let result = tool
+            .execute(
+                json!({"key": "proj_note", "content": "Uses async runtime", "category": "project"}),
+            )
+            .await
+            .unwrap();
+        assert!(result.success);
+
+        let entry = mem.get("proj_note").await.unwrap().unwrap();
+        assert_eq!(entry.content, "Uses async runtime");
+        assert_eq!(entry.category, MemoryCategory::Custom("project".into()));
+    }
+
     #[tokio::test]
     async fn store_missing_key() {
         let (_tmp, mem) = test_mem();