From f79da22f0559a137d0c6320116fcc1360e397640 Mon Sep 17 00:00:00 2001 From: Felix dingfeli Date: Fri, 11 Jul 2025 15:20:43 -0700 Subject: [PATCH 01/15] adds alternative format for mcp server field in agent config --- .../chat-cli/src/cli/agent/context_migrate.rs | 5 +- crates/chat-cli/src/cli/agent/mod.rs | 136 ++++++++++--- .../chat-cli/src/cli/agent/wrapper_types.rs | 33 +++ crates/chat-cli/src/cli/chat/mod.rs | 5 +- crates/chat-cli/src/cli/chat/tool_manager.rs | 10 +- crates/chat-cli/src/cli/mcp.rs | 190 ++++++++++++------ 6 files changed, 288 insertions(+), 91 deletions(-) diff --git a/crates/chat-cli/src/cli/agent/context_migrate.rs b/crates/chat-cli/src/cli/agent/context_migrate.rs index 5efc63c552..fbebcb2459 100644 --- a/crates/chat-cli/src/cli/agent/context_migrate.rs +++ b/crates/chat-cli/src/cli/agent/context_migrate.rs @@ -14,6 +14,7 @@ use super::{ }; use crate::cli::agent::{ CreateHooks, + McpServerConfigWrapper, PromptHooks, }; use crate::cli::chat::cli::hooks::{ @@ -180,7 +181,7 @@ impl ContextMigrate<'c'> { included_files: context.paths, create_hooks: CreateHooks::Map(create_hooks), prompt_hooks: PromptHooks::Map(prompt_hooks), - mcp_servers: mcp_servers.clone().unwrap_or_default(), + mcp_servers: McpServerConfigWrapper::Map(mcp_servers.clone().unwrap_or_default()), ..Default::default() }); } @@ -204,7 +205,7 @@ impl ContextMigrate<'c'> { included_files: context.paths, create_hooks: CreateHooks::Map(create_hooks), prompt_hooks: PromptHooks::Map(prompt_hooks), - mcp_servers: mcp_servers.clone().unwrap_or_default(), + mcp_servers: McpServerConfigWrapper::Map(mcp_servers.clone().unwrap_or_default()), ..Default::default() }); } diff --git a/crates/chat-cli/src/cli/agent/mod.rs b/crates/chat-cli/src/cli/agent/mod.rs index ae33c888d2..5e4bfc7a8a 100644 --- a/crates/chat-cli/src/cli/agent/mod.rs +++ b/crates/chat-cli/src/cli/agent/mod.rs @@ -38,6 +38,7 @@ use tracing::{ }; pub use wrapper_types::{ CreateHooks, + McpServerConfigWrapper, OriginalToolName, PromptHooks, ToolSettingTarget, @@ -64,8 +65,30 @@ mod wrapper_types; /// An [Agent] is a declarative way of configuring a given instance of q chat. Currently, it is /// impacting q chat in via influenicng [ContextManager] and [ToolManager]. /// Changes made to [ContextManager] and [ToolManager] do not persist across sessions. +/// +/// To increase the usability of the agent config, (both from the perspective of CLI and the users +/// who would need to write these config), the agent config has two states of existence: "cold" and +/// "warm". +/// +/// A "cold" state describes the config as it is written. And a "warm" state is an alternate form +/// of the same config, modified for the convenience of the business logic that relies on it in the +/// application. +/// +/// For example, the "cold" state does not require the field of "path" to be populated. This is +/// because it would be redundant and tedious for user to have to write the path of the file they +/// had created in said file. This field is thus populated during its parsing. +/// +/// Another example is the mcp config. To support backwards compatibility of users existing global +/// mcp.json, we allow users to supply a list of mcp server names they want the agent to +/// instantiate. But in order for this config to actually be useful to the CLI, it needs to contain +/// actual information about the comamnd, not just the list of names. Thus the "warm" state in this +/// field would be a filtered version of [McpServerConfig], while the "cold" state could be either. +/// +/// Where agents are instantiated from their config, we would need to convert them from "cold" to +/// "warm". #[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, JsonSchema)] #[serde(rename_all = "camelCase")] +#[schemars(description = "An Agent is a declarative way of configuring a given instance of q chat.")] pub struct Agent { /// Agent names are derived from the file name. Thus they are skipped for /// serializing @@ -80,7 +103,7 @@ pub struct Agent { pub prompt: Option, /// Configuration for Model Context Protocol (MCP) servers #[serde(default)] - pub mcp_servers: McpServerConfig, + pub mcp_servers: McpServerConfigWrapper, /// List of tools the agent can see. Use \"@{MCP_SERVER_NAME}/tool_name\" to specify tools from /// mcp servers. To include all tools from a server, use \"@{MCP_SERVER_NAME}\" #[serde(default)] @@ -143,9 +166,9 @@ impl Agent { pub async fn get_agent_by_name(os: &Os, agent_name: &str) -> eyre::Result<(Agent, PathBuf)> { let config_path: Result = 'config: { // local first, and then fall back to looking at global - let local_config_dir = directories::chat_local_agent_dir()?.join(agent_name); + let local_config_dir = directories::chat_local_agent_dir()?.join(format!("{agent_name}.json")); if os.fs.exists(&local_config_dir) { - break 'config Ok::(local_config_dir); + break 'config Ok(local_config_dir); } let global_config_dir = directories::chat_global_agent_path(os)?.join(format!("{agent_name}.json")); @@ -159,23 +182,35 @@ impl Agent { match config_path { Ok(config_path) => { let content = os.fs.read(&config_path).await?; - Ok((serde_json::from_slice::(&content)?, config_path)) - }, - Err(global_config_dir) if agent_name == "default" => { - os.fs - .create_dir_all( - global_config_dir - .parent() - .ok_or(eyre::eyre!("Failed to retrieve global agent config parent path"))?, - ) - .await?; - os.fs.create_new(&global_config_dir).await?; - - let default_agent = Agent::default(); - let content = serde_json::to_string_pretty(&default_agent)?; - os.fs.write(&global_config_dir, content.as_bytes()).await?; - - Ok((default_agent, global_config_dir)) + let mut agent = serde_json::from_slice::(&content)?; + + if let McpServerConfigWrapper::List(selected_servers) = &agent.mcp_servers { + let global_mcp_config_path = directories::chat_legacy_mcp_config(os)?; + let global_mcp_config = McpServerConfig::load_from_file(os, global_mcp_config_path).await?; + + let filtered_config = global_mcp_config + .mcp_servers + .iter() + .filter_map(|(name, config)| { + if selected_servers.contains(name) + || (selected_servers.len() == 1 + && selected_servers.first().is_some_and(|s| s.as_str() == "*")) + { + Some((name.clone(), config.clone())) + } else { + None + } + }) + .collect::>(); + + agent.mcp_servers = McpServerConfigWrapper::Map(McpServerConfig { + mcp_servers: filtered_config, + }); + } + + agent.name = agent_name.to_string(); + + Ok((agent, config_path)) }, _ => bail!("Agent {agent_name} does not exist"), } @@ -324,6 +359,8 @@ impl Agents { agent_name.replace(name.as_str()); } + let mut global_mcp_config = None::; + let mut local_agents = 'local: { // We could be launching from the home dir, in which case the global and local agents // are the same set of agents. If that is the case, we simply skip this. @@ -341,7 +378,7 @@ impl Agents { let Ok(files) = os.fs.read_dir(path).await else { break 'local Vec::::new(); }; - load_agents_from_entries(files).await + load_agents_from_entries(files, os, &mut global_mcp_config).await }; let mut global_agents = 'global: { @@ -359,7 +396,7 @@ impl Agents { break 'global Vec::::new(); }, }; - load_agents_from_entries(files).await + load_agents_from_entries(files, os, &mut global_mcp_config).await } .into_iter() .chain(new_agents) @@ -539,7 +576,11 @@ impl Agents { } } -async fn load_agents_from_entries(mut files: ReadDir) -> Vec { +async fn load_agents_from_entries( + mut files: ReadDir, + os: &Os, + global_mcp_config: &mut Option, +) -> Vec { let mut res = Vec::::new(); while let Ok(Some(file)) = files.next_entry().await { let file_path = &file.path(); @@ -556,6 +597,7 @@ async fn load_agents_from_entries(mut files: ReadDir) -> Vec { continue; }, }; + let mut agent = match serde_json::from_slice::(&content) { Ok(mut agent) => { agent.path = Some(file_path.clone()); @@ -567,6 +609,47 @@ async fn load_agents_from_entries(mut files: ReadDir) -> Vec { continue; }, }; + + if let McpServerConfigWrapper::List(selected_servers) = &agent.mcp_servers { + let global_mcp_config = match global_mcp_config.as_mut() { + Some(config) => config, + None => { + let Ok(global_mcp_path) = directories::chat_legacy_mcp_config(os) else { + tracing::error!("Error obtaining legacy mcp json path. Skipping"); + continue; + }; + let legacy_mcp_config = match McpServerConfig::load_from_file(os, global_mcp_path).await { + Ok(config) => config, + Err(e) => { + tracing::error!("Error loading global mcp json path: {e}. Skipping"); + continue; + }, + }; + global_mcp_config.replace(legacy_mcp_config); + global_mcp_config.as_mut().unwrap() + }, + }; + + let filtered_config = global_mcp_config + .mcp_servers + .iter() + .filter_map(|(name, config)| { + if selected_servers.contains(name) + || (selected_servers.len() == 1 + && selected_servers.first().is_some_and(|s| s.as_str() == "*")) + { + Some((name.clone(), config.clone())) + } else { + None + } + }) + .collect::>(); + + agent.mcp_servers = McpServerConfigWrapper::Map(McpServerConfig { + mcp_servers: filtered_config, + }); + } + if let Some(name) = Path::new(&file.file_name()).file_stem() { agent.name = name.to_string_lossy().to_string(); res.push(agent); @@ -662,8 +745,11 @@ mod tests { #[test] fn test_deser() { let agent = serde_json::from_str::(INPUT).expect("Deserializtion failed"); - assert!(agent.mcp_servers.mcp_servers.contains_key("fetch")); - assert!(agent.mcp_servers.mcp_servers.contains_key("git")); + assert!(matches!(agent.mcp_servers, McpServerConfigWrapper::Map(_))); + if let McpServerConfigWrapper::Map(config) = &agent.mcp_servers { + assert!(config.mcp_servers.contains_key("fetch")); + assert!(config.mcp_servers.contains_key("git")); + } assert!(agent.alias.contains_key("@gits/some_tool")); } diff --git a/crates/chat-cli/src/cli/agent/wrapper_types.rs b/crates/chat-cli/src/cli/agent/wrapper_types.rs index a9fefba32e..08b7d0fdf2 100644 --- a/crates/chat-cli/src/cli/agent/wrapper_types.rs +++ b/crates/chat-cli/src/cli/agent/wrapper_types.rs @@ -13,6 +13,7 @@ use serde::{ Serialize, }; +use super::McpServerConfig; use crate::cli::chat::cli::hooks::Hook; /// Subject of the tool name change. For tools in mcp servers, you would need to prefix them with @@ -122,3 +123,35 @@ pub fn tool_settings_schema(generator: &mut SchemaGenerator) -> Schema { } }) } + +#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, JsonSchema)] +#[serde(untagged)] +pub enum McpServerConfigWrapper { + /// Array of servers to instantiate in accordance to what is in the global mcp.json + List(Vec), + /// A map of mcp server configurations, identical in format to the value of mcpServers in + /// the global mcp.json + Map(McpServerConfig), +} + +impl Default for McpServerConfigWrapper { + fn default() -> Self { + Self::Map(Default::default()) + } +} + +impl McpServerConfigWrapper { + pub fn get_server_names(&self) -> Vec<&str> { + match self { + Self::List(list) => list.iter().map(|name| name.as_str()).collect::>(), + Self::Map(config) => config.mcp_servers.keys().map(|name| name.as_str()).collect::>(), + } + } + + pub fn is_empty(&self) -> bool { + match self { + Self::List(list) => list.is_empty(), + Self::Map(config) => config.mcp_servers.is_empty(), + } + } +} diff --git a/crates/chat-cli/src/cli/chat/mod.rs b/crates/chat-cli/src/cli/chat/mod.rs index 2c60d5a5fb..56ab3a0e04 100644 --- a/crates/chat-cli/src/cli/chat/mod.rs +++ b/crates/chat-cli/src/cli/chat/mod.rs @@ -233,10 +233,7 @@ impl ChatArgs { let mut agents = Agents::load(os, self.agent.as_deref(), skip_migration, &mut stderr).await; agents.trust_all_tools = self.trust_all_tools; - if agents - .get_active() - .is_some_and(|a| !a.mcp_servers.mcp_servers.is_empty()) - { + if agents.get_active().is_some_and(|a| !a.mcp_servers.is_empty()) { if !self.no_interactive && !os.database.settings.get_bool(Setting::McpLoadedBefore).unwrap_or(false) { execute!( stderr, diff --git a/crates/chat-cli/src/cli/chat/tool_manager.rs b/crates/chat-cli/src/cli/chat/tool_manager.rs index a672033632..a14198ea1e 100644 --- a/crates/chat-cli/src/cli/chat/tool_manager.rs +++ b/crates/chat-cli/src/cli/chat/tool_manager.rs @@ -61,6 +61,7 @@ use crate::api_client::model::{ use crate::cli::agent::{ Agent, McpServerConfig, + McpServerConfigWrapper, }; use crate::cli::chat::cli::prompts::GetPromptError; use crate::cli::chat::message::AssistantToolUse; @@ -173,7 +174,14 @@ impl ToolManagerBuilder { } pub fn agent(mut self, agent: Agent) -> Self { - self.mcp_server_config.replace(agent.mcp_servers.clone()); + if let McpServerConfigWrapper::Map(config) = &agent.mcp_servers { + self.mcp_server_config.replace(config.clone()); + } else { + error!( + "No valid mcp config valid in agent {}, no mcp config loaded", + &agent.name + ); + } self.agent.replace(agent); self } diff --git a/crates/chat-cli/src/cli/mcp.rs b/crates/chat-cli/src/cli/mcp.rs index d1cf785159..142244d66a 100644 --- a/crates/chat-cli/src/cli/mcp.rs +++ b/crates/chat-cli/src/cli/mcp.rs @@ -22,6 +22,7 @@ use super::agent::{ Agents, McpServerConfig, }; +use crate::cli::agent::McpServerConfigWrapper; use crate::cli::chat::tool_manager::{ global_mcp_config_path, workspace_mcp_config_path, @@ -89,7 +90,8 @@ pub struct AddArgs { /// Arguments to pass to the command #[arg(long, action = ArgAction::Append, allow_hyphen_values = true, value_delimiter = ',')] pub args: Vec, - /// Where to add the server to. + /// Where to add the server to. If an agent name is not supplied, the changes shall be made to + /// the global mcp.json #[arg(long)] pub agent: Option, /// Environment variables to use when launching the server @@ -108,37 +110,71 @@ pub struct AddArgs { impl AddArgs { pub async fn execute(self, os: &Os, output: &mut impl Write) -> Result<()> { - let agent_name = self.agent.as_deref().unwrap_or("default"); - let (mut agent, config_path) = Agent::get_agent_by_name(os, agent_name).await?; - - let mcp_servers = &mut agent.mcp_servers.mcp_servers; - if mcp_servers.contains_key(&self.name) && !self.force { - bail!( - "\nMCP server '{}' already exists in agent {} (path {}). Use --force to overwrite.", - self.name, - agent_name, - config_path.display(), - ); - } - - let merged_env = self.env.into_iter().flatten().collect::>(); - let tool: CustomToolConfig = serde_json::from_value(serde_json::json!({ - "command": self.command, - "args": self.args, - "env": merged_env, - "timeout": self.timeout.unwrap_or(default_timeout()), - "disabled": self.disabled, - }))?; - - writeln!( - output, - "\nTo learn more about MCP safety, see https://docs.aws.amazon.com/amazonq/latest/qdeveloper-ug/command-line-mcp-security.html\n\n" - )?; + match self.agent.as_deref() { + Some(agent_name) => { + let (mut agent, config_path) = Agent::get_agent_by_name(os, agent_name).await?; + let mcp_servers = if let McpServerConfigWrapper::Map(config) = &mut agent.mcp_servers { + &mut config.mcp_servers + } else { + bail!("Mcp server config not found on agent"); + }; + + if mcp_servers.contains_key(&self.name) && !self.force { + bail!( + "\nMCP server '{}' already exists in agent {} (path {}). Use --force to overwrite.", + self.name, + agent_name, + config_path.display(), + ); + } + + let merged_env = self.env.into_iter().flatten().collect::>(); + let tool: CustomToolConfig = serde_json::from_value(serde_json::json!({ + "command": self.command, + "args": self.args, + "env": merged_env, + "timeout": self.timeout.unwrap_or(default_timeout()), + "disabled": self.disabled, + }))?; + + mcp_servers.insert(self.name.clone(), tool); + let json = serde_json::to_string_pretty(&agent)?; + os.fs.write(config_path, json).await?; + writeln!(output, "✓ Added MCP server '{}' to agent {}\n", self.name, agent_name)?; + }, + None => { + let global_config_path = directories::chat_legacy_mcp_config(os)?; + let mut mcp_servers = McpServerConfig::load_from_file(os, &global_config_path).await?; + + if mcp_servers.mcp_servers.contains_key(&self.name) && !self.force { + bail!( + "\nMCP server '{}' already exists in global config (path {}). Use --force to overwrite.", + self.name, + &global_config_path.display(), + ); + } + + let merged_env = self.env.into_iter().flatten().collect::>(); + let tool: CustomToolConfig = serde_json::from_value(serde_json::json!({ + "command": self.command, + "args": self.args, + "env": merged_env, + "timeout": self.timeout.unwrap_or(default_timeout()), + "disabled": self.disabled, + }))?; + + mcp_servers.mcp_servers.insert(self.name.clone(), tool); + let json = serde_json::to_string_pretty(&mcp_servers)?; + os.fs.write(&global_config_path, json).await?; + writeln!( + output, + "✓ Added MCP server '{}' to global config in {}\n", + self.name, + global_config_path.display() + )?; + }, + }; - mcp_servers.insert(self.name.clone(), tool); - let json = serde_json::to_string_pretty(&agent)?; - os.fs.write(config_path, json).await?; - writeln!(output, "✓ Added MCP server '{}' to agent {}\n", self.name, agent_name)?; Ok(()) } } @@ -153,33 +189,67 @@ pub struct RemoveArgs { impl RemoveArgs { pub async fn execute(self, os: &Os, output: &mut impl Write) -> Result<()> { - let agent_name = self.agent.as_deref().unwrap_or("default"); - let (mut agent, config_path) = Agent::get_agent_by_name(os, agent_name).await?; - - if !os.fs.exists(&config_path) { - writeln!(output, "\nNo MCP server configurations found.\n")?; - return Ok(()); - } - - let config = &mut agent.mcp_servers.mcp_servers; - match config.remove(&self.name) { - Some(_) => { - let json = serde_json::to_string_pretty(&agent)?; - os.fs.write(config_path, json).await?; - writeln!( - output, - "\n✓ Removed MCP server '{}' from agent {}\n", - self.name, agent_name, - )?; + match self.agent.as_deref() { + Some(agent_name) => { + let (mut agent, config_path) = Agent::get_agent_by_name(os, agent_name).await?; + + if !os.fs.exists(&config_path) { + writeln!(output, "\nNo MCP server configurations found.\n")?; + return Ok(()); + } + + let config = if let McpServerConfigWrapper::Map(config) = &mut agent.mcp_servers { + &mut config.mcp_servers + } else { + bail!("Mcp server config not found on agent"); + }; + + match config.remove(&self.name) { + Some(_) => { + let json = serde_json::to_string_pretty(&agent)?; + os.fs.write(config_path, json).await?; + writeln!( + output, + "\n✓ Removed MCP server '{}' from agent {}\n", + self.name, agent_name, + )?; + }, + None => { + writeln!( + output, + "\nNo MCP server named '{}' found in agent {}\n", + self.name, agent_name, + )?; + }, + } }, None => { - writeln!( - output, - "\nNo MCP server named '{}' found in agent {}\n", - self.name, agent_name, - )?; + let global_config_path = directories::chat_legacy_mcp_config(os)?; + let mut config = McpServerConfig::load_from_file(os, &global_config_path).await?; + + match config.mcp_servers.remove(&self.name) { + Some(_) => { + let json = serde_json::to_string_pretty(&config)?; + os.fs.write(&global_config_path, json).await?; + writeln!( + output, + "\n✓ Removed MCP server '{}' from global config (path {})\n", + self.name, + &global_config_path.display(), + )?; + }, + None => { + writeln!( + output, + "\nNo MCP server named '{}' found in global config (path {})\n", + self.name, + &global_config_path.display(), + )?; + }, + } }, } + Ok(()) } } @@ -335,11 +405,13 @@ async fn get_mcp_server_configs( } else { Scope::Workspace }; - results.push(( - scope, - agent.path.ok_or(eyre::eyre!("Agent missing path info"))?, - Some(agent.mcp_servers), - )); + if let McpServerConfigWrapper::Map(config) = agent.mcp_servers { + results.push(( + scope, + agent.path.ok_or(eyre::eyre!("Agent missing path info"))?, + Some(config), + )); + } } Ok(results) } From c42006e3dab69290525bc246cda37410dbc930fb Mon Sep 17 00:00:00 2001 From: Felix dingfeli Date: Fri, 11 Jul 2025 15:52:31 -0700 Subject: [PATCH 02/15] adjusts migration flow for mcp backwards compatibility --- crates/chat-cli/src/cli/agent/context_migrate.rs | 15 +-------------- crates/chat-cli/src/cli/agent/mod.rs | 2 +- crates/chat-cli/src/cli/agent/wrapper_types.rs | 2 +- 3 files changed, 3 insertions(+), 16 deletions(-) diff --git a/crates/chat-cli/src/cli/agent/context_migrate.rs b/crates/chat-cli/src/cli/agent/context_migrate.rs index fbebcb2459..7efef32763 100644 --- a/crates/chat-cli/src/cli/agent/context_migrate.rs +++ b/crates/chat-cli/src/cli/agent/context_migrate.rs @@ -181,7 +181,6 @@ impl ContextMigrate<'c'> { included_files: context.paths, create_hooks: CreateHooks::Map(create_hooks), prompt_hooks: PromptHooks::Map(prompt_hooks), - mcp_servers: McpServerConfigWrapper::Map(mcp_servers.clone().unwrap_or_default()), ..Default::default() }); } @@ -205,7 +204,6 @@ impl ContextMigrate<'c'> { included_files: context.paths, create_hooks: CreateHooks::Map(create_hooks), prompt_hooks: PromptHooks::Map(prompt_hooks), - mcp_servers: McpServerConfigWrapper::Map(mcp_servers.clone().unwrap_or_default()), ..Default::default() }); } @@ -214,19 +212,7 @@ impl ContextMigrate<'c'> { os.fs.create_dir_all(&global_agent_path).await?; } - let formatted_server_list = mcp_servers - .map(|config| { - config - .mcp_servers - .keys() - .map(|server_name| format!("@{server_name}")) - .collect::>() - }) - .unwrap_or_default(); - for agent in &mut new_agents { - agent.tools.extend(formatted_server_list.clone()); - let content = serde_json::to_string_pretty(agent)?; if let Some(path) = agent.path.as_ref() { info!("Agent {} peristed in path {}", agent.name, path.to_string_lossy()); @@ -237,6 +223,7 @@ impl ContextMigrate<'c'> { agent.name ); } + agent.mcp_servers = McpServerConfigWrapper::Map(mcp_servers.clone().unwrap_or_default()); } let legacy_profile_config_path = directories::chat_profiles_dir(os)?; diff --git a/crates/chat-cli/src/cli/agent/mod.rs b/crates/chat-cli/src/cli/agent/mod.rs index 5e4bfc7a8a..76b0cfbd9c 100644 --- a/crates/chat-cli/src/cli/agent/mod.rs +++ b/crates/chat-cli/src/cli/agent/mod.rs @@ -140,7 +140,7 @@ impl Default for Agent { description: Some("Default agent".to_string()), prompt: Default::default(), mcp_servers: Default::default(), - tools: NATIVE_TOOLS.iter().copied().map(str::to_string).collect::>(), + tools: vec!["*".to_string()], alias: Default::default(), allowed_tools: { let mut set = HashSet::::new(); diff --git a/crates/chat-cli/src/cli/agent/wrapper_types.rs b/crates/chat-cli/src/cli/agent/wrapper_types.rs index 08b7d0fdf2..2c763a6e4f 100644 --- a/crates/chat-cli/src/cli/agent/wrapper_types.rs +++ b/crates/chat-cli/src/cli/agent/wrapper_types.rs @@ -136,7 +136,7 @@ pub enum McpServerConfigWrapper { impl Default for McpServerConfigWrapper { fn default() -> Self { - Self::Map(Default::default()) + Self::List(vec!["*".to_string()]) } } From 9221f4423ad6612f8452944e30e4a4d765990095 Mon Sep 17 00:00:00 2001 From: Felix dingfeli Date: Fri, 11 Jul 2025 15:54:15 -0700 Subject: [PATCH 03/15] fixes typo --- crates/chat-cli/src/cli/agent/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/chat-cli/src/cli/agent/mod.rs b/crates/chat-cli/src/cli/agent/mod.rs index 76b0cfbd9c..be6f7ae39e 100644 --- a/crates/chat-cli/src/cli/agent/mod.rs +++ b/crates/chat-cli/src/cli/agent/mod.rs @@ -81,7 +81,7 @@ mod wrapper_types; /// Another example is the mcp config. To support backwards compatibility of users existing global /// mcp.json, we allow users to supply a list of mcp server names they want the agent to /// instantiate. But in order for this config to actually be useful to the CLI, it needs to contain -/// actual information about the comamnd, not just the list of names. Thus the "warm" state in this +/// actual information about the command, not just the list of names. Thus the "warm" state in this /// field would be a filtered version of [McpServerConfig], while the "cold" state could be either. /// /// Where agents are instantiated from their config, we would need to convert them from "cold" to From 73c1e28bbeabdf3fc1377de4841af7a8a5259f73 Mon Sep 17 00:00:00 2001 From: Felix dingfeli Date: Fri, 11 Jul 2025 16:22:46 -0700 Subject: [PATCH 04/15] implements @native --- crates/chat-cli/src/cli/chat/tool_manager.rs | 34 ++++++++++++++++---- 1 file changed, 28 insertions(+), 6 deletions(-) diff --git a/crates/chat-cli/src/cli/chat/tool_manager.rs b/crates/chat-cli/src/cli/chat/tool_manager.rs index a14198ea1e..803ca5bea6 100644 --- a/crates/chat-cli/src/cli/chat/tool_manager.rs +++ b/crates/chat-cli/src/cli/chat/tool_manager.rs @@ -213,9 +213,31 @@ impl ToolManagerBuilder { if server_name.contains(MCP_SERVER_TOOL_DELIMITER) { let _ = queue!( output, - style::Print(format!( - "Invalid server name {server_name}. Server name cannot contain {MCP_SERVER_TOOL_DELIMITER}\n" - )) + style::SetForegroundColor(style::Color::Red), + style::Print("✗ Invalid server name "), + style::SetForegroundColor(style::Color::Blue), + style::Print(&server_name), + style::ResetColor, + style::Print(". Server name cannot contain "), + style::SetForegroundColor(style::Color::Yellow), + style::Print(MCP_SERVER_TOOL_DELIMITER), + style::ResetColor, + style::Print("\n") + ); + None + } else if server_name == "native" { + let _ = queue!( + output, + style::SetForegroundColor(style::Color::Red), + style::Print("✗ Invalid server name "), + style::SetForegroundColor(style::Color::Blue), + style::Print(&server_name), + style::ResetColor, + style::Print(". Server name cannot contain reserved word "), + style::SetForegroundColor(style::Color::Yellow), + style::Print("native"), + style::ResetColor, + style::Print(" (it is used to denote native tools)\n") ); None } else { @@ -873,12 +895,12 @@ impl ToolManager { let notify = self.notify.take(); self.schema = { let tool_list = &self.agent.lock().await.tools; + let is_allow_all = tool_list.len() == 1 && tool_list.first().is_some_and(|n| n == "*"); + let is_allow_native = tool_list.iter().any(|t| t.as_str() == "@native"); let mut tool_specs = serde_json::from_str::>(include_str!("tools/tool_index.json"))? .into_iter() - .filter(|(name, _)| { - tool_list.len() == 1 && tool_list.first().is_some_and(|n| n == "*") || tool_list.contains(name) - }) + .filter(|(name, _)| is_allow_all || is_allow_native || tool_list.contains(name)) .collect::>(); if !crate::cli::chat::tools::thinking::Thinking::is_enabled(os) { tool_specs.remove("thinking"); From f94736fb0e42ee3a0180f0f9f3896e2c6aa9820b Mon Sep 17 00:00:00 2001 From: Felix dingfeli Date: Mon, 14 Jul 2025 10:28:44 -0700 Subject: [PATCH 05/15] WIP: utilizes type check to ensure invariant --- crates/chat-cli/src/cli/agent/mod.rs | 33 ++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/crates/chat-cli/src/cli/agent/mod.rs b/crates/chat-cli/src/cli/agent/mod.rs index be6f7ae39e..df249d9450 100644 --- a/crates/chat-cli/src/cli/agent/mod.rs +++ b/crates/chat-cli/src/cli/agent/mod.rs @@ -939,3 +939,36 @@ mod tests { println!("Schema for agent: {}", serde_json::to_string_pretty(&schema).unwrap()); } } + +#[cfg(test)] +mod test_test { + use super::*; + trait A: for<'de> Deserialize<'de> + Serialize {} + #[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, JsonSchema)] + #[serde(rename_all = "camelCase", bound = "T: A")] + struct Testtest { + hello: String, + #[serde(skip)] + phantom: std::marker::PhantomData, + } + + #[derive(Serialize, Deserialize)] + struct ConcreteA; + + impl A for ConcreteA {} + + #[derive(Serialize, Deserialize)] + struct ConcreteB; + + const INPUT: &str = r#" + { + "hello": "world" + } + "#; + + #[test] + fn test_deser_err() { + let deser_res = serde_json::from_str::>(INPUT); + assert!(deser_res.is_ok()); + } +} From 5a3d0ac3218cb477ca919562278097be5d078e6b Mon Sep 17 00:00:00 2001 From: Felix dingfeli Date: Tue, 15 Jul 2025 17:27:02 -0700 Subject: [PATCH 06/15] makes agent generic to prevent api misuse --- .../chat-cli/src/cli/agent/context_migrate.rs | 7 +- crates/chat-cli/src/cli/agent/mcp_config.rs | 17 ++ crates/chat-cli/src/cli/agent/mod.rs | 260 +++++++++++------- .../chat-cli/src/cli/agent/wrapper_types.rs | 11 + crates/chat-cli/src/cli/chat/cli/profile.rs | 3 +- crates/chat-cli/src/cli/chat/cli/tools.rs | 7 +- crates/chat-cli/src/cli/chat/context.rs | 13 +- crates/chat-cli/src/cli/chat/tool_manager.rs | 10 +- .../src/cli/chat/tools/custom_tool.rs | 3 +- .../src/cli/chat/tools/execute/mod.rs | 3 +- crates/chat-cli/src/cli/chat/tools/fs_read.rs | 3 +- .../chat-cli/src/cli/chat/tools/fs_write.rs | 3 +- crates/chat-cli/src/cli/chat/tools/mod.rs | 3 +- crates/chat-cli/src/cli/chat/tools/use_aws.rs | 3 +- crates/chat-cli/src/cli/mcp.rs | 4 +- 15 files changed, 232 insertions(+), 118 deletions(-) diff --git a/crates/chat-cli/src/cli/agent/context_migrate.rs b/crates/chat-cli/src/cli/agent/context_migrate.rs index 7efef32763..e8cfd5c523 100644 --- a/crates/chat-cli/src/cli/agent/context_migrate.rs +++ b/crates/chat-cli/src/cli/agent/context_migrate.rs @@ -8,6 +8,7 @@ use tracing::{ warn, }; +use super::wrapper_types::Warm; use super::{ Agent, McpServerConfig, @@ -30,7 +31,7 @@ pub(in crate::cli::agent) struct ContextMigrate { legacy_global_context: Option, legacy_profiles: HashMap, mcp_servers: Option, - new_agents: Vec, + new_agents: Vec>, } impl ContextMigrate<'a'> { @@ -213,7 +214,7 @@ impl ContextMigrate<'c'> { } for agent in &mut new_agents { - let content = serde_json::to_string_pretty(agent)?; + let content = serde_json::to_string_pretty(&agent.clone().freeze())?; if let Some(path) = agent.path.as_ref() { info!("Agent {} peristed in path {}", agent.name, path.to_string_lossy()); os.fs.write(path, content).await?; @@ -258,7 +259,7 @@ impl ContextMigrate<'c'> { } impl ContextMigrate<'d'> { - pub async fn prompt_set_default(self, os: &mut Os) -> eyre::Result<(Option, Vec)> { + pub async fn prompt_set_default(self, os: &mut Os) -> eyre::Result<(Option, Vec>)> { let ContextMigrate { new_agents, .. } = self; let labels = new_agents diff --git a/crates/chat-cli/src/cli/agent/mcp_config.rs b/crates/chat-cli/src/cli/agent/mcp_config.rs index e6adfc10cd..c848904391 100644 --- a/crates/chat-cli/src/cli/agent/mcp_config.rs +++ b/crates/chat-cli/src/cli/agent/mcp_config.rs @@ -10,11 +10,28 @@ use serde::{ use crate::cli::chat::tools::custom_tool::CustomToolConfig; use crate::os::Os; +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub enum OriginalForm { + /// The mcp config has been written in isolation, with no reference to global mcp.json + #[default] + Map, + /// The mcp config has been written as an include-all array, i.e. '["*"]' + All, + /// The mcp config has been written as an array containing a list of mcp servers to include + /// from global mcp.json + List, +} + // This is to mirror claude's config set up #[derive(Clone, Serialize, Deserialize, Debug, Default, Eq, PartialEq, JsonSchema)] #[serde(rename_all = "camelCase", transparent)] pub struct McpServerConfig { pub mcp_servers: HashMap, + /// This field keeps track of how the mcp server field is configured as seen on the user's + /// config file. This is important because when we write the config to file, we want to + /// preserve its behavior the next time it is read. + #[serde(skip)] + pub original_config_form: OriginalForm, } impl McpServerConfig { diff --git a/crates/chat-cli/src/cli/agent/mod.rs b/crates/chat-cli/src/cli/agent/mod.rs index df249d9450..7e1744bc11 100644 --- a/crates/chat-cli/src/cli/agent/mod.rs +++ b/crates/chat-cli/src/cli/agent/mod.rs @@ -9,10 +9,8 @@ use std::io::{ self, Write, }; -use std::path::{ - Path, - PathBuf, -}; +use std::marker::PhantomData; +use std::path::PathBuf; use context_migrate::ContextMigrate; use crossterm::style::{ @@ -25,6 +23,7 @@ use crossterm::{ }; use eyre::bail; pub use mcp_config::McpServerConfig; +use mcp_config::OriginalForm; use regex::Regex; use schemars::JsonSchema; use serde::{ @@ -36,11 +35,16 @@ use tracing::{ error, warn, }; +use wrapper_types::{ + Cold, + Warm, +}; pub use wrapper_types::{ CreateHooks, McpServerConfigWrapper, OriginalToolName, PromptHooks, + SerdeReady, ToolSettingTarget, alias_schema, tool_settings_schema, @@ -60,7 +64,7 @@ use crate::util::{ mod context_migrate; mod mcp_config; -mod wrapper_types; +pub mod wrapper_types; /// An [Agent] is a declarative way of configuring a given instance of q chat. Currently, it is /// impacting q chat in via influenicng [ContextManager] and [ToolManager]. @@ -87,9 +91,9 @@ mod wrapper_types; /// Where agents are instantiated from their config, we would need to convert them from "cold" to /// "warm". #[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, JsonSchema)] -#[serde(rename_all = "camelCase")] +#[serde(rename_all = "camelCase", bound = "T: SerdeReady")] #[schemars(description = "An Agent is a declarative way of configuring a given instance of q chat.")] -pub struct Agent { +pub struct Agent { /// Agent names are derived from the file name. Thus they are skipped for /// serializing #[serde(skip)] @@ -131,9 +135,11 @@ pub struct Agent { pub tools_settings: HashMap, #[serde(skip)] pub path: Option, + #[serde(skip)] + phantom: std::marker::PhantomData, } -impl Default for Agent { +impl Default for Agent { fn default() -> Self { Self { name: "default".to_string(), @@ -156,14 +162,116 @@ impl Default for Agent { prompt_hooks: Default::default(), tools_settings: Default::default(), path: None, + phantom: PhantomData, } } } -impl Agent { +impl Agent { + pub fn thaw(mut self, path: &PathBuf, global_mcp_config: Option<&McpServerConfig>) -> eyre::Result> { + let Self { mcp_servers, .. } = &mut self; + + let name = path + .file_stem() + .ok_or(eyre::eyre!("Missing valid file name"))? + .to_string_lossy() + .to_string(); + + if let McpServerConfigWrapper::List(selected_servers) = mcp_servers { + let include_all = + selected_servers.len() == 1 && selected_servers.first().is_some_and(|s| s.as_str() == "*"); + let filtered_config = if let Some(global_mcp_config) = global_mcp_config { + global_mcp_config + .mcp_servers + .iter() + .filter_map(|(name, config)| { + if selected_servers.contains(name) || include_all { + Some((name.clone(), config.clone())) + } else { + None + } + }) + .collect::>() + } else { + bail!( + "Error converting agent {name} to usable state. Mcp config is given as list but global mcp config is missing. Aborting" + ); + }; + + *mcp_servers = McpServerConfigWrapper::Map(McpServerConfig { + mcp_servers: filtered_config, + original_config_form: if include_all { + OriginalForm::All + } else { + OriginalForm::Map + }, + }); + } + + Ok(Agent { + name, + description: self.description, + prompt: self.prompt, + mcp_servers: self.mcp_servers, + tools: self.tools, + alias: self.alias, + allowed_tools: self.allowed_tools, + included_files: self.included_files, + create_hooks: self.create_hooks, + prompt_hooks: self.prompt_hooks, + tools_settings: self.tools_settings, + path: Some(path.clone()), + phantom: PhantomData, + }) + } +} + +impl Agent { + pub fn freeze(mut self) -> Agent { + let Self { mcp_servers, .. } = &mut self; + + if let McpServerConfigWrapper::Map(servers) = mcp_servers { + match servers.original_config_form { + OriginalForm::All => { + *mcp_servers = McpServerConfigWrapper::List(vec!["*".to_string()]); + }, + OriginalForm::List => { + *mcp_servers = McpServerConfigWrapper::List({ + servers + .mcp_servers + .iter() + .fold(Vec::::new(), |mut acc, (name, _)| { + acc.push(name.clone()); + acc + }) + }); + }, + OriginalForm::Map => { + // noop + }, + } + } + + Agent { + name: self.name, + description: self.description, + prompt: self.prompt, + mcp_servers: self.mcp_servers, + tools: self.tools, + alias: self.alias, + allowed_tools: self.allowed_tools, + included_files: self.included_files, + create_hooks: self.create_hooks, + prompt_hooks: self.prompt_hooks, + tools_settings: self.tools_settings, + path: self.path, + phantom: PhantomData, + } + } + /// Retrieves an agent by name. It does so via first seeking the given agent under local dir, /// and falling back to global dir if it does not exist in local. - pub async fn get_agent_by_name(os: &Os, agent_name: &str) -> eyre::Result<(Agent, PathBuf)> { + pub async fn get_agent_by_name(os: &Os, agent_name: &str) -> eyre::Result<(Agent, PathBuf)> { let config_path: Result = 'config: { // local first, and then fall back to looking at global let local_config_dir = directories::chat_local_agent_dir()?.join(format!("{agent_name}.json")); @@ -182,35 +290,17 @@ impl Agent { match config_path { Ok(config_path) => { let content = os.fs.read(&config_path).await?; - let mut agent = serde_json::from_slice::(&content)?; - - if let McpServerConfigWrapper::List(selected_servers) = &agent.mcp_servers { - let global_mcp_config_path = directories::chat_legacy_mcp_config(os)?; - let global_mcp_config = McpServerConfig::load_from_file(os, global_mcp_config_path).await?; - - let filtered_config = global_mcp_config - .mcp_servers - .iter() - .filter_map(|(name, config)| { - if selected_servers.contains(name) - || (selected_servers.len() == 1 - && selected_servers.first().is_some_and(|s| s.as_str() == "*")) - { - Some((name.clone(), config.clone())) - } else { - None - } - }) - .collect::>(); - - agent.mcp_servers = McpServerConfigWrapper::Map(McpServerConfig { - mcp_servers: filtered_config, - }); - } - - agent.name = agent_name.to_string(); - - Ok((agent, config_path)) + let agent = serde_json::from_slice::>(&content)?; + + let global_mcp_path = directories::chat_legacy_mcp_config(os)?; + let global_mcp_config = match McpServerConfig::load_from_file(os, global_mcp_path).await { + Ok(config) => Some(config), + Err(e) => { + tracing::error!("Error loading global mcp json path: {e}."); + None + }, + }; + Ok((agent.thaw(&config_path, global_mcp_config.as_ref())?, config_path)) }, _ => bail!("Agent {agent_name} does not exist"), } @@ -226,7 +316,7 @@ pub enum PermissionEvalResult { #[derive(Clone, Default, Debug)] pub struct Agents { - pub agents: HashMap, + pub agents: HashMap>, pub active_idx: String, pub trust_all_tools: bool, } @@ -250,15 +340,15 @@ impl Agents { } } - pub fn get_active(&self) -> Option<&Agent> { + pub fn get_active(&self) -> Option<&Agent> { self.agents.get(&self.active_idx) } - pub fn get_active_mut(&mut self) -> Option<&mut Agent> { + pub fn get_active_mut(&mut self) -> Option<&mut Agent> { self.agents.get_mut(&self.active_idx) } - pub fn switch(&mut self, name: &str) -> eyre::Result<&Agent> { + pub fn switch(&mut self, name: &str) -> eyre::Result<&Agent> { if !self.agents.contains_key(name) { eyre::bail!("No agent with name {name} found"); } @@ -296,7 +386,7 @@ impl Agents { path: Some(agent_path.clone()), ..Default::default() }; - let contents = serde_json::to_string_pretty(&agent) + let contents = serde_json::to_string_pretty(&agent.clone().freeze()) .map_err(|e| eyre::eyre!("Failed to serialize profile configuration: {}", e))?; if let Some(parent) = agent_path.parent() { @@ -365,7 +455,7 @@ impl Agents { // We could be launching from the home dir, in which case the global and local agents // are the same set of agents. If that is the case, we simply skip this. match (std::env::current_dir(), directories::home_dir(os)) { - (Ok(cwd), Ok(home_dir)) if cwd == home_dir => break 'local Vec::::new(), + (Ok(cwd), Ok(home_dir)) if cwd == home_dir => break 'local Vec::>::new(), _ => { // noop, we keep going with the extraction of local agents (even if we have an // error retrieving cwd or home_dir) @@ -373,17 +463,17 @@ impl Agents { } let Ok(path) = directories::chat_local_agent_dir() else { - break 'local Vec::::new(); + break 'local Vec::>::new(); }; let Ok(files) = os.fs.read_dir(path).await else { - break 'local Vec::::new(); + break 'local Vec::>::new(); }; load_agents_from_entries(files, os, &mut global_mcp_config).await }; let mut global_agents = 'global: { let Ok(path) = directories::chat_global_agent_path(os) else { - break 'global Vec::::new(); + break 'global Vec::>::new(); }; let files = match os.fs.read_dir(&path).await { Ok(files) => files, @@ -393,7 +483,7 @@ impl Agents { error!("Error creating global agent dir: {:?}", e); } } - break 'global Vec::::new(); + break 'global Vec::>::new(); }, }; load_agents_from_entries(files, os, &mut global_mcp_config).await @@ -439,7 +529,7 @@ impl Agents { }, ..Default::default() }; - let Ok(content) = serde_json::to_string_pretty(&example_agent) else { + let Ok(content) = serde_json::to_string_pretty(&example_agent.freeze()) else { error!("Error serializing example agent config"); break 'example_config; }; @@ -580,8 +670,9 @@ async fn load_agents_from_entries( mut files: ReadDir, os: &Os, global_mcp_config: &mut Option, -) -> Vec { - let mut res = Vec::::new(); +) -> Vec> { + let mut res = Vec::>::new(); + while let Ok(Some(file)) = files.next_entry().await { let file_path = &file.path(); if file_path @@ -598,7 +689,7 @@ async fn load_agents_from_entries( }, }; - let mut agent = match serde_json::from_slice::(&content) { + let agent = match serde_json::from_slice::>(&content) { Ok(mut agent) => { agent.path = Some(file_path.clone()); agent @@ -610,52 +701,29 @@ async fn load_agents_from_entries( }, }; - if let McpServerConfigWrapper::List(selected_servers) = &agent.mcp_servers { - let global_mcp_config = match global_mcp_config.as_mut() { - Some(config) => config, - None => { - let Ok(global_mcp_path) = directories::chat_legacy_mcp_config(os) else { - tracing::error!("Error obtaining legacy mcp json path. Skipping"); - continue; - }; - let legacy_mcp_config = match McpServerConfig::load_from_file(os, global_mcp_path).await { - Ok(config) => config, - Err(e) => { - tracing::error!("Error loading global mcp json path: {e}. Skipping"); - continue; - }, - }; - global_mcp_config.replace(legacy_mcp_config); - global_mcp_config.as_mut().unwrap() + if matches!(agent.mcp_servers, McpServerConfigWrapper::List(_)) && global_mcp_config.is_none() { + let Ok(global_mcp_path) = directories::chat_legacy_mcp_config(os) else { + tracing::error!("Error obtaining legacy mcp json path. Skipping"); + continue; + }; + let legacy_mcp_config = match McpServerConfig::load_from_file(os, global_mcp_path).await { + Ok(config) => config, + Err(e) => { + tracing::error!("Error loading global mcp json path: {e}. Skipping"); + continue; }, }; - - let filtered_config = global_mcp_config - .mcp_servers - .iter() - .filter_map(|(name, config)| { - if selected_servers.contains(name) - || (selected_servers.len() == 1 - && selected_servers.first().is_some_and(|s| s.as_str() == "*")) - { - Some((name.clone(), config.clone())) - } else { - None - } - }) - .collect::>(); - - agent.mcp_servers = McpServerConfigWrapper::Map(McpServerConfig { - mcp_servers: filtered_config, - }); + global_mcp_config.replace(legacy_mcp_config); } - if let Some(name) = Path::new(&file.file_name()).file_stem() { - agent.name = name.to_string_lossy().to_string(); - res.push(agent); - } else { - let file_path = file_path.to_string_lossy(); - tracing::error!("Unable to determine agent name from config file at {file_path}, skipping"); + match agent.thaw(file_path, global_mcp_config.as_ref()) { + Ok(agent) => res.push(agent), + Err(e) => { + tracing::error!( + "Error transforming agent at {} to usable state: {e}. Skipping", + file_path.display() + ); + }, } } } @@ -679,7 +747,7 @@ fn validate_agent_name(name: &str) -> eyre::Result<()> { Ok(()) } -async fn migrate(os: &mut Os) -> eyre::Result<(Option, Vec)> { +async fn migrate(os: &mut Os) -> eyre::Result<(Option, Vec>)> { ContextMigrate::<'a'>::scan(os) .await? .prompt_migrate() diff --git a/crates/chat-cli/src/cli/agent/wrapper_types.rs b/crates/chat-cli/src/cli/agent/wrapper_types.rs index 2c763a6e4f..a9e35db9ac 100644 --- a/crates/chat-cli/src/cli/agent/wrapper_types.rs +++ b/crates/chat-cli/src/cli/agent/wrapper_types.rs @@ -155,3 +155,14 @@ impl McpServerConfigWrapper { } } } + +/// This is a marker trait to be used in conjunction with [Cold] +pub trait SerdeReady: for<'de> Deserialize<'de> + Serialize {} + +#[derive(Deserialize, Serialize)] +pub struct Cold; + +impl SerdeReady for Cold {} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct Warm; diff --git a/crates/chat-cli/src/cli/chat/cli/profile.rs b/crates/chat-cli/src/cli/chat/cli/profile.rs index e8737db8b3..b499373650 100644 --- a/crates/chat-cli/src/cli/chat/cli/profile.rs +++ b/crates/chat-cli/src/cli/chat/cli/profile.rs @@ -22,6 +22,7 @@ use syntect::util::{ }; use crate::cli::agent::Agent; +use crate::cli::agent::wrapper_types::Cold; use crate::cli::chat::{ ChatError, ChatSession, @@ -106,7 +107,7 @@ impl AgentSubcommand { Self::Schema => { use schemars::schema_for; - let schema = schema_for!(Agent); + let schema = schema_for!(Agent); let pretty = serde_json::to_string_pretty(&schema) .map_err(|e| ChatError::Custom(format!("Failed to convert agent schema to string: {e}").into()))?; highlight_json(&mut session.stderr, pretty.as_str()) diff --git a/crates/chat-cli/src/cli/chat/cli/tools.rs b/crates/chat-cli/src/cli/chat/cli/tools.rs index 0876e79915..1fe9627f2d 100644 --- a/crates/chat-cli/src/cli/chat/cli/tools.rs +++ b/crates/chat-cli/src/cli/chat/cli/tools.rs @@ -19,6 +19,7 @@ use crossterm::{ use crate::api_client::model::Tool as FigTool; use crate::cli::agent::Agent; +use crate::cli::agent::wrapper_types::Cold; use crate::cli::chat::consts::DUMMY_TOOL_NAME; use crate::cli::chat::tools::ToolOrigin; use crate::cli::chat::{ @@ -351,8 +352,10 @@ impl ToolsSubcommand { if let Some(path) = active_agent_path { let result = async { let content = tokio::fs::read(&path).await?; - let orig_agent: Agent = serde_json::from_slice(&content)?; - Ok::>(orig_agent) + let orig_agent = serde_json::from_slice::>(&content)?; + // since all we're doing here is swapping the tool list, it's okay if we + // don't thaw it here + Ok::, Box>(orig_agent) } .await; diff --git a/crates/chat-cli/src/cli/chat/context.rs b/crates/chat-cli/src/cli/chat/context.rs index 447a56f8ad..7a9b6205c5 100644 --- a/crates/chat-cli/src/cli/chat/context.rs +++ b/crates/chat-cli/src/cli/chat/context.rs @@ -14,6 +14,7 @@ use serde::{ use super::consts::CONTEXT_FILES_MAX_SIZE; use super::util::drop_matched_context_files; +use crate::cli::agent::wrapper_types::Warm; use crate::cli::agent::{ Agent, CreateHooks, @@ -38,10 +39,10 @@ pub struct ContextConfig { pub hooks: HashMap, } -impl TryFrom<&Agent> for ContextConfig { +impl TryFrom<&Agent> for ContextConfig { type Error = eyre::Report; - fn try_from(value: &Agent) -> Result { + fn try_from(value: &Agent) -> Result { Ok(Self { paths: value.included_files.clone(), hooks: { @@ -100,7 +101,7 @@ pub struct ContextManager { } impl ContextManager { - pub fn from_agent(agent: &Agent, max_context_files_size: Option) -> Result { + pub fn from_agent(agent: &Agent, max_context_files_size: Option) -> Result { let max_context_files_size = max_context_files_size.unwrap_or(CONTEXT_FILES_MAX_SIZE); let current_profile = agent.name.clone(); @@ -272,7 +273,11 @@ impl ContextManager { /// Run all the currently enabled hooks from both the global and profile contexts. /// # Returns /// A vector containing pairs of a [`Hook`] definition and its execution output - pub async fn run_hooks(&mut self, output: &mut impl Write, prompt: Option<&str>) -> Result, ChatError> { + pub async fn run_hooks( + &mut self, + output: &mut impl Write, + prompt: Option<&str>, + ) -> Result, ChatError> { let hooks = self .profile_config .hooks diff --git a/crates/chat-cli/src/cli/chat/tool_manager.rs b/crates/chat-cli/src/cli/chat/tool_manager.rs index 803ca5bea6..b9e5a49e76 100644 --- a/crates/chat-cli/src/cli/chat/tool_manager.rs +++ b/crates/chat-cli/src/cli/chat/tool_manager.rs @@ -58,6 +58,7 @@ use crate::api_client::model::{ ToolResultContentBlock, ToolResultStatus, }; +use crate::cli::agent::wrapper_types::Warm; use crate::cli::agent::{ Agent, McpServerConfig, @@ -154,7 +155,7 @@ pub struct ToolManagerBuilder { prompt_list_sender: Option>>, prompt_list_receiver: Option>>, conversation_id: Option, - agent: Option, + agent: Option>, } impl ToolManagerBuilder { @@ -173,7 +174,7 @@ impl ToolManagerBuilder { self } - pub fn agent(mut self, agent: Agent) -> Self { + pub fn agent(mut self, agent: Agent) -> Self { if let McpServerConfigWrapper::Map(config) = &agent.mcp_servers { self.mcp_server_config.replace(config.clone()); } else { @@ -192,7 +193,8 @@ impl ToolManagerBuilder { mut output: Box, interactive: bool, ) -> eyre::Result { - let McpServerConfig { mcp_servers } = self.mcp_server_config.ok_or(eyre::eyre!("Missing mcp server config"))?; + let McpServerConfig { mcp_servers, .. } = + self.mcp_server_config.ok_or(eyre::eyre!("Missing mcp server config"))?; debug_assert!(self.conversation_id.is_some()); let conversation_id = self.conversation_id.ok_or(eyre::eyre!("Missing conversation id"))?; @@ -864,7 +866,7 @@ pub struct ToolManager { /// A collection of preferences that pertains to the conversation. /// As far as tool manager goes, this is relevant for tool and server filters - pub agent: Arc>, + pub agent: Arc>>, } impl Clone for ToolManager { diff --git a/crates/chat-cli/src/cli/chat/tools/custom_tool.rs b/crates/chat-cli/src/cli/chat/tools/custom_tool.rs index cf25b32365..eae2d439d7 100644 --- a/crates/chat-cli/src/cli/chat/tools/custom_tool.rs +++ b/crates/chat-cli/src/cli/chat/tools/custom_tool.rs @@ -17,6 +17,7 @@ use tokio::sync::RwLock; use tracing::warn; use super::InvokeOutput; +use crate::cli::agent::wrapper_types::Warm; use crate::cli::agent::{ Agent, PermissionEvalResult, @@ -255,7 +256,7 @@ impl CustomTool { + TokenCounter::count_tokens(self.params.as_ref().map_or("", |p| p.as_str().unwrap_or_default())) } - pub fn eval_perm(&self, agent: &Agent) -> PermissionEvalResult { + pub fn eval_perm(&self, agent: &Agent) -> PermissionEvalResult { use crate::util::MCP_SERVER_TOOL_DELIMITER; let Self { name: tool_name, diff --git a/crates/chat-cli/src/cli/chat/tools/execute/mod.rs b/crates/chat-cli/src/cli/chat/tools/execute/mod.rs index 9b7cd6e76b..8d073e31cd 100644 --- a/crates/chat-cli/src/cli/chat/tools/execute/mod.rs +++ b/crates/chat-cli/src/cli/chat/tools/execute/mod.rs @@ -9,6 +9,7 @@ use eyre::Result; use serde::Deserialize; use tracing::error; +use crate::cli::agent::wrapper_types::Warm; use crate::cli::agent::{ Agent, PermissionEvalResult, @@ -154,7 +155,7 @@ impl ExecuteCommand { Ok(()) } - pub fn eval_perm(&self, agent: &Agent) -> PermissionEvalResult { + pub fn eval_perm(&self, agent: &Agent) -> PermissionEvalResult { #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] struct Settings { diff --git a/crates/chat-cli/src/cli/chat/tools/fs_read.rs b/crates/chat-cli/src/cli/chat/tools/fs_read.rs index 6fa2ca97e4..90845940ee 100644 --- a/crates/chat-cli/src/cli/chat/tools/fs_read.rs +++ b/crates/chat-cli/src/cli/chat/tools/fs_read.rs @@ -33,6 +33,7 @@ use super::{ format_path, sanitize_path_tool_arg, }; +use crate::cli::agent::wrapper_types::Warm; use crate::cli::agent::{ Agent, PermissionEvalResult, @@ -99,7 +100,7 @@ impl FsRead { } } - pub fn eval_perm(&self, agent: &Agent) -> PermissionEvalResult { + pub fn eval_perm(&self, agent: &Agent) -> PermissionEvalResult { #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] struct Settings { diff --git a/crates/chat-cli/src/cli/chat/tools/fs_write.rs b/crates/chat-cli/src/cli/chat/tools/fs_write.rs index 3f7c5d41c4..9ce28cde1b 100644 --- a/crates/chat-cli/src/cli/chat/tools/fs_write.rs +++ b/crates/chat-cli/src/cli/chat/tools/fs_write.rs @@ -37,6 +37,7 @@ use super::{ sanitize_path_tool_arg, supports_truecolor, }; +use crate::cli::agent::wrapper_types::Warm; use crate::cli::agent::{ Agent, PermissionEvalResult, @@ -346,7 +347,7 @@ impl FsWrite { } } - pub fn eval_perm(&self, agent: &Agent) -> PermissionEvalResult { + pub fn eval_perm(&self, agent: &Agent) -> PermissionEvalResult { #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] struct Settings { diff --git a/crates/chat-cli/src/cli/chat/tools/mod.rs b/crates/chat-cli/src/cli/chat/tools/mod.rs index 4b0a121dae..34e7c31910 100644 --- a/crates/chat-cli/src/cli/chat/tools/mod.rs +++ b/crates/chat-cli/src/cli/chat/tools/mod.rs @@ -35,6 +35,7 @@ use use_aws::UseAws; use super::consts::MAX_TOOL_RESPONSE_SIZE; use super::util::images::RichImageBlocks; +use crate::cli::agent::wrapper_types::Warm; use crate::cli::agent::{ Agent, PermissionEvalResult, @@ -89,7 +90,7 @@ impl Tool { } /// Whether or not the tool should prompt the user to accept before [Self::invoke] is called. - pub fn requires_acceptance(&self, agent: &Agent) -> PermissionEvalResult { + pub fn requires_acceptance(&self, agent: &Agent) -> PermissionEvalResult { match self { Tool::FsRead(fs_read) => fs_read.eval_perm(agent), Tool::FsWrite(fs_write) => fs_write.eval_perm(agent), diff --git a/crates/chat-cli/src/cli/chat/tools/use_aws.rs b/crates/chat-cli/src/cli/chat/tools/use_aws.rs index 59b41e8b0d..1a017782e8 100644 --- a/crates/chat-cli/src/cli/chat/tools/use_aws.rs +++ b/crates/chat-cli/src/cli/chat/tools/use_aws.rs @@ -23,6 +23,7 @@ use super::{ MAX_TOOL_RESPONSE_SIZE, OutputKind, }; +use crate::cli::agent::wrapper_types::Warm; use crate::cli::agent::{ Agent, PermissionEvalResult, @@ -194,7 +195,7 @@ impl UseAws { } } - pub fn eval_perm(&self, agent: &Agent) -> PermissionEvalResult { + pub fn eval_perm(&self, agent: &Agent) -> PermissionEvalResult { #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] struct Settings { diff --git a/crates/chat-cli/src/cli/mcp.rs b/crates/chat-cli/src/cli/mcp.rs index 142244d66a..56e3215eff 100644 --- a/crates/chat-cli/src/cli/mcp.rs +++ b/crates/chat-cli/src/cli/mcp.rs @@ -138,7 +138,7 @@ impl AddArgs { }))?; mcp_servers.insert(self.name.clone(), tool); - let json = serde_json::to_string_pretty(&agent)?; + let json = serde_json::to_string_pretty(&agent.clone().freeze())?; os.fs.write(config_path, json).await?; writeln!(output, "✓ Added MCP server '{}' to agent {}\n", self.name, agent_name)?; }, @@ -206,7 +206,7 @@ impl RemoveArgs { match config.remove(&self.name) { Some(_) => { - let json = serde_json::to_string_pretty(&agent)?; + let json = serde_json::to_string_pretty(&agent.clone().freeze())?; os.fs.write(config_path, json).await?; writeln!( output, From 58c5c52e8baea6b3aac074a2444260bf560057eb Mon Sep 17 00:00:00 2001 From: Felix dingfeli Date: Tue, 15 Jul 2025 17:54:15 -0700 Subject: [PATCH 07/15] allows optional @builtin prefix --- crates/chat-cli/src/cli/agent/mod.rs | 28 ++++++++----------- .../src/cli/agent/root_command_args.rs | 5 ++-- crates/chat-cli/src/cli/chat/cli/profile.rs | 2 +- crates/chat-cli/src/cli/chat/mod.rs | 3 +- crates/chat-cli/src/cli/chat/tool_manager.rs | 13 ++++++--- 5 files changed, 26 insertions(+), 25 deletions(-) diff --git a/crates/chat-cli/src/cli/agent/mod.rs b/crates/chat-cli/src/cli/agent/mod.rs index 6d910907ca..b682c64e85 100644 --- a/crates/chat-cli/src/cli/agent/mod.rs +++ b/crates/chat-cli/src/cli/agent/mod.rs @@ -10,7 +10,10 @@ use std::io::{ Write, }; use std::marker::PhantomData; -use std::path::PathBuf; +use std::path::{ + Path, + PathBuf, +}; use context_migrate::ContextMigrate; use crossterm::style::{ @@ -35,17 +38,15 @@ use tracing::{ error, warn, }; -use wrapper_types::{ - Cold, - Warm, -}; pub use wrapper_types::{ + Cold, CreateHooks, McpServerConfigWrapper, OriginalToolName, PromptHooks, SerdeReady, ToolSettingTarget, + Warm, alias_schema, tool_settings_schema, }; @@ -65,7 +66,7 @@ use crate::util::{ mod context_migrate; mod mcp_config; mod root_command_args; -mod wrapper_types; +pub mod wrapper_types; pub use root_command_args::*; @@ -139,7 +140,7 @@ pub struct Agent { #[serde(skip)] pub path: Option, #[serde(skip)] - phantom: std::marker::PhantomData, + pub phantom: std::marker::PhantomData, } impl Default for Agent { @@ -171,7 +172,7 @@ impl Default for Agent { } impl Agent { - pub fn thaw(mut self, path: &PathBuf, global_mcp_config: Option<&McpServerConfig>) -> eyre::Result> { + pub fn thaw(mut self, path: &Path, global_mcp_config: Option<&McpServerConfig>) -> eyre::Result> { let Self { mcp_servers, .. } = &mut self; let name = path @@ -223,7 +224,7 @@ impl Agent { create_hooks: self.create_hooks, prompt_hooks: self.prompt_hooks, tools_settings: self.tools_settings, - path: Some(path.clone()), + path: Some(path.to_path_buf()), phantom: PhantomData, }) } @@ -815,7 +816,7 @@ mod tests { #[test] fn test_deser() { - let agent = serde_json::from_str::(INPUT).expect("Deserializtion failed"); + let agent = serde_json::from_str::>(INPUT).expect("Deserializtion failed"); assert!(matches!(agent.mcp_servers, McpServerConfigWrapper::Map(_))); if let McpServerConfigWrapper::Map(config) = &agent.mcp_servers { assert!(config.mcp_servers.contains_key("fetch")); @@ -1002,13 +1003,6 @@ mod tests { assert!(validate_agent_name("invalid!").is_err()); assert!(validate_agent_name("invalid space").is_err()); } - - #[test] - fn test_schema_gen() { - use schemars::schema_for; - let schema = schema_for!(Agent); - println!("Schema for agent: {}", serde_json::to_string_pretty(&schema).unwrap()); - } } #[cfg(test)] diff --git a/crates/chat-cli/src/cli/agent/root_command_args.rs b/crates/chat-cli/src/cli/agent/root_command_args.rs index 86dfadf30c..034261a9cd 100644 --- a/crates/chat-cli/src/cli/agent/root_command_args.rs +++ b/crates/chat-cli/src/cli/agent/root_command_args.rs @@ -14,6 +14,7 @@ use eyre::{ use super::{ Agent, Agents, + Cold, }; use crate::database::settings::Setting; use crate::os::Os; @@ -103,7 +104,7 @@ impl AgentArgs { path_with_file_name.display() ); }; - if let Err(e) = serde_json::from_slice::(&content) { + if let Err(e) = serde_json::from_slice::>(&content) { bail!( "Post write validation failed for agent '{name}' at path: {}. Malformed config detected: {e}", path_with_file_name.display() @@ -163,7 +164,7 @@ pub async fn create_agent( let prepopulated_content = if let Some(from) = from { let agent_to_copy = agents.switch(from.as_str())?; - serde_json::to_string_pretty(agent_to_copy)? + serde_json::to_string_pretty(&agent_to_copy.clone().freeze())? } else { Default::default() }; diff --git a/crates/chat-cli/src/cli/chat/cli/profile.rs b/crates/chat-cli/src/cli/chat/cli/profile.rs index 54c29dd17d..cfea635ffa 100644 --- a/crates/chat-cli/src/cli/chat/cli/profile.rs +++ b/crates/chat-cli/src/cli/chat/cli/profile.rs @@ -155,7 +155,7 @@ impl AgentSubcommand { .into(), )); }; - if let Err(e) = serde_json::from_slice::(&content) { + if let Err(e) = serde_json::from_slice::>(&content) { return Err(ChatError::Custom( format!("Post write validation failed for agent '{name}'. Malformed config detected: {e}") .into(), diff --git a/crates/chat-cli/src/cli/chat/mod.rs b/crates/chat-cli/src/cli/chat/mod.rs index c58159325a..0bce2778bd 100644 --- a/crates/chat-cli/src/cli/chat/mod.rs +++ b/crates/chat-cli/src/cli/chat/mod.rs @@ -2476,7 +2476,8 @@ mod tests { ..Default::default() }; if let Ok(false) = os.fs.try_exists(AGENT_PATH).await { - let content = serde_json::to_string_pretty(&agent).expect("Failed to serialize test agent to file"); + let content = + serde_json::to_string_pretty(&agent.clone().freeze()).expect("Failed to serialize test agent to file"); let agent_path = PathBuf::from(AGENT_PATH); os.fs .create_dir_all( diff --git a/crates/chat-cli/src/cli/chat/tool_manager.rs b/crates/chat-cli/src/cli/chat/tool_manager.rs index b9e5a49e76..f6999a7846 100644 --- a/crates/chat-cli/src/cli/chat/tool_manager.rs +++ b/crates/chat-cli/src/cli/chat/tool_manager.rs @@ -227,7 +227,7 @@ impl ToolManagerBuilder { style::Print("\n") ); None - } else if server_name == "native" { + } else if server_name == "builtin" { let _ = queue!( output, style::SetForegroundColor(style::Color::Red), @@ -237,7 +237,7 @@ impl ToolManagerBuilder { style::ResetColor, style::Print(". Server name cannot contain reserved word "), style::SetForegroundColor(style::Color::Yellow), - style::Print("native"), + style::Print("builtin"), style::ResetColor, style::Print(" (it is used to denote native tools)\n") ); @@ -898,11 +898,16 @@ impl ToolManager { self.schema = { let tool_list = &self.agent.lock().await.tools; let is_allow_all = tool_list.len() == 1 && tool_list.first().is_some_and(|n| n == "*"); - let is_allow_native = tool_list.iter().any(|t| t.as_str() == "@native"); + let is_allow_native = tool_list.iter().any(|t| t.as_str() == "@builtin"); let mut tool_specs = serde_json::from_str::>(include_str!("tools/tool_index.json"))? .into_iter() - .filter(|(name, _)| is_allow_all || is_allow_native || tool_list.contains(name)) + .filter(|(name, _)| { + is_allow_all + || is_allow_native + || tool_list.contains(name) + || tool_list.contains(&format!("@builtin/{name}")) + }) .collect::>(); if !crate::cli::chat::tools::thinking::Thinking::is_enabled(os) { tool_specs.remove("thinking"); From a1e164b06133daf8c13be2dfe4a8d2fe0897d060 Mon Sep 17 00:00:00 2001 From: Felix dingfeli Date: Wed, 16 Jul 2025 15:03:21 -0700 Subject: [PATCH 08/15] reverts typestate pattern --- .../chat-cli/src/cli/agent/context_migrate.rs | 7 +- crates/chat-cli/src/cli/agent/mod.rs | 199 ++++++------------ .../src/cli/agent/root_command_args.rs | 5 +- .../chat-cli/src/cli/agent/wrapper_types.rs | 11 - crates/chat-cli/src/cli/chat/cli/profile.rs | 5 +- crates/chat-cli/src/cli/chat/cli/tools.rs | 5 +- crates/chat-cli/src/cli/chat/context.rs | 7 +- crates/chat-cli/src/cli/chat/mod.rs | 3 +- crates/chat-cli/src/cli/chat/tool_manager.rs | 10 +- .../src/cli/chat/tools/custom_tool.rs | 3 +- .../src/cli/chat/tools/execute/mod.rs | 3 +- crates/chat-cli/src/cli/chat/tools/fs_read.rs | 3 +- .../chat-cli/src/cli/chat/tools/fs_write.rs | 3 +- crates/chat-cli/src/cli/chat/tools/mod.rs | 3 +- crates/chat-cli/src/cli/chat/tools/use_aws.rs | 3 +- crates/chat-cli/src/cli/mcp.rs | 4 +- 16 files changed, 93 insertions(+), 181 deletions(-) diff --git a/crates/chat-cli/src/cli/agent/context_migrate.rs b/crates/chat-cli/src/cli/agent/context_migrate.rs index e8cfd5c523..0711dc860b 100644 --- a/crates/chat-cli/src/cli/agent/context_migrate.rs +++ b/crates/chat-cli/src/cli/agent/context_migrate.rs @@ -8,7 +8,6 @@ use tracing::{ warn, }; -use super::wrapper_types::Warm; use super::{ Agent, McpServerConfig, @@ -31,7 +30,7 @@ pub(in crate::cli::agent) struct ContextMigrate { legacy_global_context: Option, legacy_profiles: HashMap, mcp_servers: Option, - new_agents: Vec>, + new_agents: Vec, } impl ContextMigrate<'a'> { @@ -214,7 +213,7 @@ impl ContextMigrate<'c'> { } for agent in &mut new_agents { - let content = serde_json::to_string_pretty(&agent.clone().freeze())?; + let content = agent.to_str_pretty()?; if let Some(path) = agent.path.as_ref() { info!("Agent {} peristed in path {}", agent.name, path.to_string_lossy()); os.fs.write(path, content).await?; @@ -259,7 +258,7 @@ impl ContextMigrate<'c'> { } impl ContextMigrate<'d'> { - pub async fn prompt_set_default(self, os: &mut Os) -> eyre::Result<(Option, Vec>)> { + pub async fn prompt_set_default(self, os: &mut Os) -> eyre::Result<(Option, Vec)> { let ContextMigrate { new_agents, .. } = self; let labels = new_agents diff --git a/crates/chat-cli/src/cli/agent/mod.rs b/crates/chat-cli/src/cli/agent/mod.rs index b682c64e85..51a4922ed2 100644 --- a/crates/chat-cli/src/cli/agent/mod.rs +++ b/crates/chat-cli/src/cli/agent/mod.rs @@ -9,7 +9,6 @@ use std::io::{ self, Write, }; -use std::marker::PhantomData; use std::path::{ Path, PathBuf, @@ -39,14 +38,11 @@ use tracing::{ warn, }; pub use wrapper_types::{ - Cold, CreateHooks, McpServerConfigWrapper, OriginalToolName, PromptHooks, - SerdeReady, ToolSettingTarget, - Warm, alias_schema, tool_settings_schema, }; @@ -95,9 +91,9 @@ pub use root_command_args::*; /// Where agents are instantiated from their config, we would need to convert them from "cold" to /// "warm". #[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, JsonSchema)] -#[serde(rename_all = "camelCase", bound = "T: SerdeReady")] +#[serde(rename_all = "camelCase")] #[schemars(description = "An Agent is a declarative way of configuring a given instance of q chat.")] -pub struct Agent { +pub struct Agent { /// Agent names are derived from the file name. Thus they are skipped for /// serializing #[serde(skip)] @@ -139,11 +135,9 @@ pub struct Agent { pub tools_settings: HashMap, #[serde(skip)] pub path: Option, - #[serde(skip)] - pub phantom: std::marker::PhantomData, } -impl Default for Agent { +impl Default for Agent { fn default() -> Self { Self { name: "default".to_string(), @@ -166,14 +160,41 @@ impl Default for Agent { prompt_hooks: Default::default(), tools_settings: Default::default(), path: None, - phantom: PhantomData, } } } -impl Agent { - pub fn thaw(mut self, path: &Path, global_mcp_config: Option<&McpServerConfig>) -> eyre::Result> { - let Self { mcp_servers, .. } = &mut self; +impl Agent { + pub fn freeze(&mut self) -> eyre::Result<()> { + let Self { mcp_servers, .. } = self; + + if let McpServerConfigWrapper::Map(servers) = mcp_servers { + match servers.original_config_form { + OriginalForm::All => { + *mcp_servers = McpServerConfigWrapper::List(vec!["*".to_string()]); + }, + OriginalForm::List => { + *mcp_servers = McpServerConfigWrapper::List({ + servers + .mcp_servers + .iter() + .fold(Vec::::new(), |mut acc, (name, _)| { + acc.push(name.clone()); + acc + }) + }); + }, + OriginalForm::Map => { + // noop + }, + } + } + + Ok(()) + } + + pub fn thaw(&mut self, path: &Path, global_mcp_config: Option<&McpServerConfig>) -> eyre::Result<()> { + let Self { mcp_servers, .. } = self; let name = path .file_stem() @@ -212,70 +233,18 @@ impl Agent { }); } - Ok(Agent { - name, - description: self.description, - prompt: self.prompt, - mcp_servers: self.mcp_servers, - tools: self.tools, - alias: self.alias, - allowed_tools: self.allowed_tools, - included_files: self.included_files, - create_hooks: self.create_hooks, - prompt_hooks: self.prompt_hooks, - tools_settings: self.tools_settings, - path: Some(path.to_path_buf()), - phantom: PhantomData, - }) + Ok(()) } -} -impl Agent { - pub fn freeze(mut self) -> Agent { - let Self { mcp_servers, .. } = &mut self; - - if let McpServerConfigWrapper::Map(servers) = mcp_servers { - match servers.original_config_form { - OriginalForm::All => { - *mcp_servers = McpServerConfigWrapper::List(vec!["*".to_string()]); - }, - OriginalForm::List => { - *mcp_servers = McpServerConfigWrapper::List({ - servers - .mcp_servers - .iter() - .fold(Vec::::new(), |mut acc, (name, _)| { - acc.push(name.clone()); - acc - }) - }); - }, - OriginalForm::Map => { - // noop - }, - } - } - - Agent { - name: self.name, - description: self.description, - prompt: self.prompt, - mcp_servers: self.mcp_servers, - tools: self.tools, - alias: self.alias, - allowed_tools: self.allowed_tools, - included_files: self.included_files, - create_hooks: self.create_hooks, - prompt_hooks: self.prompt_hooks, - tools_settings: self.tools_settings, - path: self.path, - phantom: PhantomData, - } + pub fn to_str_pretty(&self) -> eyre::Result { + let mut agent_clone = self.clone(); + agent_clone.freeze()?; + Ok(serde_json::to_string_pretty(&agent_clone)?) } /// Retrieves an agent by name. It does so via first seeking the given agent under local dir, /// and falling back to global dir if it does not exist in local. - pub async fn get_agent_by_name(os: &Os, agent_name: &str) -> eyre::Result<(Agent, PathBuf)> { + pub async fn get_agent_by_name(os: &Os, agent_name: &str) -> eyre::Result<(Agent, PathBuf)> { let config_path: Result = 'config: { // local first, and then fall back to looking at global let local_config_dir = directories::chat_local_agent_dir()?.join(format!("{agent_name}.json")); @@ -294,7 +263,7 @@ impl Agent { match config_path { Ok(config_path) => { let content = os.fs.read(&config_path).await?; - let agent = serde_json::from_slice::>(&content)?; + let mut agent = serde_json::from_slice::(&content)?; let global_mcp_path = directories::chat_legacy_mcp_config(os)?; let global_mcp_config = match McpServerConfig::load_from_file(os, global_mcp_path).await { @@ -304,7 +273,8 @@ impl Agent { None }, }; - Ok((agent.thaw(&config_path, global_mcp_config.as_ref())?, config_path)) + agent.thaw(&config_path, global_mcp_config.as_ref())?; + Ok((agent, config_path)) }, _ => bail!("Agent {agent_name} does not exist"), } @@ -320,7 +290,7 @@ pub enum PermissionEvalResult { #[derive(Clone, Default, Debug)] pub struct Agents { - pub agents: HashMap>, + pub agents: HashMap, pub active_idx: String, pub trust_all_tools: bool, } @@ -344,15 +314,15 @@ impl Agents { } } - pub fn get_active(&self) -> Option<&Agent> { + pub fn get_active(&self) -> Option<&Agent> { self.agents.get(&self.active_idx) } - pub fn get_active_mut(&mut self) -> Option<&mut Agent> { + pub fn get_active_mut(&mut self) -> Option<&mut Agent> { self.agents.get_mut(&self.active_idx) } - pub fn switch(&mut self, name: &str) -> eyre::Result<&Agent> { + pub fn switch(&mut self, name: &str) -> eyre::Result<&Agent> { if !self.agents.contains_key(name) { eyre::bail!("No agent with name {name} found"); } @@ -390,7 +360,8 @@ impl Agents { path: Some(agent_path.clone()), ..Default::default() }; - let contents = serde_json::to_string_pretty(&agent.clone().freeze()) + let contents = agent + .to_str_pretty() .map_err(|e| eyre::eyre!("Failed to serialize profile configuration: {}", e))?; if let Some(parent) = agent_path.parent() { @@ -459,7 +430,7 @@ impl Agents { // We could be launching from the home dir, in which case the global and local agents // are the same set of agents. If that is the case, we simply skip this. match (std::env::current_dir(), directories::home_dir(os)) { - (Ok(cwd), Ok(home_dir)) if cwd == home_dir => break 'local Vec::>::new(), + (Ok(cwd), Ok(home_dir)) if cwd == home_dir => break 'local Vec::::new(), _ => { // noop, we keep going with the extraction of local agents (even if we have an // error retrieving cwd or home_dir) @@ -467,17 +438,17 @@ impl Agents { } let Ok(path) = directories::chat_local_agent_dir() else { - break 'local Vec::>::new(); + break 'local Vec::::new(); }; let Ok(files) = os.fs.read_dir(path).await else { - break 'local Vec::>::new(); + break 'local Vec::::new(); }; load_agents_from_entries(files, os, &mut global_mcp_config).await }; let mut global_agents = 'global: { let Ok(path) = directories::chat_global_agent_path(os) else { - break 'global Vec::>::new(); + break 'global Vec::::new(); }; let files = match os.fs.read_dir(&path).await { Ok(files) => files, @@ -487,7 +458,7 @@ impl Agents { error!("Error creating global agent dir: {:?}", e); } } - break 'global Vec::>::new(); + break 'global Vec::::new(); }, }; load_agents_from_entries(files, os, &mut global_mcp_config).await @@ -533,7 +504,7 @@ impl Agents { }, ..Default::default() }; - let Ok(content) = serde_json::to_string_pretty(&example_agent.freeze()) else { + let Ok(content) = example_agent.to_str_pretty() else { error!("Error serializing example agent config"); break 'example_config; }; @@ -674,8 +645,8 @@ async fn load_agents_from_entries( mut files: ReadDir, os: &Os, global_mcp_config: &mut Option, -) -> Vec> { - let mut res = Vec::>::new(); +) -> Vec { + let mut res = Vec::::new(); while let Ok(Some(file)) = files.next_entry().await { let file_path = &file.path(); @@ -693,7 +664,7 @@ async fn load_agents_from_entries( }, }; - let agent = match serde_json::from_slice::>(&content) { + let mut agent = match serde_json::from_slice::(&content) { Ok(mut agent) => { agent.path = Some(file_path.clone()); agent @@ -720,15 +691,14 @@ async fn load_agents_from_entries( global_mcp_config.replace(legacy_mcp_config); } - match agent.thaw(file_path, global_mcp_config.as_ref()) { - Ok(agent) => res.push(agent), - Err(e) => { - tracing::error!( - "Error transforming agent at {} to usable state: {e}. Skipping", - file_path.display() - ); - }, - } + if let Err(e) = agent.thaw(file_path, global_mcp_config.as_ref()) { + tracing::error!( + "Error transforming agent at {} to usable state: {e}. Skipping", + file_path.display() + ); + }; + + res.push(agent); } } res @@ -751,7 +721,7 @@ fn validate_agent_name(name: &str) -> eyre::Result<()> { Ok(()) } -async fn migrate(os: &mut Os) -> eyre::Result<(Option, Vec>)> { +async fn migrate(os: &mut Os) -> eyre::Result<(Option, Vec)> { ContextMigrate::<'a'>::scan(os) .await? .prompt_migrate() @@ -816,7 +786,7 @@ mod tests { #[test] fn test_deser() { - let agent = serde_json::from_str::>(INPUT).expect("Deserializtion failed"); + let agent = serde_json::from_str::(INPUT).expect("Deserializtion failed"); assert!(matches!(agent.mcp_servers, McpServerConfigWrapper::Map(_))); if let McpServerConfigWrapper::Map(config) = &agent.mcp_servers { assert!(config.mcp_servers.contains_key("fetch")); @@ -1004,36 +974,3 @@ mod tests { assert!(validate_agent_name("invalid space").is_err()); } } - -#[cfg(test)] -mod test_test { - use super::*; - trait A: for<'de> Deserialize<'de> + Serialize {} - #[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, JsonSchema)] - #[serde(rename_all = "camelCase", bound = "T: A")] - struct Testtest { - hello: String, - #[serde(skip)] - phantom: std::marker::PhantomData, - } - - #[derive(Serialize, Deserialize)] - struct ConcreteA; - - impl A for ConcreteA {} - - #[derive(Serialize, Deserialize)] - struct ConcreteB; - - const INPUT: &str = r#" - { - "hello": "world" - } - "#; - - #[test] - fn test_deser_err() { - let deser_res = serde_json::from_str::>(INPUT); - assert!(deser_res.is_ok()); - } -} diff --git a/crates/chat-cli/src/cli/agent/root_command_args.rs b/crates/chat-cli/src/cli/agent/root_command_args.rs index 034261a9cd..7e1262b369 100644 --- a/crates/chat-cli/src/cli/agent/root_command_args.rs +++ b/crates/chat-cli/src/cli/agent/root_command_args.rs @@ -14,7 +14,6 @@ use eyre::{ use super::{ Agent, Agents, - Cold, }; use crate::database::settings::Setting; use crate::os::Os; @@ -104,7 +103,7 @@ impl AgentArgs { path_with_file_name.display() ); }; - if let Err(e) = serde_json::from_slice::>(&content) { + if let Err(e) = serde_json::from_slice::(&content) { bail!( "Post write validation failed for agent '{name}' at path: {}. Malformed config detected: {e}", path_with_file_name.display() @@ -164,7 +163,7 @@ pub async fn create_agent( let prepopulated_content = if let Some(from) = from { let agent_to_copy = agents.switch(from.as_str())?; - serde_json::to_string_pretty(&agent_to_copy.clone().freeze())? + agent_to_copy.to_str_pretty()? } else { Default::default() }; diff --git a/crates/chat-cli/src/cli/agent/wrapper_types.rs b/crates/chat-cli/src/cli/agent/wrapper_types.rs index a9e35db9ac..2c763a6e4f 100644 --- a/crates/chat-cli/src/cli/agent/wrapper_types.rs +++ b/crates/chat-cli/src/cli/agent/wrapper_types.rs @@ -155,14 +155,3 @@ impl McpServerConfigWrapper { } } } - -/// This is a marker trait to be used in conjunction with [Cold] -pub trait SerdeReady: for<'de> Deserialize<'de> + Serialize {} - -#[derive(Deserialize, Serialize)] -pub struct Cold; - -impl SerdeReady for Cold {} - -#[derive(Clone, Debug, Eq, PartialEq)] -pub struct Warm; diff --git a/crates/chat-cli/src/cli/chat/cli/profile.rs b/crates/chat-cli/src/cli/chat/cli/profile.rs index cfea635ffa..8e60732fbd 100644 --- a/crates/chat-cli/src/cli/chat/cli/profile.rs +++ b/crates/chat-cli/src/cli/chat/cli/profile.rs @@ -25,7 +25,6 @@ use syntect::util::{ use crate::cli::agent::{ Agent, Agents, - Cold, create_agent, rename_agent, }; @@ -127,7 +126,7 @@ impl AgentSubcommand { Self::Schema => { use schemars::schema_for; - let schema = schema_for!(Agent); + let schema = schema_for!(Agent); let pretty = serde_json::to_string_pretty(&schema) .map_err(|e| ChatError::Custom(format!("Failed to convert agent schema to string: {e}").into()))?; highlight_json(&mut session.stderr, pretty.as_str()) @@ -155,7 +154,7 @@ impl AgentSubcommand { .into(), )); }; - if let Err(e) = serde_json::from_slice::>(&content) { + if let Err(e) = serde_json::from_slice::(&content) { return Err(ChatError::Custom( format!("Post write validation failed for agent '{name}'. Malformed config detected: {e}") .into(), diff --git a/crates/chat-cli/src/cli/chat/cli/tools.rs b/crates/chat-cli/src/cli/chat/cli/tools.rs index 1fe9627f2d..515983919d 100644 --- a/crates/chat-cli/src/cli/chat/cli/tools.rs +++ b/crates/chat-cli/src/cli/chat/cli/tools.rs @@ -19,7 +19,6 @@ use crossterm::{ use crate::api_client::model::Tool as FigTool; use crate::cli::agent::Agent; -use crate::cli::agent::wrapper_types::Cold; use crate::cli::chat::consts::DUMMY_TOOL_NAME; use crate::cli::chat::tools::ToolOrigin; use crate::cli::chat::{ @@ -352,10 +351,10 @@ impl ToolsSubcommand { if let Some(path) = active_agent_path { let result = async { let content = tokio::fs::read(&path).await?; - let orig_agent = serde_json::from_slice::>(&content)?; + let orig_agent = serde_json::from_slice::(&content)?; // since all we're doing here is swapping the tool list, it's okay if we // don't thaw it here - Ok::, Box>(orig_agent) + Ok::>(orig_agent) } .await; diff --git a/crates/chat-cli/src/cli/chat/context.rs b/crates/chat-cli/src/cli/chat/context.rs index 7a9b6205c5..c71ff148fb 100644 --- a/crates/chat-cli/src/cli/chat/context.rs +++ b/crates/chat-cli/src/cli/chat/context.rs @@ -14,7 +14,6 @@ use serde::{ use super::consts::CONTEXT_FILES_MAX_SIZE; use super::util::drop_matched_context_files; -use crate::cli::agent::wrapper_types::Warm; use crate::cli::agent::{ Agent, CreateHooks, @@ -39,10 +38,10 @@ pub struct ContextConfig { pub hooks: HashMap, } -impl TryFrom<&Agent> for ContextConfig { +impl TryFrom<&Agent> for ContextConfig { type Error = eyre::Report; - fn try_from(value: &Agent) -> Result { + fn try_from(value: &Agent) -> Result { Ok(Self { paths: value.included_files.clone(), hooks: { @@ -101,7 +100,7 @@ pub struct ContextManager { } impl ContextManager { - pub fn from_agent(agent: &Agent, max_context_files_size: Option) -> Result { + pub fn from_agent(agent: &Agent, max_context_files_size: Option) -> Result { let max_context_files_size = max_context_files_size.unwrap_or(CONTEXT_FILES_MAX_SIZE); let current_profile = agent.name.clone(); diff --git a/crates/chat-cli/src/cli/chat/mod.rs b/crates/chat-cli/src/cli/chat/mod.rs index 0bce2778bd..127216f4ef 100644 --- a/crates/chat-cli/src/cli/chat/mod.rs +++ b/crates/chat-cli/src/cli/chat/mod.rs @@ -2476,8 +2476,7 @@ mod tests { ..Default::default() }; if let Ok(false) = os.fs.try_exists(AGENT_PATH).await { - let content = - serde_json::to_string_pretty(&agent.clone().freeze()).expect("Failed to serialize test agent to file"); + let content = agent.to_str_pretty().expect("Failed to serialize test agent to file"); let agent_path = PathBuf::from(AGENT_PATH); os.fs .create_dir_all( diff --git a/crates/chat-cli/src/cli/chat/tool_manager.rs b/crates/chat-cli/src/cli/chat/tool_manager.rs index f6999a7846..15dbf7af52 100644 --- a/crates/chat-cli/src/cli/chat/tool_manager.rs +++ b/crates/chat-cli/src/cli/chat/tool_manager.rs @@ -58,7 +58,6 @@ use crate::api_client::model::{ ToolResultContentBlock, ToolResultStatus, }; -use crate::cli::agent::wrapper_types::Warm; use crate::cli::agent::{ Agent, McpServerConfig, @@ -155,7 +154,7 @@ pub struct ToolManagerBuilder { prompt_list_sender: Option>>, prompt_list_receiver: Option>>, conversation_id: Option, - agent: Option>, + agent: Option, } impl ToolManagerBuilder { @@ -174,7 +173,7 @@ impl ToolManagerBuilder { self } - pub fn agent(mut self, agent: Agent) -> Self { + pub fn agent(mut self, agent: Agent) -> Self { if let McpServerConfigWrapper::Map(config) = &agent.mcp_servers { self.mcp_server_config.replace(config.clone()); } else { @@ -193,8 +192,7 @@ impl ToolManagerBuilder { mut output: Box, interactive: bool, ) -> eyre::Result { - let McpServerConfig { mcp_servers, .. } = - self.mcp_server_config.ok_or(eyre::eyre!("Missing mcp server config"))?; + let McpServerConfig { mcp_servers, .. } = self.mcp_server_config.unwrap_or_default(); debug_assert!(self.conversation_id.is_some()); let conversation_id = self.conversation_id.ok_or(eyre::eyre!("Missing conversation id"))?; @@ -866,7 +864,7 @@ pub struct ToolManager { /// A collection of preferences that pertains to the conversation. /// As far as tool manager goes, this is relevant for tool and server filters - pub agent: Arc>>, + pub agent: Arc>, } impl Clone for ToolManager { diff --git a/crates/chat-cli/src/cli/chat/tools/custom_tool.rs b/crates/chat-cli/src/cli/chat/tools/custom_tool.rs index eae2d439d7..cf25b32365 100644 --- a/crates/chat-cli/src/cli/chat/tools/custom_tool.rs +++ b/crates/chat-cli/src/cli/chat/tools/custom_tool.rs @@ -17,7 +17,6 @@ use tokio::sync::RwLock; use tracing::warn; use super::InvokeOutput; -use crate::cli::agent::wrapper_types::Warm; use crate::cli::agent::{ Agent, PermissionEvalResult, @@ -256,7 +255,7 @@ impl CustomTool { + TokenCounter::count_tokens(self.params.as_ref().map_or("", |p| p.as_str().unwrap_or_default())) } - pub fn eval_perm(&self, agent: &Agent) -> PermissionEvalResult { + pub fn eval_perm(&self, agent: &Agent) -> PermissionEvalResult { use crate::util::MCP_SERVER_TOOL_DELIMITER; let Self { name: tool_name, diff --git a/crates/chat-cli/src/cli/chat/tools/execute/mod.rs b/crates/chat-cli/src/cli/chat/tools/execute/mod.rs index 8d073e31cd..9b7cd6e76b 100644 --- a/crates/chat-cli/src/cli/chat/tools/execute/mod.rs +++ b/crates/chat-cli/src/cli/chat/tools/execute/mod.rs @@ -9,7 +9,6 @@ use eyre::Result; use serde::Deserialize; use tracing::error; -use crate::cli::agent::wrapper_types::Warm; use crate::cli::agent::{ Agent, PermissionEvalResult, @@ -155,7 +154,7 @@ impl ExecuteCommand { Ok(()) } - pub fn eval_perm(&self, agent: &Agent) -> PermissionEvalResult { + pub fn eval_perm(&self, agent: &Agent) -> PermissionEvalResult { #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] struct Settings { diff --git a/crates/chat-cli/src/cli/chat/tools/fs_read.rs b/crates/chat-cli/src/cli/chat/tools/fs_read.rs index 90845940ee..6fa2ca97e4 100644 --- a/crates/chat-cli/src/cli/chat/tools/fs_read.rs +++ b/crates/chat-cli/src/cli/chat/tools/fs_read.rs @@ -33,7 +33,6 @@ use super::{ format_path, sanitize_path_tool_arg, }; -use crate::cli::agent::wrapper_types::Warm; use crate::cli::agent::{ Agent, PermissionEvalResult, @@ -100,7 +99,7 @@ impl FsRead { } } - pub fn eval_perm(&self, agent: &Agent) -> PermissionEvalResult { + pub fn eval_perm(&self, agent: &Agent) -> PermissionEvalResult { #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] struct Settings { diff --git a/crates/chat-cli/src/cli/chat/tools/fs_write.rs b/crates/chat-cli/src/cli/chat/tools/fs_write.rs index 9ce28cde1b..3f7c5d41c4 100644 --- a/crates/chat-cli/src/cli/chat/tools/fs_write.rs +++ b/crates/chat-cli/src/cli/chat/tools/fs_write.rs @@ -37,7 +37,6 @@ use super::{ sanitize_path_tool_arg, supports_truecolor, }; -use crate::cli::agent::wrapper_types::Warm; use crate::cli::agent::{ Agent, PermissionEvalResult, @@ -347,7 +346,7 @@ impl FsWrite { } } - pub fn eval_perm(&self, agent: &Agent) -> PermissionEvalResult { + pub fn eval_perm(&self, agent: &Agent) -> PermissionEvalResult { #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] struct Settings { diff --git a/crates/chat-cli/src/cli/chat/tools/mod.rs b/crates/chat-cli/src/cli/chat/tools/mod.rs index 34e7c31910..4b0a121dae 100644 --- a/crates/chat-cli/src/cli/chat/tools/mod.rs +++ b/crates/chat-cli/src/cli/chat/tools/mod.rs @@ -35,7 +35,6 @@ use use_aws::UseAws; use super::consts::MAX_TOOL_RESPONSE_SIZE; use super::util::images::RichImageBlocks; -use crate::cli::agent::wrapper_types::Warm; use crate::cli::agent::{ Agent, PermissionEvalResult, @@ -90,7 +89,7 @@ impl Tool { } /// Whether or not the tool should prompt the user to accept before [Self::invoke] is called. - pub fn requires_acceptance(&self, agent: &Agent) -> PermissionEvalResult { + pub fn requires_acceptance(&self, agent: &Agent) -> PermissionEvalResult { match self { Tool::FsRead(fs_read) => fs_read.eval_perm(agent), Tool::FsWrite(fs_write) => fs_write.eval_perm(agent), diff --git a/crates/chat-cli/src/cli/chat/tools/use_aws.rs b/crates/chat-cli/src/cli/chat/tools/use_aws.rs index 1a017782e8..59b41e8b0d 100644 --- a/crates/chat-cli/src/cli/chat/tools/use_aws.rs +++ b/crates/chat-cli/src/cli/chat/tools/use_aws.rs @@ -23,7 +23,6 @@ use super::{ MAX_TOOL_RESPONSE_SIZE, OutputKind, }; -use crate::cli::agent::wrapper_types::Warm; use crate::cli::agent::{ Agent, PermissionEvalResult, @@ -195,7 +194,7 @@ impl UseAws { } } - pub fn eval_perm(&self, agent: &Agent) -> PermissionEvalResult { + pub fn eval_perm(&self, agent: &Agent) -> PermissionEvalResult { #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] struct Settings { diff --git a/crates/chat-cli/src/cli/mcp.rs b/crates/chat-cli/src/cli/mcp.rs index 56e3215eff..99ac099aab 100644 --- a/crates/chat-cli/src/cli/mcp.rs +++ b/crates/chat-cli/src/cli/mcp.rs @@ -138,7 +138,7 @@ impl AddArgs { }))?; mcp_servers.insert(self.name.clone(), tool); - let json = serde_json::to_string_pretty(&agent.clone().freeze())?; + let json = agent.to_str_pretty()?; os.fs.write(config_path, json).await?; writeln!(output, "✓ Added MCP server '{}' to agent {}\n", self.name, agent_name)?; }, @@ -206,7 +206,7 @@ impl RemoveArgs { match config.remove(&self.name) { Some(_) => { - let json = serde_json::to_string_pretty(&agent.clone().freeze())?; + let json = agent.to_str_pretty()?; os.fs.write(config_path, json).await?; writeln!( output, From f18b6a3b60ac0db0caeda422df67ec0e62004415 Mon Sep 17 00:00:00 2001 From: Felix dingfeli Date: Wed, 16 Jul 2025 17:13:46 -0700 Subject: [PATCH 09/15] assign agent name to agent during thaw --- crates/chat-cli/src/cli/agent/mod.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/crates/chat-cli/src/cli/agent/mod.rs b/crates/chat-cli/src/cli/agent/mod.rs index 51a4922ed2..482ae9dd6b 100644 --- a/crates/chat-cli/src/cli/agent/mod.rs +++ b/crates/chat-cli/src/cli/agent/mod.rs @@ -202,6 +202,8 @@ impl Agent { .to_string_lossy() .to_string(); + self.name = name.clone(); + if let McpServerConfigWrapper::List(selected_servers) = mcp_servers { let include_all = selected_servers.len() == 1 && selected_servers.first().is_some_and(|s| s.as_str() == "*"); From 5e561c908735fb33ae3a815a8d4aa867b3dd2ba1 Mon Sep 17 00:00:00 2001 From: Felix dingfeli Date: Wed, 16 Jul 2025 17:23:01 -0700 Subject: [PATCH 10/15] adds comment for freeze and thaw --- crates/chat-cli/src/cli/agent/mod.rs | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/crates/chat-cli/src/cli/agent/mod.rs b/crates/chat-cli/src/cli/agent/mod.rs index 482ae9dd6b..bf4b25057a 100644 --- a/crates/chat-cli/src/cli/agent/mod.rs +++ b/crates/chat-cli/src/cli/agent/mod.rs @@ -165,6 +165,9 @@ impl Default for Agent { } impl Agent { + /// This function mutates the agent to a state that is writable. + /// Practically this means reverting some fields back to their original values as they were + /// written in the config. pub fn freeze(&mut self) -> eyre::Result<()> { let Self { mcp_servers, .. } = self; @@ -193,6 +196,10 @@ impl Agent { Ok(()) } + /// This function mutates the agent to a state that is usable for runtime. + /// Practically this means to convert some of the fields value to their usable counterpart. + /// For example, we populate the agent with its file name, convert the mcp array to actual + /// mcp config and populate the agent file path. pub fn thaw(&mut self, path: &Path, global_mcp_config: Option<&McpServerConfig>) -> eyre::Result<()> { let Self { mcp_servers, .. } = self; From 5e8b3dd1be2aa0c3f0f9b603649d0a118fec573e Mon Sep 17 00:00:00 2001 From: Felix dingfeli Date: Thu, 17 Jul 2025 14:17:32 -0700 Subject: [PATCH 11/15] monomorphisizes mcp server field in agent config --- .../chat-cli/src/cli/agent/context_migrate.rs | 3 +- crates/chat-cli/src/cli/agent/mcp_config.rs | 17 --- crates/chat-cli/src/cli/agent/mod.rs | 123 ++++++++---------- .../chat-cli/src/cli/agent/wrapper_types.rs | 33 ----- crates/chat-cli/src/cli/chat/mod.rs | 5 +- crates/chat-cli/src/cli/chat/tool_manager.rs | 12 +- .../src/cli/chat/tools/custom_tool.rs | 4 + crates/chat-cli/src/cli/mcp.rs | 26 ++-- 8 files changed, 70 insertions(+), 153 deletions(-) diff --git a/crates/chat-cli/src/cli/agent/context_migrate.rs b/crates/chat-cli/src/cli/agent/context_migrate.rs index 0711dc860b..7a4d9fb4ce 100644 --- a/crates/chat-cli/src/cli/agent/context_migrate.rs +++ b/crates/chat-cli/src/cli/agent/context_migrate.rs @@ -14,7 +14,6 @@ use super::{ }; use crate::cli::agent::{ CreateHooks, - McpServerConfigWrapper, PromptHooks, }; use crate::cli::chat::cli::hooks::{ @@ -223,7 +222,7 @@ impl ContextMigrate<'c'> { agent.name ); } - agent.mcp_servers = McpServerConfigWrapper::Map(mcp_servers.clone().unwrap_or_default()); + agent.mcp_servers = mcp_servers.clone().unwrap_or_default(); } let legacy_profile_config_path = directories::chat_profiles_dir(os)?; diff --git a/crates/chat-cli/src/cli/agent/mcp_config.rs b/crates/chat-cli/src/cli/agent/mcp_config.rs index c848904391..e6adfc10cd 100644 --- a/crates/chat-cli/src/cli/agent/mcp_config.rs +++ b/crates/chat-cli/src/cli/agent/mcp_config.rs @@ -10,28 +10,11 @@ use serde::{ use crate::cli::chat::tools::custom_tool::CustomToolConfig; use crate::os::Os; -#[derive(Clone, Debug, Default, Eq, PartialEq)] -pub enum OriginalForm { - /// The mcp config has been written in isolation, with no reference to global mcp.json - #[default] - Map, - /// The mcp config has been written as an include-all array, i.e. '["*"]' - All, - /// The mcp config has been written as an array containing a list of mcp servers to include - /// from global mcp.json - List, -} - // This is to mirror claude's config set up #[derive(Clone, Serialize, Deserialize, Debug, Default, Eq, PartialEq, JsonSchema)] #[serde(rename_all = "camelCase", transparent)] pub struct McpServerConfig { pub mcp_servers: HashMap, - /// This field keeps track of how the mcp server field is configured as seen on the user's - /// config file. This is important because when we write the config to file, we want to - /// preserve its behavior the next time it is read. - #[serde(skip)] - pub original_config_form: OriginalForm, } impl McpServerConfig { diff --git a/crates/chat-cli/src/cli/agent/mod.rs b/crates/chat-cli/src/cli/agent/mod.rs index bf4b25057a..15a8c820e8 100644 --- a/crates/chat-cli/src/cli/agent/mod.rs +++ b/crates/chat-cli/src/cli/agent/mod.rs @@ -25,7 +25,6 @@ use crossterm::{ }; use eyre::bail; pub use mcp_config::McpServerConfig; -use mcp_config::OriginalForm; use regex::Regex; use schemars::JsonSchema; use serde::{ @@ -39,7 +38,6 @@ use tracing::{ }; pub use wrapper_types::{ CreateHooks, - McpServerConfigWrapper, OriginalToolName, PromptHooks, ToolSettingTarget, @@ -107,7 +105,7 @@ pub struct Agent { pub prompt: Option, /// Configuration for Model Context Protocol (MCP) servers #[serde(default)] - pub mcp_servers: McpServerConfigWrapper, + pub mcp_servers: McpServerConfig, /// List of tools the agent can see. Use \"@{MCP_SERVER_NAME}/tool_name\" to specify tools from /// mcp servers. To include all tools from a server, use \"@{MCP_SERVER_NAME}\" #[serde(default)] @@ -133,6 +131,11 @@ pub struct Agent { #[serde(default)] #[schemars(schema_with = "tool_settings_schema")] pub tools_settings: HashMap, + /// Whether or not to include the legacy ~/.aws/amazonq/mcp.json in the agent + /// You can reference tools brought in by these servers as just as you would with the servers + /// you configure in the mcpServers field in this config + #[serde(default)] + pub use_legacy_mcp_json: bool, #[serde(skip)] pub path: Option, } @@ -159,6 +162,7 @@ impl Default for Agent { create_hooks: Default::default(), prompt_hooks: Default::default(), tools_settings: Default::default(), + use_legacy_mcp_json: true, path: None, } } @@ -171,27 +175,9 @@ impl Agent { pub fn freeze(&mut self) -> eyre::Result<()> { let Self { mcp_servers, .. } = self; - if let McpServerConfigWrapper::Map(servers) = mcp_servers { - match servers.original_config_form { - OriginalForm::All => { - *mcp_servers = McpServerConfigWrapper::List(vec!["*".to_string()]); - }, - OriginalForm::List => { - *mcp_servers = McpServerConfigWrapper::List({ - servers - .mcp_servers - .iter() - .fold(Vec::::new(), |mut acc, (name, _)| { - acc.push(name.clone()); - acc - }) - }); - }, - OriginalForm::Map => { - // noop - }, - } - } + mcp_servers + .mcp_servers + .retain(|_name, config| !config.is_from_legacy_mcp_json); Ok(()) } @@ -211,35 +197,29 @@ impl Agent { self.name = name.clone(); - if let McpServerConfigWrapper::List(selected_servers) = mcp_servers { - let include_all = - selected_servers.len() == 1 && selected_servers.first().is_some_and(|s| s.as_str() == "*"); - let filtered_config = if let Some(global_mcp_config) = global_mcp_config { - global_mcp_config - .mcp_servers - .iter() - .filter_map(|(name, config)| { - if selected_servers.contains(name) || include_all { - Some((name.clone(), config.clone())) - } else { - None - } - }) - .collect::>() - } else { - bail!( - "Error converting agent {name} to usable state. Mcp config is given as list but global mcp config is missing. Aborting" - ); - }; - - *mcp_servers = McpServerConfigWrapper::Map(McpServerConfig { - mcp_servers: filtered_config, - original_config_form: if include_all { - OriginalForm::All - } else { - OriginalForm::Map - }, - }); + if let (true, Some(global_mcp_config)) = (self.use_legacy_mcp_json, global_mcp_config) { + let mut stderr = std::io::stderr(); + for (name, legacy_server) in &global_mcp_config.mcp_servers { + if mcp_servers.mcp_servers.contains_key(name) { + let _ = queue!( + stderr, + style::SetForegroundColor(Color::Yellow), + style::Print("WARNING: "), + style::ResetColor, + style::Print("MCP server '"), + style::SetForegroundColor(Color::Green), + style::Print(name), + style::ResetColor, + style::Print( + "' is already configured in agent config. Skipping duplicate from legacy mcp.json.\n" + ) + ); + continue; + } + let mut server_clone = legacy_server.clone(); + server_clone.is_from_legacy_mcp_json = true; + mcp_servers.mcp_servers.insert(name.clone(), server_clone); + } } Ok(()) @@ -685,19 +665,23 @@ async fn load_agents_from_entries( }, }; - if matches!(agent.mcp_servers, McpServerConfigWrapper::List(_)) && global_mcp_config.is_none() { - let Ok(global_mcp_path) = directories::chat_legacy_mcp_config(os) else { - tracing::error!("Error obtaining legacy mcp json path. Skipping"); - continue; - }; - let legacy_mcp_config = match McpServerConfig::load_from_file(os, global_mcp_path).await { - Ok(config) => config, - Err(e) => { - tracing::error!("Error loading global mcp json path: {e}. Skipping"); - continue; - }, - }; - global_mcp_config.replace(legacy_mcp_config); + // The agent config could have use_legacy_mcp_json set to true but not have a valid + // global mcp.json. We would still need to carry on loading the config. + 'load_legacy_mcp_json: { + if agent.use_legacy_mcp_json && global_mcp_config.is_none() { + let Ok(global_mcp_path) = directories::chat_legacy_mcp_config(os) else { + tracing::error!("Error obtaining legacy mcp json path. Skipping"); + break 'load_legacy_mcp_json; + }; + let legacy_mcp_config = match McpServerConfig::load_from_file(os, global_mcp_path).await { + Ok(config) => config, + Err(e) => { + tracing::error!("Error loading global mcp json path: {e}. Skipping"); + break 'load_legacy_mcp_json; + }, + }; + global_mcp_config.replace(legacy_mcp_config); + } } if let Err(e) = agent.thaw(file_path, global_mcp_config.as_ref()) { @@ -796,11 +780,8 @@ mod tests { #[test] fn test_deser() { let agent = serde_json::from_str::(INPUT).expect("Deserializtion failed"); - assert!(matches!(agent.mcp_servers, McpServerConfigWrapper::Map(_))); - if let McpServerConfigWrapper::Map(config) = &agent.mcp_servers { - assert!(config.mcp_servers.contains_key("fetch")); - assert!(config.mcp_servers.contains_key("git")); - } + assert!(agent.mcp_servers.mcp_servers.contains_key("fetch")); + assert!(agent.mcp_servers.mcp_servers.contains_key("git")); assert!(agent.alias.contains_key("@gits/some_tool")); } diff --git a/crates/chat-cli/src/cli/agent/wrapper_types.rs b/crates/chat-cli/src/cli/agent/wrapper_types.rs index 2c763a6e4f..a9fefba32e 100644 --- a/crates/chat-cli/src/cli/agent/wrapper_types.rs +++ b/crates/chat-cli/src/cli/agent/wrapper_types.rs @@ -13,7 +13,6 @@ use serde::{ Serialize, }; -use super::McpServerConfig; use crate::cli::chat::cli::hooks::Hook; /// Subject of the tool name change. For tools in mcp servers, you would need to prefix them with @@ -123,35 +122,3 @@ pub fn tool_settings_schema(generator: &mut SchemaGenerator) -> Schema { } }) } - -#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, JsonSchema)] -#[serde(untagged)] -pub enum McpServerConfigWrapper { - /// Array of servers to instantiate in accordance to what is in the global mcp.json - List(Vec), - /// A map of mcp server configurations, identical in format to the value of mcpServers in - /// the global mcp.json - Map(McpServerConfig), -} - -impl Default for McpServerConfigWrapper { - fn default() -> Self { - Self::List(vec!["*".to_string()]) - } -} - -impl McpServerConfigWrapper { - pub fn get_server_names(&self) -> Vec<&str> { - match self { - Self::List(list) => list.iter().map(|name| name.as_str()).collect::>(), - Self::Map(config) => config.mcp_servers.keys().map(|name| name.as_str()).collect::>(), - } - } - - pub fn is_empty(&self) -> bool { - match self { - Self::List(list) => list.is_empty(), - Self::Map(config) => config.mcp_servers.is_empty(), - } - } -} diff --git a/crates/chat-cli/src/cli/chat/mod.rs b/crates/chat-cli/src/cli/chat/mod.rs index 127216f4ef..b3a694e2d9 100644 --- a/crates/chat-cli/src/cli/chat/mod.rs +++ b/crates/chat-cli/src/cli/chat/mod.rs @@ -233,7 +233,10 @@ impl ChatArgs { let mut agents = Agents::load(os, self.agent.as_deref(), skip_migration, &mut stderr).await; agents.trust_all_tools = self.trust_all_tools; - if agents.get_active().is_some_and(|a| !a.mcp_servers.is_empty()) { + if agents + .get_active() + .is_some_and(|a| !a.mcp_servers.mcp_servers.is_empty()) + { if !self.no_interactive && !os.database.settings.get_bool(Setting::McpLoadedBefore).unwrap_or(false) { execute!( stderr, diff --git a/crates/chat-cli/src/cli/chat/tool_manager.rs b/crates/chat-cli/src/cli/chat/tool_manager.rs index 15dbf7af52..99a52e8c35 100644 --- a/crates/chat-cli/src/cli/chat/tool_manager.rs +++ b/crates/chat-cli/src/cli/chat/tool_manager.rs @@ -61,7 +61,6 @@ use crate::api_client::model::{ use crate::cli::agent::{ Agent, McpServerConfig, - McpServerConfigWrapper, }; use crate::cli::chat::cli::prompts::GetPromptError; use crate::cli::chat::message::AssistantToolUse; @@ -150,7 +149,6 @@ pub enum LoadingRecord { #[derive(Default)] pub struct ToolManagerBuilder { - mcp_server_config: Option, prompt_list_sender: Option>>, prompt_list_receiver: Option>>, conversation_id: Option, @@ -174,14 +172,6 @@ impl ToolManagerBuilder { } pub fn agent(mut self, agent: Agent) -> Self { - if let McpServerConfigWrapper::Map(config) = &agent.mcp_servers { - self.mcp_server_config.replace(config.clone()); - } else { - error!( - "No valid mcp config valid in agent {}, no mcp config loaded", - &agent.name - ); - } self.agent.replace(agent); self } @@ -192,7 +182,7 @@ impl ToolManagerBuilder { mut output: Box, interactive: bool, ) -> eyre::Result { - let McpServerConfig { mcp_servers, .. } = self.mcp_server_config.unwrap_or_default(); + let McpServerConfig { mcp_servers } = self.agent.as_ref().map(|a| a.mcp_servers.clone()).unwrap_or_default(); debug_assert!(self.conversation_id.is_some()); let conversation_id = self.conversation_id.ok_or(eyre::eyre!("Missing conversation id"))?; diff --git a/crates/chat-cli/src/cli/chat/tools/custom_tool.rs b/crates/chat-cli/src/cli/chat/tools/custom_tool.rs index cf25b32365..65b66bd04b 100644 --- a/crates/chat-cli/src/cli/chat/tools/custom_tool.rs +++ b/crates/chat-cli/src/cli/chat/tools/custom_tool.rs @@ -54,6 +54,9 @@ pub struct CustomToolConfig { /// A boolean flag to denote whether or not to load this mcp server #[serde(default)] pub disabled: bool, + /// A flag to denote whether this is a server from the legacy mcp.json + #[serde(skip)] + pub is_from_legacy_mcp_json: bool, } pub fn default_timeout() -> u64 { @@ -79,6 +82,7 @@ impl CustomToolClient { env, timeout, disabled: _, + .. } = config; let mcp_client_config = McpClientConfig { server_name: server_name.clone(), diff --git a/crates/chat-cli/src/cli/mcp.rs b/crates/chat-cli/src/cli/mcp.rs index 99ac099aab..14e4bf5352 100644 --- a/crates/chat-cli/src/cli/mcp.rs +++ b/crates/chat-cli/src/cli/mcp.rs @@ -22,7 +22,6 @@ use super::agent::{ Agents, McpServerConfig, }; -use crate::cli::agent::McpServerConfigWrapper; use crate::cli::chat::tool_manager::{ global_mcp_config_path, workspace_mcp_config_path, @@ -113,11 +112,7 @@ impl AddArgs { match self.agent.as_deref() { Some(agent_name) => { let (mut agent, config_path) = Agent::get_agent_by_name(os, agent_name).await?; - let mcp_servers = if let McpServerConfigWrapper::Map(config) = &mut agent.mcp_servers { - &mut config.mcp_servers - } else { - bail!("Mcp server config not found on agent"); - }; + let mcp_servers = &mut agent.mcp_servers.mcp_servers; if mcp_servers.contains_key(&self.name) && !self.force { bail!( @@ -198,11 +193,7 @@ impl RemoveArgs { return Ok(()); } - let config = if let McpServerConfigWrapper::Map(config) = &mut agent.mcp_servers { - &mut config.mcp_servers - } else { - bail!("Mcp server config not found on agent"); - }; + let config = &mut agent.mcp_servers.mcp_servers; match config.remove(&self.name) { Some(_) => { @@ -405,13 +396,12 @@ async fn get_mcp_server_configs( } else { Scope::Workspace }; - if let McpServerConfigWrapper::Map(config) = agent.mcp_servers { - results.push(( - scope, - agent.path.ok_or(eyre::eyre!("Agent missing path info"))?, - Some(config), - )); - } + + results.push(( + scope, + agent.path.ok_or(eyre::eyre!("Agent missing path info"))?, + Some(agent.mcp_servers), + )); } Ok(results) } From 3a82b555cb34f1edbe565bc257e85d6e5999419f Mon Sep 17 00:00:00 2001 From: Felix dingfeli Date: Thu, 17 Jul 2025 15:18:58 -0700 Subject: [PATCH 12/15] modifies migration flow to use database for migration status --- .../chat-cli/src/cli/agent/context_migrate.rs | 42 +++++++++---------- crates/chat-cli/src/cli/chat/mod.rs | 5 +-- crates/chat-cli/src/cli/mod.rs | 9 ---- crates/chat-cli/src/database/mod.rs | 11 +++++ 4 files changed, 31 insertions(+), 36 deletions(-) diff --git a/crates/chat-cli/src/cli/agent/context_migrate.rs b/crates/chat-cli/src/cli/agent/context_migrate.rs index 7a4d9fb4ce..3843eda6d4 100644 --- a/crates/chat-cli/src/cli/agent/context_migrate.rs +++ b/crates/chat-cli/src/cli/agent/context_migrate.rs @@ -1,5 +1,9 @@ use std::collections::HashMap; +use crossterm::{ + execute, + style, +}; use dialoguer::Select; use eyre::bail; use tracing::{ @@ -34,6 +38,11 @@ pub(in crate::cli::agent) struct ContextMigrate { impl ContextMigrate<'a'> { pub async fn scan(os: &Os) -> eyre::Result> { + let has_migrated = os.database.get_has_migrated()?; + if has_migrated.is_some_and(|has_migrated| has_migrated) { + bail!("Nothing to migrate"); + } + let legacy_global_context_path = directories::chat_global_context_path(os)?; let legacy_global_context: Option = 'global: { let Ok(content) = os.fs.read(&legacy_global_context_path).await else { @@ -143,6 +152,15 @@ impl ContextMigrate<'b'> { new_agents, }) } else { + let _ = execute!( + std::io::stdout(), + style::SetForegroundColor(style::Color::Yellow), + style::Print("WARNING: "), + style::ResetColor, + style::Print( + "Migration aborted. Context in profiles will need to be migrated to agents manually as profiles are no longer supported." + ), + ); bail!("Aborting migration") } } @@ -161,8 +179,6 @@ impl ContextMigrate<'c'> { mut new_agents, } = self; - let has_global_context = legacy_global_context.is_some(); - // Migration of global context if let Some(context) = legacy_global_context { let (create_hooks, prompt_hooks) = @@ -225,27 +241,7 @@ impl ContextMigrate<'c'> { agent.mcp_servers = mcp_servers.clone().unwrap_or_default(); } - let legacy_profile_config_path = directories::chat_profiles_dir(os)?; - let profile_backup_path = legacy_profile_config_path - .parent() - .ok_or(eyre::eyre!("Failed to obtain profile config parent path"))? - .join("profiles.bak"); - os.fs.rename(legacy_profile_config_path, profile_backup_path).await?; - - if has_global_context { - let legacy_global_config_path = directories::chat_global_context_path(os)?; - let legacy_global_config_file_name = legacy_global_config_path - .file_name() - .ok_or(eyre::eyre!("Failed to obtain legacy global config name"))? - .to_string_lossy(); - let global_context_backup_path = legacy_global_config_path - .parent() - .ok_or(eyre::eyre!("Failed to obtain parent path for global context"))? - .join(format!("{}.bak", legacy_global_config_file_name)); - os.fs - .rename(legacy_global_config_path, global_context_backup_path) - .await?; - } + os.database.set_has_migrated()?; Ok(ContextMigrate { legacy_global_context: None, diff --git a/crates/chat-cli/src/cli/chat/mod.rs b/crates/chat-cli/src/cli/chat/mod.rs index b3a694e2d9..b77d3425b0 100644 --- a/crates/chat-cli/src/cli/chat/mod.rs +++ b/crates/chat-cli/src/cli/chat/mod.rs @@ -178,9 +178,6 @@ pub struct ChatArgs { pub no_interactive: bool, /// The first question to ask pub input: Option, - /// Run migration of legacy profiles to agents if applicable - #[arg(long)] - pub migrate: bool, } impl ChatArgs { @@ -229,7 +226,7 @@ impl ChatArgs { } let agents = { - let skip_migration = self.no_interactive || !self.migrate; + let skip_migration = self.no_interactive; let mut agents = Agents::load(os, self.agent.as_deref(), skip_migration, &mut stderr).await; agents.trust_all_tools = self.trust_all_tools; diff --git a/crates/chat-cli/src/cli/mod.rs b/crates/chat-cli/src/cli/mod.rs index 3d90a9e79b..c51e5df3eb 100644 --- a/crates/chat-cli/src/cli/mod.rs +++ b/crates/chat-cli/src/cli/mod.rs @@ -362,7 +362,6 @@ mod test { trust_all_tools: false, trust_tools: None, no_interactive: false, - migrate: false, })), verbose: 2, help_all: false, @@ -402,7 +401,6 @@ mod test { trust_all_tools: false, trust_tools: None, no_interactive: false, - migrate: false, }) ); } @@ -419,7 +417,6 @@ mod test { trust_all_tools: false, trust_tools: None, no_interactive: false, - migrate: false, }) ); } @@ -436,7 +433,6 @@ mod test { trust_all_tools: true, trust_tools: None, no_interactive: false, - migrate: false, }) ); } @@ -453,7 +449,6 @@ mod test { trust_all_tools: false, trust_tools: None, no_interactive: true, - migrate: false, }) ); assert_parse!( @@ -466,7 +461,6 @@ mod test { trust_all_tools: false, trust_tools: None, no_interactive: true, - migrate: false, }) ); } @@ -483,7 +477,6 @@ mod test { trust_all_tools: true, trust_tools: None, no_interactive: false, - migrate: false, }) ); } @@ -500,7 +493,6 @@ mod test { trust_all_tools: false, trust_tools: Some(vec!["".to_string()]), no_interactive: false, - migrate: false, }) ); } @@ -517,7 +509,6 @@ mod test { trust_all_tools: false, trust_tools: Some(vec!["fs_read".to_string(), "fs_write".to_string()]), no_interactive: false, - migrate: false, }) ); } diff --git a/crates/chat-cli/src/database/mod.rs b/crates/chat-cli/src/database/mod.rs index bccee45fd4..9b5a48ee10 100644 --- a/crates/chat-cli/src/database/mod.rs +++ b/crates/chat-cli/src/database/mod.rs @@ -60,6 +60,7 @@ const START_URL_KEY: &str = "auth.idc.start-url"; const IDC_REGION_KEY: &str = "auth.idc.region"; // We include this key to remove for backwards compatibility const CUSTOMIZATION_STATE_KEY: &str = "api.selectedCustomization"; +const PROFILE_MIGRATION_KEY: &str = "profile.Migrated"; const MIGRATIONS: &[Migration] = migrations![ "000_migration_table", @@ -300,6 +301,16 @@ impl Database { self.set_json_entry(Table::State, IDC_REGION_KEY, region) } + /// Get if user has already completed a migration + pub fn get_has_migrated(&self) -> Result, DatabaseError> { + self.get_entry::(Table::State, PROFILE_MIGRATION_KEY) + } + + /// Set if user has already completed a migration + pub fn set_has_migrated(&self) -> Result { + self.set_entry(Table::State, PROFILE_MIGRATION_KEY, true) + } + // /// Get the model id used for last conversation state. // pub fn get_last_used_model_id(&self) -> Result, DatabaseError> { // self.get_json_entry::(Table::State, LAST_USED_MODEL_ID) From 987fc79df8930b31efa60014d92a8e22f1a4920c Mon Sep 17 00:00:00 2001 From: Felix dingfeli Date: Thu, 17 Jul 2025 15:32:47 -0700 Subject: [PATCH 13/15] updates agent doc comments --- crates/chat-cli/src/cli/agent/mod.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/crates/chat-cli/src/cli/agent/mod.rs b/crates/chat-cli/src/cli/agent/mod.rs index 15a8c820e8..0d3b5d0244 100644 --- a/crates/chat-cli/src/cli/agent/mod.rs +++ b/crates/chat-cli/src/cli/agent/mod.rs @@ -81,10 +81,10 @@ pub use root_command_args::*; /// had created in said file. This field is thus populated during its parsing. /// /// Another example is the mcp config. To support backwards compatibility of users existing global -/// mcp.json, we allow users to supply a list of mcp server names they want the agent to -/// instantiate. But in order for this config to actually be useful to the CLI, it needs to contain -/// actual information about the command, not just the list of names. Thus the "warm" state in this -/// field would be a filtered version of [McpServerConfig], while the "cold" state could be either. +/// mcp.json, we allow users to supply a flag to denote whether they would want to include servers +/// from the legacy global mcp.json. If this flag exists, we would need to read the legacy mcp +/// config and merge it with what is in the agent mcp servers field. Conversely, when we write this +/// config to file, we would want to filter out the servers that belong only in the mcp.json. /// /// Where agents are instantiated from their config, we would need to convert them from "cold" to /// "warm". From 518cb86f426320f3f9207925d9d3a896f49e91ce Mon Sep 17 00:00:00 2001 From: Felix dingfeli Date: Fri, 18 Jul 2025 10:16:53 -0700 Subject: [PATCH 14/15] fixes errors from merge --- crates/chat-cli/src/cli/agent/legacy/mod.rs | 51 ++++++--------------- crates/chat-cli/src/cli/agent/mod.rs | 2 - 2 files changed, 14 insertions(+), 39 deletions(-) diff --git a/crates/chat-cli/src/cli/agent/legacy/mod.rs b/crates/chat-cli/src/cli/agent/legacy/mod.rs index 1039357f34..41ecf39c4c 100644 --- a/crates/chat-cli/src/cli/agent/legacy/mod.rs +++ b/crates/chat-cli/src/cli/agent/legacy/mod.rs @@ -21,6 +21,11 @@ use crate::os::Os; use crate::util::directories; pub async fn migrate(os: &mut Os) -> eyre::Result> { + let has_migrated = os.database.get_has_migrated()?; + if has_migrated.is_some_and(|has_migrated| has_migrated) { + bail!("Nothing to migrate"); + } + let legacy_global_context_path = directories::chat_global_context_path(os)?; let legacy_global_context: Option = 'global: { let Ok(content) = os.fs.read(&legacy_global_context_path).await else { @@ -70,7 +75,12 @@ pub async fn migrate(os: &mut Os) -> eyre::Result> { let config_path = directories::chat_legacy_mcp_config(os)?; if os.fs.exists(&config_path) { match McpServerConfig::load_from_file(os, config_path).await { - Ok(config) => Some(config), + Ok(mut config) => { + config.mcp_servers.iter_mut().for_each(|(_name, config)| { + config.is_from_legacy_mcp_json = true; + }); + Some(config) + }, Err(e) => { error!("Malformed legacy global mcp config detected: {e}. Skipping mcp migration."); None @@ -82,6 +92,7 @@ pub async fn migrate(os: &mut Os) -> eyre::Result> { }; if legacy_global_context.is_none() && legacy_profiles.is_empty() { + os.database.set_has_migrated()?; bail!("Nothing to migrate"); } @@ -113,8 +124,6 @@ pub async fn migrate(os: &mut Os) -> eyre::Result> { const DEFAULT_DESC: &str = "This is an agent migrated from global context"; const PROFILE_DESC: &str = "This is an agent migrated from profile context"; - let has_global_context = legacy_global_context.is_some(); - // Migration of global context let mut new_agents = vec![]; if let Some(context) = legacy_global_context { @@ -192,20 +201,8 @@ pub async fn migrate(os: &mut Os) -> eyre::Result> { os.fs.create_dir_all(&global_agent_path).await?; } - let formatted_server_list = mcp_servers - .map(|config| { - config - .mcp_servers - .keys() - .map(|server_name| format!("@{server_name}")) - .collect::>() - }) - .unwrap_or_default(); - for agent in &mut new_agents { - agent.tools.extend(formatted_server_list.clone()); - - let content = serde_json::to_string_pretty(agent)?; + let content = agent.to_str_pretty()?; if let Some(path) = agent.path.as_ref() { info!("Agent {} peristed in path {}", agent.name, path.to_string_lossy()); os.fs.write(path, content).await?; @@ -217,27 +214,7 @@ pub async fn migrate(os: &mut Os) -> eyre::Result> { } } - let legacy_profile_config_path = directories::chat_profiles_dir(os)?; - let profile_backup_path = legacy_profile_config_path - .parent() - .ok_or(eyre::eyre!("Failed to obtain profile config parent path"))? - .join("profiles.bak"); - os.fs.rename(legacy_profile_config_path, profile_backup_path).await?; - - if has_global_context { - let legacy_global_config_path = directories::chat_global_context_path(os)?; - let legacy_global_config_file_name = legacy_global_config_path - .file_name() - .ok_or(eyre::eyre!("Failed to obtain legacy global config name"))? - .to_string_lossy(); - let global_context_backup_path = legacy_global_config_path - .parent() - .ok_or(eyre::eyre!("Failed to obtain parent path for global context"))? - .join(format!("{}.bak", legacy_global_config_file_name)); - os.fs - .rename(legacy_global_config_path, global_context_backup_path) - .await?; - } + os.database.set_has_migrated()?; Ok(new_agents) } diff --git a/crates/chat-cli/src/cli/agent/mod.rs b/crates/chat-cli/src/cli/agent/mod.rs index c6dd94afc8..681ae39a5c 100644 --- a/crates/chat-cli/src/cli/agent/mod.rs +++ b/crates/chat-cli/src/cli/agent/mod.rs @@ -63,8 +63,6 @@ use crate::util::{ directories, }; -mod context_migrate; - /// An [Agent] is a declarative way of configuring a given instance of q chat. Currently, it is /// impacting q chat in via influenicng [ContextManager] and [ToolManager]. /// Changes made to [ContextManager] and [ToolManager] do not persist across sessions. From 3c82089e590fae216748a7c8d16dd58268025e80 Mon Sep 17 00:00:00 2001 From: Felix dingfeli Date: Fri, 18 Jul 2025 11:32:08 -0700 Subject: [PATCH 15/15] makes freeze and thaw private --- crates/chat-cli/src/cli/agent/mod.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/chat-cli/src/cli/agent/mod.rs b/crates/chat-cli/src/cli/agent/mod.rs index 681ae39a5c..6078286eaf 100644 --- a/crates/chat-cli/src/cli/agent/mod.rs +++ b/crates/chat-cli/src/cli/agent/mod.rs @@ -167,7 +167,7 @@ impl Agent { /// This function mutates the agent to a state that is writable. /// Practically this means reverting some fields back to their original values as they were /// written in the config. - pub fn freeze(&mut self) -> eyre::Result<()> { + fn freeze(&mut self) -> eyre::Result<()> { let Self { mcp_servers, .. } = self; mcp_servers @@ -181,7 +181,7 @@ impl Agent { /// Practically this means to convert some of the fields value to their usable counterpart. /// For example, we populate the agent with its file name, convert the mcp array to actual /// mcp config and populate the agent file path. - pub fn thaw(&mut self, path: &Path, global_mcp_config: Option<&McpServerConfig>) -> eyre::Result<()> { + fn thaw(&mut self, path: &Path, global_mcp_config: Option<&McpServerConfig>) -> eyre::Result<()> { let Self { mcp_servers, .. } = self; let name = path