ai_agent/tools/agent/
agent_memory.rs1#![allow(dead_code)]
3
4use std::path::{Path, PathBuf};
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub enum AgentMemoryScope {
9 User,
10 Project,
11 Local,
12}
13
14impl AgentMemoryScope {
15 pub fn from_str(s: &str) -> Option<Self> {
16 match s {
17 "user" => Some(AgentMemoryScope::User),
18 "project" => Some(AgentMemoryScope::Project),
19 "local" => Some(AgentMemoryScope::Local),
20 _ => None,
21 }
22 }
23
24 pub fn as_str(&self) -> &'static str {
25 match self {
26 AgentMemoryScope::User => "user",
27 AgentMemoryScope::Project => "project",
28 AgentMemoryScope::Local => "local",
29 }
30 }
31}
32
33fn sanitize_agent_type_for_path(agent_type: &str) -> String {
36 agent_type.replace(':', "-")
37}
38
39fn get_memory_base_dir() -> PathBuf {
42 std::env::var("CLAUDE_CODE_MEMORY_BASE_DIR")
43 .map(PathBuf::from)
44 .unwrap_or_else(|_| {
45 dirs::home_dir()
46 .unwrap_or_else(|| PathBuf::from("."))
47 .join(".claude")
48 })
49}
50
51fn get_cwd() -> PathBuf {
53 std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."))
54}
55
56fn get_local_agent_memory_dir(dir_name: &str) -> PathBuf {
58 if let Ok(remote_dir) = std::env::var("CLAUDE_CODE_REMOTE_MEMORY_DIR") {
59 let project_root = get_project_root();
60 PathBuf::from(&remote_dir)
61 .join("projects")
62 .join(sanitize_path(&project_root))
63 .join("agent-memory-local")
64 .join(dir_name)
65 } else {
66 get_cwd()
67 .join(".claude")
68 .join("agent-memory-local")
69 .join(dir_name)
70 }
71}
72
73fn sanitize_path(path: &str) -> String {
75 path.replace(
76 |c: char| !c.is_alphanumeric() && c != '/' && c != '-' && c != '_',
77 "_",
78 )
79}
80
81fn get_project_root() -> String {
83 get_cwd().to_string_lossy().to_string()
85}
86
87pub fn get_agent_memory_dir(agent_type: &str, scope: AgentMemoryScope) -> PathBuf {
89 let dir_name = sanitize_agent_type_for_path(agent_type);
90 match scope {
91 AgentMemoryScope::Project => get_cwd()
92 .join(".claude")
93 .join("agent-memory")
94 .join(dir_name),
95 AgentMemoryScope::Local => get_local_agent_memory_dir(&dir_name),
96 AgentMemoryScope::User => get_memory_base_dir().join("agent-memory").join(dir_name),
97 }
98}
99
100pub fn is_agent_memory_path(absolute_path: &str) -> bool {
102 let normalized = Path::new(absolute_path)
103 .canonicalize()
104 .unwrap_or_else(|_| absolute_path.into());
105 let normalized_str = normalized.to_string_lossy();
106 let memory_base = get_memory_base_dir();
107
108 if normalized_str.starts_with(
110 &memory_base
111 .join("agent-memory")
112 .to_string_lossy()
113 .to_string(),
114 ) {
115 return true;
116 }
117
118 let project_mem = get_cwd().join(".claude").join("agent-memory");
120 if normalized_str.starts_with(&project_mem.to_string_lossy().to_string()) {
121 return true;
122 }
123
124 if let Ok(remote_dir) = std::env::var("CLAUDE_CODE_REMOTE_MEMORY_DIR") {
126 if normalized_str.contains("agent-memory-local")
127 && normalized_str.starts_with(&format!("{}/projects", remote_dir))
128 {
129 return true;
130 }
131 } else {
132 let local_mem = get_cwd().join(".claude").join("agent-memory-local");
133 if normalized_str.starts_with(&local_mem.to_string_lossy().to_string()) {
134 return true;
135 }
136 }
137
138 false
139}
140
141pub fn get_agent_memory_entrypoint(agent_type: &str, scope: AgentMemoryScope) -> PathBuf {
143 get_agent_memory_dir(agent_type, scope).join("MEMORY.md")
144}
145
146pub fn get_memory_scope_display(scope: Option<AgentMemoryScope>) -> &'static str {
148 match scope {
149 Some(AgentMemoryScope::User) => "User (~/.claude/agent-memory/)",
150 Some(AgentMemoryScope::Project) => "Project (.claude/agent-memory/)",
151 Some(AgentMemoryScope::Local) => "Local (.claude/agent-memory-local/)",
152 None => "None",
153 }
154}
155
156pub fn load_agent_memory_prompt(agent_type: &str, scope: AgentMemoryScope) -> String {
159 let scope_note = match scope {
160 AgentMemoryScope::User => {
161 "- Since this memory is user-scope, keep learnings general since they apply across all projects"
162 }
163 AgentMemoryScope::Project => {
164 "- Since this memory is project-scope and shared with your team via version control, tailor your memories to this project"
165 }
166 AgentMemoryScope::Local => {
167 "- Since this memory is local-scope (not checked into version control), tailor your memories to this project and machine"
168 }
169 };
170
171 let memory_dir = get_agent_memory_dir(agent_type, scope);
172
173 let _ = std::fs::create_dir_all(&memory_dir);
175
176 let extra_guidelines = std::env::var("CLAUDE_COWORK_MEMORY_EXTRA_GUIDELINES").ok();
177 let extra_guidelines = extra_guidelines.as_deref().filter(|s| !s.trim().is_empty());
178
179 build_memory_prompt(
180 "Persistent Agent Memory",
181 &memory_dir,
182 if let Some(guidelines) = extra_guidelines {
183 vec![scope_note, guidelines]
184 } else {
185 vec![scope_note]
186 },
187 )
188}
189
190fn build_memory_prompt(
192 display_name: &str,
193 memory_dir: &Path,
194 extra_guidelines: Vec<&str>,
195) -> String {
196 let memory_contents = read_memory_files(memory_dir);
197 let guidelines = extra_guidelines.join("\n");
198
199 format!(
200 "# {display_name}\n\n\
201 Memory directory: {memory_dir}\n\n\
202 {guidelines}\n\n\
203 {memory_contents}",
204 memory_dir = memory_dir.display()
205 )
206}
207
208fn read_memory_files(memory_dir: &Path) -> String {
210 let mut contents = String::new();
211
212 if let Ok(entries) = std::fs::read_dir(memory_dir) {
213 for entry in entries.flatten() {
214 let path = entry.path();
215 if path.extension().and_then(|e| e.to_str()) == Some("md") {
216 if let Ok(content) = std::fs::read_to_string(&path) {
217 contents.push_str(&format!("\n--- {} ---\n{}\n", path.display(), content));
218 }
219 }
220 }
221 }
222
223 contents
224}
225
226#[cfg(test)]
227mod tests {
228 use super::*;
229
230 #[test]
231 fn test_sanitize_agent_type() {
232 assert_eq!(sanitize_agent_type_for_path("my-agent"), "my-agent");
233 assert_eq!(
234 sanitize_agent_type_for_path("my-plugin:my-agent"),
235 "my-plugin-my-agent"
236 );
237 }
238
239 #[test]
240 fn test_memory_scope_from_str() {
241 assert_eq!(
242 AgentMemoryScope::from_str("user"),
243 Some(AgentMemoryScope::User)
244 );
245 assert_eq!(
246 AgentMemoryScope::from_str("project"),
247 Some(AgentMemoryScope::Project)
248 );
249 assert_eq!(
250 AgentMemoryScope::from_str("local"),
251 Some(AgentMemoryScope::Local)
252 );
253 assert_eq!(AgentMemoryScope::from_str("invalid"), None);
254 }
255
256 #[test]
257 fn test_memory_scope_display() {
258 assert_eq!(get_memory_scope_display(None), "None");
259 assert_eq!(
260 get_memory_scope_display(Some(AgentMemoryScope::User)),
261 "User (~/.claude/agent-memory/)"
262 );
263 }
264
265 #[test]
266 fn test_get_agent_memory_entrypoint() {
267 let path = get_agent_memory_entrypoint("test-agent", AgentMemoryScope::Project);
268 assert!(path.to_string_lossy().contains("agent-memory"));
269 assert!(path.to_string_lossy().contains("MEMORY.md"));
270 }
271}