Skip to main content

vtcode_core/tools/
native_memory.rs

1use anyhow::{Context, Result, anyhow, bail};
2use serde::Deserialize;
3use serde_json::{Value, json};
4use std::path::{Component, Path, PathBuf};
5
6use crate::config::PersistentMemoryConfig;
7use crate::config::loader::VTCodeConfig;
8use crate::persistent_memory::{
9    rebuild_generated_memory_files, resolve_persistent_memory_dir, scaffold_persistent_memory,
10};
11use crate::tools::error_helpers::deserialize_tool_args;
12
13const MEMORIES_ROOT: &str = "/memories";
14
15#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
16#[serde(rename_all = "snake_case")]
17pub enum NativeMemoryCommand {
18    View,
19    Create,
20    StrReplace,
21    Insert,
22    Delete,
23    Rename,
24}
25
26#[derive(Debug, Clone, Deserialize)]
27pub struct NativeMemoryRequest {
28    pub command: NativeMemoryCommand,
29    #[serde(default)]
30    pub path: Option<String>,
31    #[serde(default)]
32    pub old_path: Option<String>,
33    #[serde(default)]
34    pub new_path: Option<String>,
35    #[serde(default)]
36    pub file_text: Option<String>,
37    #[serde(default)]
38    pub old_str: Option<String>,
39    #[serde(default)]
40    pub new_str: Option<String>,
41    #[serde(default)]
42    pub insert_line: Option<usize>,
43    #[serde(default)]
44    pub insert_text: Option<String>,
45}
46
47pub fn parameter_schema() -> Value {
48    json!({
49        "type": "object",
50        "properties": {
51            "command": {
52                "type": "string",
53                "enum": ["view", "create", "str_replace", "insert", "delete", "rename"],
54                "description": "Memory operation to perform."
55            },
56            "path": {
57                "type": "string",
58                "description": "Path under /memories for view/create/str_replace/insert/delete."
59            },
60            "old_path": {
61                "type": "string",
62                "description": "Existing path under /memories for rename."
63            },
64            "new_path": {
65                "type": "string",
66                "description": "Destination path under /memories for rename."
67            },
68            "file_text": {
69                "type": "string",
70                "description": "Full file contents for create."
71            },
72            "old_str": {
73                "type": "string",
74                "description": "Exact substring to replace for str_replace."
75            },
76            "new_str": {
77                "type": "string",
78                "description": "Replacement string for str_replace."
79            },
80            "insert_line": {
81                "type": "integer",
82                "minimum": 0,
83                "description": "Zero-based line index for insert."
84            },
85            "insert_text": {
86                "type": "string",
87                "description": "Text to insert at insert_line."
88            }
89        },
90        "required": ["command"],
91        "additionalProperties": false
92    })
93}
94
95pub async fn execute(
96    workspace_root: &Path,
97    config: &PersistentMemoryConfig,
98    args: Value,
99) -> Result<Value> {
100    let request: NativeMemoryRequest = deserialize_tool_args(&args, "memory")?;
101    let root = prepare_root(workspace_root, config).await?;
102    let output = execute_request(&root, workspace_root, config, request).await?;
103    Ok(Value::String(output))
104}
105
106pub async fn execute_with_vt_config(
107    workspace_root: &Path,
108    vt_cfg: &VTCodeConfig,
109    args: Value,
110) -> Result<Value> {
111    if !vt_cfg.persistent_memory_enabled() {
112        bail!(
113            "Persistent memory is disabled. Enable features.memories and agent.persistent_memory.enabled to use /memories"
114        );
115    }
116
117    execute(workspace_root, &vt_cfg.agent.persistent_memory, args).await
118}
119
120async fn prepare_root(workspace_root: &Path, config: &PersistentMemoryConfig) -> Result<PathBuf> {
121    scaffold_persistent_memory(config, workspace_root)
122        .await
123        .context("Failed to scaffold persistent memory layout")?;
124    let cfg = config.clone();
125    let ws = workspace_root.to_path_buf();
126    tokio::task::spawn_blocking(move || resolve_persistent_memory_dir(&cfg, &ws))
127        .await
128        .context("Persistent memory directory resolution task panicked")??
129        .ok_or_else(|| anyhow!("Persistent memory directory could not be resolved"))
130}
131
132async fn execute_request(
133    root: &Path,
134    workspace_root: &Path,
135    config: &PersistentMemoryConfig,
136    request: NativeMemoryRequest,
137) -> Result<String> {
138    match request.command {
139        NativeMemoryCommand::View => {
140            let path = request.path.as_deref().unwrap_or(MEMORIES_ROOT);
141            view(root, path).await
142        }
143        NativeMemoryCommand::Create => {
144            let path = required_text(request.path.as_deref(), "path")?;
145            let file_text = required_text(request.file_text.as_deref(), "file_text")?;
146            let resolved = resolve_virtual_path(root, path)?;
147            ensure_writable_path(&resolved.relative, path)?;
148            if let Some(parent) = resolved.absolute.parent() {
149                tokio::fs::create_dir_all(parent)
150                    .await
151                    .with_context(|| format!("Failed to create {}", parent.display()))?;
152            }
153            tokio::fs::write(&resolved.absolute, file_text)
154                .await
155                .with_context(|| format!("Failed to write {}", path))?;
156            rebuild_generated_memory_files(config, workspace_root).await?;
157            Ok(format!("Created {path}"))
158        }
159        NativeMemoryCommand::StrReplace => {
160            let path = required_text(request.path.as_deref(), "path")?;
161            let old_str = required_text(request.old_str.as_deref(), "old_str")?;
162            let new_str = request.new_str.as_deref().unwrap_or_default();
163            if old_str.is_empty() {
164                bail!("old_str must not be empty");
165            }
166            let resolved = resolve_virtual_path(root, path)?;
167            ensure_writable_path(&resolved.relative, path)?;
168            let content = tokio::fs::read_to_string(&resolved.absolute)
169                .await
170                .with_context(|| format!("Failed to read {}", path))?;
171            let matches = content.matches(old_str).count();
172            if matches == 0 {
173                bail!("old_str not found in {path}");
174            }
175            if matches > 1 {
176                bail!("old_str appears {matches} times in {path}; be more specific");
177            }
178            tokio::fs::write(&resolved.absolute, content.replacen(old_str, new_str, 1))
179                .await
180                .with_context(|| format!("Failed to write {}", path))?;
181            rebuild_generated_memory_files(config, workspace_root).await?;
182            Ok(format!("Replaced in {path}"))
183        }
184        NativeMemoryCommand::Insert => {
185            let path = required_text(request.path.as_deref(), "path")?;
186            let insert_line = request
187                .insert_line
188                .ok_or_else(|| anyhow!("insert_line is required"))?;
189            let insert_text = required_text(request.insert_text.as_deref(), "insert_text")?;
190            let resolved = resolve_virtual_path(root, path)?;
191            ensure_writable_path(&resolved.relative, path)?;
192            let content = tokio::fs::read_to_string(&resolved.absolute)
193                .await
194                .with_context(|| format!("Failed to read {}", path))?;
195            let mut lines = content
196                .split('\n')
197                .map(ToOwned::to_owned)
198                .collect::<Vec<_>>();
199            if insert_line > lines.len() {
200                bail!("insert_line {} is out of bounds for {}", insert_line, path);
201            }
202            lines.insert(insert_line, insert_text.to_string());
203            tokio::fs::write(&resolved.absolute, lines.join("\n"))
204                .await
205                .with_context(|| format!("Failed to write {}", path))?;
206            rebuild_generated_memory_files(config, workspace_root).await?;
207            Ok(format!("Inserted at line {insert_line} in {path}"))
208        }
209        NativeMemoryCommand::Delete => {
210            let path = required_text(request.path.as_deref(), "path")?;
211            let resolved = resolve_virtual_path(root, path)?;
212            ensure_writable_path(&resolved.relative, path)?;
213            let metadata = tokio::fs::metadata(&resolved.absolute)
214                .await
215                .with_context(|| format!("Failed to stat {}", path))?;
216            if metadata.is_dir() {
217                tokio::fs::remove_dir_all(&resolved.absolute)
218                    .await
219                    .with_context(|| format!("Failed to delete {}", path))?;
220            } else {
221                tokio::fs::remove_file(&resolved.absolute)
222                    .await
223                    .with_context(|| format!("Failed to delete {}", path))?;
224            }
225            rebuild_generated_memory_files(config, workspace_root).await?;
226            Ok(format!("Deleted {path}"))
227        }
228        NativeMemoryCommand::Rename => {
229            let old_path = required_text(request.old_path.as_deref(), "old_path")?;
230            let new_path = required_text(request.new_path.as_deref(), "new_path")?;
231            let old_resolved = resolve_virtual_path(root, old_path)?;
232            let new_resolved = resolve_virtual_path(root, new_path)?;
233            ensure_writable_path(&old_resolved.relative, old_path)?;
234            ensure_writable_path(&new_resolved.relative, new_path)?;
235            if let Some(parent) = new_resolved.absolute.parent() {
236                tokio::fs::create_dir_all(parent)
237                    .await
238                    .with_context(|| format!("Failed to create {}", parent.display()))?;
239            }
240            tokio::fs::rename(&old_resolved.absolute, &new_resolved.absolute)
241                .await
242                .with_context(|| format!("Failed to rename {old_path} to {new_path}"))?;
243            rebuild_generated_memory_files(config, workspace_root).await?;
244            Ok(format!("Renamed {old_path} -> {new_path}"))
245        }
246    }
247}
248
249async fn view(root: &Path, path: &str) -> Result<String> {
250    let resolved = resolve_virtual_path(root, path)?;
251    if !resolved.absolute.exists() {
252        if path == MEMORIES_ROOT {
253            return Ok("Directory /memories is empty.".to_string());
254        }
255        bail!("{path} does not exist");
256    }
257
258    let metadata = tokio::fs::metadata(&resolved.absolute)
259        .await
260        .with_context(|| format!("Failed to stat {}", path))?;
261    if metadata.is_dir() {
262        let mut entries = tokio::fs::read_dir(&resolved.absolute)
263            .await
264            .with_context(|| format!("Failed to list {}", path))?;
265        let mut names = Vec::new();
266        while let Some(entry) = entries.next_entry().await? {
267            let entry_path = entry.path();
268            let file_name = entry.file_name().to_string_lossy().to_string();
269            let suffix = if entry_path.is_dir() { "/" } else { "" };
270            names.push(format!("{file_name}{suffix}"));
271        }
272        names.sort();
273        if names.is_empty() {
274            return Ok("(empty directory)".to_string());
275        }
276        return Ok(names.join("\n"));
277    }
278
279    let content = tokio::fs::read_to_string(&resolved.absolute)
280        .await
281        .with_context(|| format!("Failed to read {}", path))?;
282    Ok(content
283        .split('\n')
284        .enumerate()
285        .map(|(idx, line)| format!("{:4}\t{}", idx + 1, line))
286        .collect::<Vec<_>>()
287        .join("\n"))
288}
289
290fn required_text<'a>(value: Option<&'a str>, field: &str) -> Result<&'a str> {
291    value
292        .map(str::trim)
293        .filter(|value| !value.is_empty())
294        .ok_or_else(|| anyhow!("{field} is required"))
295}
296
297struct ResolvedMemoryPath {
298    absolute: PathBuf,
299    relative: PathBuf,
300}
301
302fn resolve_virtual_path(root: &Path, virtual_path: &str) -> Result<ResolvedMemoryPath> {
303    let trimmed = virtual_path.trim();
304    if !trimmed.starts_with(MEMORIES_ROOT) {
305        bail!("memory paths must stay under {MEMORIES_ROOT}");
306    }
307
308    let relative_raw = trimmed
309        .strip_prefix(MEMORIES_ROOT)
310        .unwrap_or_default()
311        .trim_start_matches('/');
312    let relative = if relative_raw.is_empty() {
313        PathBuf::new()
314    } else {
315        sanitize_relative_path(relative_raw)?
316    };
317
318    Ok(ResolvedMemoryPath {
319        absolute: if relative.as_os_str().is_empty() {
320            root.to_path_buf()
321        } else {
322            root.join(&relative)
323        },
324        relative,
325    })
326}
327
328fn sanitize_relative_path(raw: &str) -> Result<PathBuf> {
329    let path = Path::new(raw);
330    let mut sanitized = PathBuf::new();
331    for component in path.components() {
332        match component {
333            Component::Normal(part) => sanitized.push(part),
334            Component::CurDir => {}
335            Component::ParentDir | Component::RootDir | Component::Prefix(_) => {
336                bail!("memory paths may not escape {MEMORIES_ROOT}");
337            }
338        }
339    }
340    Ok(sanitized)
341}
342
343fn ensure_writable_path(relative: &Path, original: &str) -> Result<()> {
344    if is_writable_relative_path(relative) {
345        return Ok(());
346    }
347
348    bail!(
349        "{original} is read-only; writable paths are /memories/preferences.md, /memories/repository-facts.md, and /memories/notes/**"
350    );
351}
352
353fn is_writable_relative_path(relative: &Path) -> bool {
354    let components = relative
355        .components()
356        .filter_map(|component| match component {
357            Component::Normal(part) => Some(part.to_string_lossy().to_string()),
358            _ => None,
359        })
360        .collect::<Vec<_>>();
361
362    matches!(
363        components.as_slice(),
364        [single] if single == "preferences.md" || single == "repository-facts.md"
365    ) || matches!(components.as_slice(), [first, ..] if first == "notes" && components.len() >= 2)
366}
367
368#[cfg(test)]
369mod tests {
370    use super::{MEMORIES_ROOT, execute, parameter_schema};
371    use crate::config::PersistentMemoryConfig;
372    use crate::persistent_memory::{
373        MEMORY_FILENAME, MEMORY_SUMMARY_FILENAME, ROLLOUT_SUMMARIES_DIRNAME,
374        resolve_persistent_memory_dir,
375    };
376    use serde_json::json;
377    use tempfile::tempdir;
378
379    #[tokio::test]
380    async fn parameter_schema_lists_supported_commands() {
381        let schema = parameter_schema();
382        assert_eq!(
383            schema["properties"]["command"]["enum"],
384            json!([
385                "view",
386                "create",
387                "str_replace",
388                "insert",
389                "delete",
390                "rename"
391            ])
392        );
393    }
394
395    #[tokio::test]
396    async fn execute_supports_crud_and_rebuilds_generated_files() {
397        let workspace = tempdir().expect("workspace");
398        let config = PersistentMemoryConfig {
399            enabled: true,
400            directory_override: Some(workspace.path().join(".memory").display().to_string()),
401            ..PersistentMemoryConfig::default()
402        };
403
404        execute(
405            workspace.path(),
406            &config,
407            json!({
408                "command": "create",
409                "path": "/memories/notes/research.md",
410                "file_text": "# Notes\n\n- First finding"
411            }),
412        )
413        .await
414        .expect("create");
415
416        let view = execute(
417            workspace.path(),
418            &config,
419            json!({
420                "command": "view",
421                "path": "/memories/notes/research.md"
422            }),
423        )
424        .await
425        .expect("view");
426        assert!(view.as_str().expect("string").contains("First finding"));
427
428        execute(
429            workspace.path(),
430            &config,
431            json!({
432                "command": "str_replace",
433                "path": "/memories/notes/research.md",
434                "old_str": "First finding",
435                "new_str": "Updated finding"
436            }),
437        )
438        .await
439        .expect("replace");
440        execute(
441            workspace.path(),
442            &config,
443            json!({
444                "command": "insert",
445                "path": "/memories/notes/research.md",
446                "insert_line": 2,
447                "insert_text": "- Follow-up"
448            }),
449        )
450        .await
451        .expect("insert");
452        execute(
453            workspace.path(),
454            &config,
455            json!({
456                "command": "rename",
457                "old_path": "/memories/notes/research.md",
458                "new_path": "/memories/notes/archive/research.md"
459            }),
460        )
461        .await
462        .expect("rename");
463
464        let memory_dir = resolve_persistent_memory_dir(&config, workspace.path())
465            .expect("dir")
466            .expect("resolved");
467        let summary =
468            std::fs::read_to_string(memory_dir.join(MEMORY_SUMMARY_FILENAME)).expect("summary");
469        assert!(summary.contains("Updated finding") || summary.contains("Follow-up"));
470
471        execute(
472            workspace.path(),
473            &config,
474            json!({
475                "command": "delete",
476                "path": "/memories/notes/archive/research.md"
477            }),
478        )
479        .await
480        .expect("delete");
481        let recreated_summary =
482            std::fs::read_to_string(memory_dir.join(MEMORY_SUMMARY_FILENAME)).expect("summary");
483        assert!(!recreated_summary.contains("Updated finding"));
484        assert!(!recreated_summary.contains("Follow-up"));
485    }
486
487    #[tokio::test]
488    async fn execute_blocks_path_traversal_and_generated_file_writes() {
489        let workspace = tempdir().expect("workspace");
490        let config = PersistentMemoryConfig {
491            enabled: true,
492            directory_override: Some(workspace.path().join(".memory").display().to_string()),
493            ..PersistentMemoryConfig::default()
494        };
495
496        let traversal = execute(
497            workspace.path(),
498            &config,
499            json!({
500                "command": "create",
501                "path": "/memories/notes/../../escape.md",
502                "file_text": "bad"
503            }),
504        )
505        .await;
506        traversal.unwrap_err();
507
508        let readonly = execute(
509            workspace.path(),
510            &config,
511            json!({
512                "command": "create",
513                "path": format!("{MEMORIES_ROOT}/{MEMORY_FILENAME}"),
514                "file_text": "bad"
515            }),
516        )
517        .await;
518        readonly.unwrap_err();
519    }
520
521    #[tokio::test]
522    async fn execute_views_root_and_rollout_directories_read_only() {
523        let workspace = tempdir().expect("workspace");
524        let config = PersistentMemoryConfig {
525            enabled: true,
526            directory_override: Some(workspace.path().join(".memory").display().to_string()),
527            ..PersistentMemoryConfig::default()
528        };
529
530        let root_listing = execute(
531            workspace.path(),
532            &config,
533            json!({
534                "command": "view",
535                "path": MEMORIES_ROOT
536            }),
537        )
538        .await
539        .expect("view root");
540        let root_listing = root_listing.as_str().expect("string");
541        assert!(root_listing.contains("notes/"));
542        assert!(root_listing.contains(&format!("{ROLLOUT_SUMMARIES_DIRNAME}/")));
543
544        let rollout_write = execute(
545            workspace.path(),
546            &config,
547            json!({
548                "command": "create",
549                "path": format!("{MEMORIES_ROOT}/{ROLLOUT_SUMMARIES_DIRNAME}/bad.md"),
550                "file_text": "bad"
551            }),
552        )
553        .await;
554        rollout_write.unwrap_err();
555    }
556}