From 218750cc1ee2ebedae44e9366fbcf2dd8e1bc479 Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Wed, 10 Jun 2026 19:24:08 -0600 Subject: [PATCH] feat: added .set memory REPL commands to control memory injection and applied formatting --- src/config/memory.rs | 44 +++++++++++++++++++------------- src/config/paths.rs | 8 +++++- src/config/request_context.rs | 47 ++++++++++++++++++++++++++++------- src/function/memory.rs | 21 +++++++--------- src/function/mod.rs | 7 +++++- 5 files changed, 87 insertions(+), 40 deletions(-) diff --git a/src/config/memory.rs b/src/config/memory.rs index abed514..3620ef7 100644 --- a/src/config/memory.rs +++ b/src/config/memory.rs @@ -5,22 +5,29 @@ use anyhow::{Context, Result}; use log::warn; use serde::{Deserialize, Serialize}; -use crate::config::{paths, MEMORY_DIR_NAME, MEMORY_INDEX_FILE_NAME, WORKSPACE_MEMORY_DIR_NAME, WORKSPACE_MEMORY_FILE_NAME}; +use crate::config::{ + MEMORY_DIR_NAME, MEMORY_INDEX_FILE_NAME, WORKSPACE_MEMORY_DIR_NAME, WORKSPACE_MEMORY_FILE_NAME, + paths, +}; pub const DEFAULT_MEMORY_CAP_WITH_TOOLS: usize = 6_000; pub const DEFAULT_MEMORY_CAP_WITHOUT_TOOLS: usize = 12_000; #[derive(Debug, Clone)] pub enum WorkspaceMemory { - Structured { workspace_root: PathBuf, dir: PathBuf }, - Lite { workspace_root: PathBuf, file: PathBuf }, + Structured { + workspace_root: PathBuf, + dir: PathBuf, + }, + Lite { + workspace_root: PathBuf, + file: PathBuf, + }, } pub fn discover_workspace_memory(start: &Path) -> Option { for dir in start.ancestors() { - let structured = dir - .join(WORKSPACE_MEMORY_DIR_NAME) - .join(MEMORY_DIR_NAME); + let structured = dir.join(WORKSPACE_MEMORY_DIR_NAME).join(MEMORY_DIR_NAME); if structured.join(MEMORY_INDEX_FILE_NAME).exists() { return Some(WorkspaceMemory::Structured { workspace_root: dir.to_path_buf(), @@ -181,7 +188,7 @@ pub fn build_memory_section( buf.push_str("\n\n"); consumed += s.chars().count(); } - + if let Some(s) = &workspace_index { buf.push_str("\n"); buf.push_str(s); @@ -193,8 +200,7 @@ pub fn build_memory_section( warn!( "memory indexes ({} chars) exceed cap ({} chars); injecting fully - \ consider raising memory_cap_* in config or shrinking MEMORY.md", - consumed, - cap + consumed, cap ); } @@ -214,7 +220,7 @@ pub fn build_memory_section( buf.push_str("\n\n"); budget = budget.saturating_sub(needed); } - + if omitted > 0 { buf.push_str(&format!( "\n", @@ -340,7 +346,11 @@ mod tests { .join(WORKSPACE_MEMORY_DIR_NAME) .join(MEMORY_DIR_NAME); fs::create_dir_all(&structured).unwrap(); - fs::write(structured.join(MEMORY_INDEX_FILE_NAME), "workspace-index-content").unwrap(); + fs::write( + structured.join(MEMORY_INDEX_FILE_NAME), + "workspace-index-content", + ) + .unwrap(); fs::write( structured.join("foo.md"), "---\nname: foo\n---\nfoo body that should not appear\n", @@ -431,9 +441,9 @@ mod tests { #[test] fn parse_frontmatter_extracts_yaml() { let raw = "---\nname: foo\ndescription: a thing\ntype: user\n---\nBody text\n"; - + let (fm, body) = parse_frontmatter(raw).unwrap(); - + assert_eq!(fm.name, "foo"); assert_eq!(fm.description.as_deref(), Some("a thing")); assert_eq!(fm.kind.as_deref(), Some("user")); @@ -443,9 +453,9 @@ mod tests { #[test] fn parse_frontmatter_handles_missing_block() { let raw = "# Just markdown, no frontmatter\nbody"; - + let (fm, body) = parse_frontmatter(raw).unwrap(); - + assert_eq!(fm.name, ""); assert!(fm.kind.is_none()); assert_eq!(body, raw); @@ -454,9 +464,9 @@ mod tests { #[test] fn parse_frontmatter_handles_unterminated_block() { let raw = "---\nname: oops\nno closing delimiter\n# rest of doc"; - + let (fm, body) = parse_frontmatter(raw).unwrap(); - + assert_eq!(fm.name, ""); assert_eq!(body, raw); } diff --git a/src/config/paths.rs b/src/config/paths.rs index 6acbacc..a5d0c40 100644 --- a/src/config/paths.rs +++ b/src/config/paths.rs @@ -1,5 +1,11 @@ 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, WORKSPACE_MEMORY_DIR_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, MEMORY_DIR_NAME, + MEMORY_INDEX_FILE_NAME, ModelsOverride, RAGS_DIR_NAME, ROLES_DIR_NAME, SKILLS_DIR_NAME, + WORKSPACE_MEMORY_DIR_NAME, +}; use crate::client::ProviderModels; use crate::utils::{get_env_name, list_file_names, normalize_env_name}; diff --git a/src/config/request_context.rs b/src/config/request_context.rs index 5b79103..864f2d7 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -5,7 +5,13 @@ use super::skill_policy::SkillPolicy; use super::skill_registry::SkillRegistry; use super::todo::TodoList; use super::tool_scope::{McpRuntime, ToolScope}; -use super::{AGENTS_DIR_NAME, Agent, AgentVariables, AppConfig, AppState, AssetCategory, CREATE_TITLE_ROLE, Input, InstallFilter, LEFT_PROMPT, LastMessage, MESSAGES_FILE_NAME, RIGHT_PROMPT, Role, RoleLike, SESSIONS_DIR_NAME, SUMMARIZATION_PROMPT, SUMMARY_CONTEXT_PROMPT, StateFlags, TEMP_ROLE_NAME, TEMP_SESSION_NAME, WorkingMode, ensure_parent_exists, list_agents, paths, memory}; +use super::{ + AGENTS_DIR_NAME, Agent, AgentVariables, AppConfig, AppState, AssetCategory, CREATE_TITLE_ROLE, + Input, InstallFilter, LEFT_PROMPT, LastMessage, MESSAGES_FILE_NAME, RIGHT_PROMPT, Role, + RoleLike, SESSIONS_DIR_NAME, SUMMARIZATION_PROMPT, SUMMARY_CONTEXT_PROMPT, StateFlags, + TEMP_ROLE_NAME, TEMP_SESSION_NAME, WorkingMode, ensure_parent_exists, list_agents, memory, + paths, +}; use super::{MessageContentToolCalls, prompts}; use crate::client::{Model, ModelType, list_models}; use crate::function::{ @@ -25,6 +31,9 @@ use crate::utils::{ list_file_names, now, render_prompt, temp_file, }; +use super::memory::{ + DEFAULT_MEMORY_CAP_WITH_TOOLS, DEFAULT_MEMORY_CAP_WITHOUT_TOOLS, MemoryStore, WorkspaceMemory, +}; use crate::graph; use anyhow::{Context, Error, Result, bail}; use gman::providers::SupportedProvider; @@ -41,7 +50,6 @@ use std::io::Write; use std::path::{Path, PathBuf}; use std::sync::Arc; use std::{env, fs}; -use super::memory::{WorkspaceMemory, MemoryStore, DEFAULT_MEMORY_CAP_WITH_TOOLS, DEFAULT_MEMORY_CAP_WITHOUT_TOOLS}; pub struct AutoContinueConfig { pub enabled: bool, @@ -677,11 +685,13 @@ impl RequestContext { } } - if self.should_inject_memory() - && let Some(cwd) = env::current_dir().ok() - { - let store = MemoryStore::new(&cwd); - let with_tools = self.should_register_memory_tools(); + let memory_config = self.memory_config(); + if memory_config.enabled { + let store = MemoryStore { + global_dir: paths::global_memory_dir(), + workspace: memory_config.workspace, + }; + let with_tools = app.function_calling_support; let cap = if with_tools { app.memory_cap_with_tools .unwrap_or(DEFAULT_MEMORY_CAP_WITH_TOOLS) @@ -2078,6 +2088,24 @@ impl RequestContext { self.update_app_config(|app| app.skill_instructions = value); } } + "memory" => { + let value: bool = value.parse().with_context(|| "Invalid value")?; + + if let Some(session) = self.session.as_mut() { + session.set_memory(Some(value)); + } else { + self.update_app_config(|app| app.memory = Some(value)); + } + + let should_register = self.should_register_memory_tools(); + let already_registered = self.tool_scope.functions.contains("memory__read"); + + if should_register && !already_registered { + self.tool_scope.functions.append_memory_functions(); + } else if !should_register && already_registered { + self.tool_scope.functions.remove_memory_functions(); + } + } _ => bail!("Unknown key '{key}'"), } Ok(()) @@ -2350,6 +2378,7 @@ impl RequestContext { super::complete_bool(config.inject) } "skill_instructions" => vec!["null".to_string()], + "memory" => super::complete_bool(self.should_inject_memory()), _ => vec![], }; values = candidates.into_iter().map(|v| (v, None)).collect(); @@ -3280,12 +3309,12 @@ mod tests { #[test] fn should_register_memory_tools_false_when_function_calling_off() { let mut ctx = create_test_ctx(); - + ctx.update_app_config(|app| { app.memory = Some(true); app.function_calling_support = false; }); - + assert!( !ctx.should_register_memory_tools(), "memory tools must require function_calling_support even when memory itself would otherwise be enabled" diff --git a/src/function/memory.rs b/src/function/memory.rs index b3e12ea..2d16982 100644 --- a/src/function/memory.rs +++ b/src/function/memory.rs @@ -253,8 +253,7 @@ fn workspace_label(w: &WorkspaceMemory) -> Value { 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 names: HashSet<&str> = files.iter().map(|f| f.frontmatter.name.as_str()).collect(); let mut oversized = Vec::new(); let mut broken_links = Vec::new(); @@ -293,12 +292,13 @@ fn extract_wikilinks(body: &str) -> Vec { 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; - } + if bytes[i] == b'[' + && bytes[i + 1] == b'[' + && 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; } @@ -466,10 +466,7 @@ mod tests { 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(); + 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")); diff --git a/src/function/mod.rs b/src/function/mod.rs index 703bed8..30b1c40 100644 --- a/src/function/mod.rs +++ b/src/function/mod.rs @@ -20,10 +20,10 @@ use crate::parsers::{bash, python, typescript}; use anyhow::{Context, Result, anyhow, bail}; use indexmap::IndexMap; use indoc::formatdoc; +use memory::MEMORY_FUNCTION_PREFIX; 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; @@ -362,6 +362,11 @@ impl Functions { .extend(memory::memory_function_declarations()); } + pub fn remove_memory_functions(&mut self) { + self.declarations + .retain(|f| !f.name.starts_with(MEMORY_FUNCTION_PREFIX)); + } + pub fn append_skill_functions(&mut self) { self.declarations .extend(skill::skill_function_declarations());