From 0717fde86ab750e9ef2b80233c75e5dadc97c210 Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Mon, 1 Jun 2026 13:22:44 -0600 Subject: [PATCH] feat: dynamic loading/unloading of skill tools and MCP servers whenever load_skill/unload_skill are invoked --- src/config/request_context.rs | 61 +++++++++++++++++++++++++++++++++-- src/config/skill_registry.rs | 15 +++++++++ src/function/mod.rs | 12 ++++--- src/function/skill.rs | 50 +++++++++++++++++----------- 4 files changed, 113 insertions(+), 25 deletions(-) diff --git a/src/config/request_context.rs b/src/config/request_context.rs index e7a08ce..9a191dd 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -1,5 +1,6 @@ use super::rag_cache::{RagCache, RagKey}; use super::session::Session; +use super::skill_policy::SkillPolicy; use super::skill_registry::SkillRegistry; use super::todo::TodoList; use super::tool_scope::{McpRuntime, ToolScope}; @@ -35,7 +36,7 @@ use indexmap::IndexMap; use indoc::formatdoc; use inquire::{Confirm, MultiSelect, Text, list_option::ListOption, validator::Validation}; use parking_lot::RwLock; -use std::collections::{HashMap, HashSet}; +use std::collections::{BTreeSet, HashMap, HashSet}; use std::fs::{File, OpenOptions, read_dir, read_to_string, remove_dir_all, remove_file}; use std::io::Write; use std::path::{Path, PathBuf}; @@ -617,7 +618,7 @@ impl RequestContext { } } - role + self.skill_registry.effective_role(&role) } pub fn auto_continue_config(&self) -> AutoContinueConfig { @@ -2067,6 +2068,35 @@ impl RequestContext { enabled_mcp_servers: Option, abort_signal: AbortSignal, ) -> Result<()> { + let policy = SkillPolicy::effective( + app, + self.role.as_ref(), + self.agent.as_ref(), + self.session.as_ref(), + )?; + + let enabled_mcp_servers = if policy.skills_enabled && app.mcp_server_support { + let skill_mcps = self.skill_registry.loaded_mcp_servers(); + match (enabled_mcp_servers.as_deref(), skill_mcps.is_empty()) { + (Some("all"), _) | (_, true) => enabled_mcp_servers, + (base, false) => { + let mut merged: BTreeSet = skill_mcps; + if let Some(s) = base { + for token in s.split(',') { + let t = token.trim(); + if !t.is_empty() { + merged.insert(t.to_string()); + } + } + } + + Some(merged.into_iter().collect::>().join(",")) + } + } + } else { + enabled_mcp_servers + }; + let mut mcp_runtime = McpRuntime::new(); if app.mcp_server_support @@ -2134,6 +2164,9 @@ impl RequestContext { if !mcp_runtime.is_empty() { functions.append_mcp_meta_functions(mcp_runtime.server_names()); } + if app.function_calling_support && policy.skills_enabled { + functions.append_skill_functions(); + } let tool_tracker = self.tool_scope.tool_tracker.clone(); self.tool_scope = ToolScope { @@ -2144,6 +2177,30 @@ impl RequestContext { Ok(()) } + pub async fn refresh_tool_scope(&mut self, abort_signal: AbortSignal) -> Result<()> { + let app = (*self.app.config).clone(); + let base_mcps = if app.mcp_server_support { + if let Some(session) = &self.session { + session.enabled_mcp_servers() + } else if let Some(agent) = &self.agent { + let names = agent.mcp_server_names(); + if names.is_empty() { + None + } else { + Some(names.join(",")) + } + } else if let Some(role) = &self.role { + role.enabled_mcp_servers() + } else { + app.enabled_mcp_servers.clone() + } + } else { + None + }; + + self.rebuild_tool_scope(&app, base_mcps, abort_signal).await + } + pub async fn use_role( &mut self, app: &AppConfig, diff --git a/src/config/skill_registry.rs b/src/config/skill_registry.rs index 7cd5f49..c43afd3 100644 --- a/src/config/skill_registry.rs +++ b/src/config/skill_registry.rs @@ -53,6 +53,21 @@ impl SkillRegistry { self.loaded.keys().cloned().collect() } + pub fn loaded_mcp_servers(&self) -> BTreeSet { + let mut out = BTreeSet::new(); + for skill in self.loaded.values() { + if let Some(csv) = skill.enabled_mcp_servers() { + for token in csv.split(',') { + let t = token.trim(); + if !t.is_empty() { + out.insert(t.to_string()); + } + } + } + } + out + } + pub fn is_loaded(&self, name: &str) -> bool { self.loaded.contains_key(name) } diff --git a/src/function/mod.rs b/src/function/mod.rs index cc858e4..487038e 100644 --- a/src/function/mod.rs +++ b/src/function/mod.rs @@ -1048,11 +1048,13 @@ impl ToolCall { }) } _ if cmd_name.starts_with(SKILL_FUNCTION_PREFIX) => { - skill::handle_skill_tool(ctx, &cmd_name, &json_data).unwrap_or_else(|e| { - let error_msg = format!("Skill tool failed: {e}"); - eprintln!("{}", warning_text(&format!("⚠️ {error_msg} ⚠️"))); - json!({"tool_call_error": error_msg}) - }) + skill::handle_skill_tool(ctx, &cmd_name, &json_data) + .await + .unwrap_or_else(|e| { + let error_msg = format!("Skill tool failed: {e}"); + eprintln!("{}", warning_text(&format!("⚠️ {error_msg} ⚠️"))); + json!({"tool_call_error": error_msg}) + }) } _ if cmd_name.starts_with(SUPERVISOR_FUNCTION_PREFIX) => { supervisor::handle_supervisor_tool(ctx, &cmd_name, &json_data) diff --git a/src/function/skill.rs b/src/function/skill.rs index 2c62941..ec663ed 100644 --- a/src/function/skill.rs +++ b/src/function/skill.rs @@ -1,5 +1,6 @@ use super::{FunctionDeclaration, JsonSchema}; use crate::config::{RequestContext, Skill, SkillPolicy, paths}; +use crate::utils::create_abort_signal; use anyhow::{Result, bail}; use indexmap::IndexMap; @@ -71,7 +72,7 @@ pub fn skill_function_declarations() -> Vec { ] } -pub fn handle_skill_tool( +pub async fn handle_skill_tool( ctx: &mut RequestContext, cmd_name: &str, args: &Value, @@ -95,8 +96,8 @@ pub fn handle_skill_tool( match action { "list" => handle_list(ctx, &policy), - "load" => handle_load(ctx, args, &policy), - "unload" => handle_unload(ctx, args), + "load" => handle_load(ctx, args, &policy).await, + "unload" => handle_unload(ctx, args).await, _ => bail!("Unknown skill action: {action}"), } } @@ -137,7 +138,7 @@ fn handle_list(ctx: &RequestContext, policy: &SkillPolicy) -> Result { Ok(json!({"skills": entries})) } -fn handle_load( +async fn handle_load( ctx: &mut RequestContext, args: &Value, policy: &SkillPolicy, @@ -189,29 +190,42 @@ fn handle_load( })); } - match ctx.skill_registry.insert(skill) { - Ok(()) => Ok(json!({ - "status": "ok", - "loaded": name, - "message": format!("Skill '{name}' loaded") - })), - Err(e) => Ok(json!({"error": e.to_string()})), + if let Err(e) = ctx.skill_registry.insert(skill) { + return Ok(json!({"error": e.to_string()})); } + + if let Err(e) = ctx.refresh_tool_scope(create_abort_signal()).await { + let _ = ctx.skill_registry.unload(name); + return Ok(json!({ + "error": format!("Loaded skill '{name}' but failed to refresh tool scope: {e}") + })); + } + + Ok(json!({ + "status": "ok", + "loaded": name, + "message": format!("Skill '{name}' loaded") + })) } -fn handle_unload(ctx: &mut RequestContext, args: &Value) -> Result { +async fn handle_unload(ctx: &mut RequestContext, args: &Value) -> Result { let name = match args.get("name").and_then(Value::as_str) { Some(n) if !n.is_empty() => n, _ => return Ok(json!({"error": "name is required"})), }; - match ctx.skill_registry.unload(name) { - Ok(()) => Ok(json!({ - "status": "ok", - "unloaded": name - })), - Err(e) => Ok(json!({"error": e.to_string()})), + if let Err(e) = ctx.skill_registry.unload(name) { + return Ok(json!({"error": e.to_string()})); } + + if let Err(e) = ctx.refresh_tool_scope(create_abort_signal()).await { + warn!("Unloaded skill '{name}' but failed to refresh tool scope: {e}"); + } + + Ok(json!({ + "status": "ok", + "unloaded": name + })) } fn csv_to_vec(csv: Option<&str>) -> Vec {