From 876c5556c5bf5da2112900e0e1a323ec0d506ef7 Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Mon, 18 May 2026 14:59:02 -0600 Subject: [PATCH] feat: Created a new .install command to install bundled assets on-demand --- src/cli/mod.rs | 4 ++ src/config/agent.rs | 4 +- src/config/macros.rs | 4 +- src/config/mod.rs | 77 ++++++++++++++++++++++++- src/config/request_context.rs | 102 ++++++++++++++++++++++++++++++++-- src/function/mod.rs | 21 ++++++- src/main.rs | 3 + src/repl/mod.rs | 25 +++++++-- 8 files changed, 222 insertions(+), 18 deletions(-) diff --git a/src/cli/mod.rs b/src/cli/mod.rs index 1b90753..f6aa29b 100644 --- a/src/cli/mod.rs +++ b/src/cli/mod.rs @@ -4,6 +4,7 @@ use crate::cli::completer::{ ShellCompletion, agent_completer, macro_completer, model_completer, rag_completer, role_completer, secrets_completer, session_completer, }; +use crate::config::AssetCategory; use anyhow::{Context, Result}; use clap::ValueHint; use clap::{Parser, crate_authors, crate_description, crate_version}; @@ -82,6 +83,9 @@ pub struct Cli { /// Build all configured Bash tool scripts #[arg(long)] pub build_tools: bool, + /// Reinstall bundled assets, overwriting any local changes + #[arg(long, value_name = "CATEGORY", value_enum)] + pub install: Option, /// Sync models updates #[arg(long)] pub sync_models: bool, diff --git a/src/config/agent.rs b/src/config/agent.rs index 742ad1b..2940952 100644 --- a/src/config/agent.rs +++ b/src/config/agent.rs @@ -45,7 +45,7 @@ pub struct Agent { } impl Agent { - pub fn install_builtin_agents() -> Result<()> { + pub fn install_builtin_agents(force: bool) -> Result<()> { info!( "Installing built-in agents in {}", paths::agents_data_dir().display() @@ -65,7 +65,7 @@ impl Agent { #[cfg_attr(not(unix), expect(unused))] let is_script = matches!(file_extension.as_deref(), Some("sh") | Some("py")); - if file_path.exists() { + if file_path.exists() && !force { debug!( "Agent file already exists, skipping: {}", file_path.display() diff --git a/src/config/macros.rs b/src/config/macros.rs index fcdef5d..289952d 100644 --- a/src/config/macros.rs +++ b/src/config/macros.rs @@ -85,7 +85,7 @@ impl Macro { Ok(value) } - pub fn install_macros() -> Result<()> { + pub fn install_macros(force: bool) -> Result<()> { info!( "Installing built-in macros in {}", paths::macros_dir().display() @@ -98,7 +98,7 @@ impl Macro { let content = unsafe { std::str::from_utf8_unchecked(&embedded_file.data) }; let file_path = paths::macros_dir().join(file.as_ref()); - if file_path.exists() { + if file_path.exists() && !force { debug!( "Macro file already exists, skipping: {}", file_path.display() diff --git a/src/config/mod.rs b/src/config/mod.rs index 32dfa32..0474c62 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -221,12 +221,83 @@ impl Default for Config { } pub fn install_builtins() -> Result<()> { - Functions::install_builtin_global_tools()?; - Agent::install_builtin_agents()?; - Macro::install_macros()?; + Functions::install_builtin_global_tools(false)?; + Agent::install_builtin_agents(false)?; + Macro::install_macros(false)?; Ok(()) } +#[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)] +pub enum AssetCategory { + Agents, + Macros, + Functions, + #[value(name = "mcp_config")] + McpConfig, +} + +impl AssetCategory { + pub const NAMES: [&'static str; 4] = ["agents", "macros", "functions", "mcp_config"]; + + pub fn parse(name: &str) -> Option { + match name { + "agents" => Some(Self::Agents), + "macros" => Some(Self::Macros), + "functions" => Some(Self::Functions), + "mcp_config" => Some(Self::McpConfig), + _ => None, + } + } +} + +pub fn install_assets(category: AssetCategory) -> Result<()> { + let (label, target) = match category { + AssetCategory::Agents => ("agents", paths::agents_data_dir()), + AssetCategory::Macros => ("macros", paths::macros_dir()), + AssetCategory::Functions => ("functions", paths::functions_dir()), + AssetCategory::McpConfig => ("MCP config", paths::mcp_config_file()), + }; + + if !confirm_asset_overwrite(category, label, &target)? { + println!("Aborted. No files were changed."); + return Ok(()); + } + + match category { + AssetCategory::Agents => Agent::install_builtin_agents(true)?, + AssetCategory::Macros => Macro::install_macros(true)?, + AssetCategory::Functions => Functions::install_builtin_global_tools(true)?, + AssetCategory::McpConfig => Functions::install_mcp_config()?, + } + + println!("Reinstalled bundled {label} ({})", target.display()); + Ok(()) +} + +fn confirm_asset_overwrite(category: AssetCategory, label: &str, target: &Path) -> Result { + if !*IS_STDOUT_TERMINAL { + return Ok(true); + } + let body = match category { + AssetCategory::McpConfig => format!( + "This replaces your MCP server configuration at {} with this \ + build's bundled template. Your configured MCP servers (and any \ + custom secret references they contain) will be lost.", + target.display() + ), + _ => format!( + "Reinstalling bundled {label} overwrites every bundled {label} in \ + {} with this build's packaged versions. Local changes to bundled \ + {label} will be lost; {label} you created yourself are left \ + untouched.", + target.display() + ), + }; + let prompt = format!("{} {body}\nContinue? [y/N] ", warning_text("WARNING:")); + let answer = read_single_key(&['y', 'Y', 'n', 'N'], 'n', &prompt)?; + Ok(matches!(answer, 'y' | 'Y')) +} + pub fn default_sessions_dir() -> PathBuf { match env::var(get_env_name("sessions_dir")) { Ok(value) => PathBuf::from(value), diff --git a/src/config/request_context.rs b/src/config/request_context.rs index f80db3a..24f48a8 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -3,10 +3,10 @@ use super::session::Session; use super::todo::TodoList; use super::tool_scope::{McpRuntime, ToolScope}; use super::{ - AGENTS_DIR_NAME, Agent, AgentVariables, AppConfig, AppState, CREATE_TITLE_ROLE, Input, - 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, + AGENTS_DIR_NAME, Agent, AgentVariables, AppConfig, AppState, AssetCategory, CREATE_TITLE_ROLE, + Input, 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, }; use super::{MessageContentToolCalls, prompts}; use crate::client::{Model, ModelType, list_models}; @@ -1784,6 +1784,9 @@ impl RequestContext { } ".rag" => super::map_completion_values(paths::list_rags()), ".agent" => super::map_completion_values(list_agents()), + ".install" => super::map_completion_values( + AssetCategory::NAMES.iter().map(|s| s.to_string()).collect(), + ), ".macro" => super::map_completion_values(paths::list_macros()), ".starter" => match &self.agent { Some(agent) => agent @@ -3676,4 +3679,95 @@ mod tests { "Graph agent must not engage a session, not even an inherited default" ); } + + fn first_file(dir: &Path) -> Option { + for entry in read_dir(dir).ok()?.flatten() { + let path = entry.path(); + if path.is_dir() { + if let Some(found) = first_file(&path) { + return Some(found); + } + } else { + return Some(path); + } + } + None + } + + #[test] + fn asset_category_parse_maps_known_names() { + assert_eq!(AssetCategory::parse("agents"), Some(AssetCategory::Agents)); + assert_eq!(AssetCategory::parse("macros"), Some(AssetCategory::Macros)); + assert_eq!( + AssetCategory::parse("functions"), + Some(AssetCategory::Functions) + ); + assert_eq!( + AssetCategory::parse("mcp_config"), + Some(AssetCategory::McpConfig) + ); + assert_eq!(AssetCategory::parse("roles"), None); + assert_eq!(AssetCategory::parse(""), None); + } + + #[test] + #[serial] + fn install_builtin_agents_force_overwrites_only_with_force() { + let _guard = TestConfigDirGuard::new(); + + Agent::install_builtin_agents(false).unwrap(); + let file = + first_file(&paths::agents_data_dir()).expect("bundled agents should be installed"); + + write(&file, "SENTINEL").unwrap(); + Agent::install_builtin_agents(false).unwrap(); + assert_eq!( + read_to_string(&file).unwrap(), + "SENTINEL", + "non-force install must not overwrite an existing file" + ); + + Agent::install_builtin_agents(true).unwrap(); + assert_ne!( + read_to_string(&file).unwrap(), + "SENTINEL", + "force install must overwrite the existing file" + ); + } + + #[test] + #[serial] + fn install_functions_force_preserves_user_mcp_json() { + let _guard = TestConfigDirGuard::new(); + + Functions::install_builtin_global_tools(false).unwrap(); + let mcp = paths::mcp_config_file(); + assert!(mcp.exists(), "mcp.json should be installed on first run"); + + write(&mcp, "USER_MCP_CONFIG").unwrap(); + Functions::install_builtin_global_tools(true).unwrap(); + assert_eq!( + read_to_string(&mcp).unwrap(), + "USER_MCP_CONFIG", + "force install must NOT overwrite the user's mcp.json" + ); + } + + #[test] + #[serial] + fn install_mcp_config_overwrites_existing() { + let _guard = TestConfigDirGuard::new(); + + Functions::install_mcp_config().unwrap(); + let mcp = paths::mcp_config_file(); + assert!(mcp.exists(), "install_mcp_config should create mcp.json"); + + write(&mcp, "USER_MCP_CONFIG").unwrap(); + Functions::install_mcp_config().unwrap(); + assert_ne!( + read_to_string(&mcp).unwrap(), + "USER_MCP_CONFIG", + "install_mcp_config must overwrite the existing mcp.json" + ); + } } diff --git a/src/function/mod.rs b/src/function/mod.rs index 2fad1fa..0fab009 100644 --- a/src/function/mod.rs +++ b/src/function/mod.rs @@ -203,7 +203,7 @@ pub struct Functions { } impl Functions { - pub fn install_builtin_global_tools() -> Result<()> { + pub fn install_builtin_global_tools(force: bool) -> Result<()> { info!( "Installing global built-in functions in {}", paths::functions_dir().display() @@ -228,7 +228,8 @@ impl Functions { #[cfg_attr(not(unix), expect(unused))] let is_script = matches!(file_extension.as_deref(), Some("sh") | Some("py")); - if file_path.exists() { + let force_this = force && file.as_ref() != "mcp.json"; + if file_path.exists() && !force_this { debug!( "Function file already exists, skipping: {}", file_path.display() @@ -251,6 +252,22 @@ impl Functions { Ok(()) } + pub fn install_mcp_config() -> Result<()> { + let file_path = paths::mcp_config_file(); + let embedded = FunctionAssets::get("mcp.json") + .ok_or_else(|| anyhow!("Failed to load embedded mcp.json"))?; + let content = unsafe { std::str::from_utf8_unchecked(&embedded.data) }; + + ensure_parent_exists(&file_path)?; + + info!("Reinstalling MCP config file: {}", file_path.display()); + + let mut config_file = File::create(&file_path)?; + config_file.write_all(content.as_bytes())?; + + Ok(()) + } + pub fn init(visible_tools: &[String]) -> Result { Self::clear_global_functions_bin_dir()?; diff --git a/src/main.rs b/src/main.rs index 6596898..d6c84c7 100644 --- a/src/main.rs +++ b/src/main.rs @@ -145,6 +145,9 @@ async fn run( text: Option, abort_signal: AbortSignal, ) -> Result<()> { + if let Some(category) = cli.install { + return config::install_assets(category); + } if cli.sync_models { let url = ctx.app.config.sync_models_url(); return sync_models(&url, abort_signal.clone()).await; diff --git a/src/repl/mod.rs b/src/repl/mod.rs index 2d11d35..e47ef9e 100644 --- a/src/repl/mod.rs +++ b/src/repl/mod.rs @@ -7,17 +7,17 @@ use self::highlighter::ReplHighlighter; use self::prompt::ReplPrompt; use crate::client::{call_chat_completions, call_chat_completions_streaming, init_client, oauth}; -use crate::config::paths; use crate::config::{ AgentVariables, AppConfig, AssertState, Input, LastMessage, RequestContext, StateFlags, macro_execute, }; +use crate::config::{AssetCategory, paths}; use crate::render::render_error; use crate::utils::{ AbortSignal, abortable_run_with_spinner, create_abort_signal, dimmed_text, set_text, temp_file, }; -use crate::{graph, resolve_oauth_client}; +use crate::{config, graph, resolve_oauth_client}; use anyhow::{Context, Result, bail}; use crossterm::cursor::SetCursorStyle; use fancy_regex::Regex; @@ -45,7 +45,7 @@ pub const DEFAULT_CONTINUATION_PROMPT: &str = indoc! {" 4. Continue with the next pending item now. Call tools immediately." }; -static REPL_COMMANDS: LazyLock<[ReplCommand; 39]> = LazyLock::new(|| { +static REPL_COMMANDS: LazyLock<[ReplCommand; 40]> = LazyLock::new(|| { [ ReplCommand::new(".help", "Show this help guide", AssertState::pass()), ReplCommand::new(".info", "Show system info", AssertState::pass()), @@ -212,6 +212,11 @@ static REPL_COMMANDS: LazyLock<[ReplCommand; 39]> = LazyLock::new(|| { "View or modify the Loki vault", AssertState::pass(), ), + ReplCommand::new( + ".install", + "Reinstall bundled agents, macros, functions, or MCP config (overwrites local changes)", + AssertState::pass(), + ), ReplCommand::new(".exit", "Exit REPL", AssertState::pass()), ] }); @@ -522,6 +527,16 @@ pub async fn run_repl_command( } } } + ".install" => match args.map(str::trim) { + Some(name) if !name.is_empty() => match AssetCategory::parse(name) { + Some(category) => config::install_assets(category)?, + None => println!( + "Unknown asset category '{name}'. Valid categories: {}", + AssetCategory::NAMES.join(", ") + ), + }, + _ => println!("Usage: .install <{}>", AssetCategory::NAMES.join("|")), + }, ".rag" => { ctx.use_rag(args, abort_signal.clone()).await?; } @@ -1216,8 +1231,8 @@ mod tests { } #[test] - fn repl_commands_has_39_entries() { - assert_eq!(REPL_COMMANDS.len(), 39); + fn repl_commands_has_40_entries() { + assert_eq!(REPL_COMMANDS.len(), 40); } #[test]