feat: dynamic loading/unloading of skill tools and MCP servers whenever load_skill/unload_skill are invoked
This commit is contained in:
@@ -1,5 +1,6 @@
|
|||||||
use super::rag_cache::{RagCache, RagKey};
|
use super::rag_cache::{RagCache, RagKey};
|
||||||
use super::session::Session;
|
use super::session::Session;
|
||||||
|
use super::skill_policy::SkillPolicy;
|
||||||
use super::skill_registry::SkillRegistry;
|
use super::skill_registry::SkillRegistry;
|
||||||
use super::todo::TodoList;
|
use super::todo::TodoList;
|
||||||
use super::tool_scope::{McpRuntime, ToolScope};
|
use super::tool_scope::{McpRuntime, ToolScope};
|
||||||
@@ -35,7 +36,7 @@ use indexmap::IndexMap;
|
|||||||
use indoc::formatdoc;
|
use indoc::formatdoc;
|
||||||
use inquire::{Confirm, MultiSelect, Text, list_option::ListOption, validator::Validation};
|
use inquire::{Confirm, MultiSelect, Text, list_option::ListOption, validator::Validation};
|
||||||
use parking_lot::RwLock;
|
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::fs::{File, OpenOptions, read_dir, read_to_string, remove_dir_all, remove_file};
|
||||||
use std::io::Write;
|
use std::io::Write;
|
||||||
use std::path::{Path, PathBuf};
|
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 {
|
pub fn auto_continue_config(&self) -> AutoContinueConfig {
|
||||||
@@ -2067,6 +2068,35 @@ impl RequestContext {
|
|||||||
enabled_mcp_servers: Option<String>,
|
enabled_mcp_servers: Option<String>,
|
||||||
abort_signal: AbortSignal,
|
abort_signal: AbortSignal,
|
||||||
) -> Result<()> {
|
) -> 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<String> = 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::<Vec<_>>().join(","))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
enabled_mcp_servers
|
||||||
|
};
|
||||||
|
|
||||||
let mut mcp_runtime = McpRuntime::new();
|
let mut mcp_runtime = McpRuntime::new();
|
||||||
|
|
||||||
if app.mcp_server_support
|
if app.mcp_server_support
|
||||||
@@ -2134,6 +2164,9 @@ impl RequestContext {
|
|||||||
if !mcp_runtime.is_empty() {
|
if !mcp_runtime.is_empty() {
|
||||||
functions.append_mcp_meta_functions(mcp_runtime.server_names());
|
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();
|
let tool_tracker = self.tool_scope.tool_tracker.clone();
|
||||||
self.tool_scope = ToolScope {
|
self.tool_scope = ToolScope {
|
||||||
@@ -2144,6 +2177,30 @@ impl RequestContext {
|
|||||||
Ok(())
|
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(
|
pub async fn use_role(
|
||||||
&mut self,
|
&mut self,
|
||||||
app: &AppConfig,
|
app: &AppConfig,
|
||||||
|
|||||||
@@ -53,6 +53,21 @@ impl SkillRegistry {
|
|||||||
self.loaded.keys().cloned().collect()
|
self.loaded.keys().cloned().collect()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn loaded_mcp_servers(&self) -> BTreeSet<String> {
|
||||||
|
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 {
|
pub fn is_loaded(&self, name: &str) -> bool {
|
||||||
self.loaded.contains_key(name)
|
self.loaded.contains_key(name)
|
||||||
}
|
}
|
||||||
|
|||||||
+7
-5
@@ -1048,11 +1048,13 @@ impl ToolCall {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
_ if cmd_name.starts_with(SKILL_FUNCTION_PREFIX) => {
|
_ if cmd_name.starts_with(SKILL_FUNCTION_PREFIX) => {
|
||||||
skill::handle_skill_tool(ctx, &cmd_name, &json_data).unwrap_or_else(|e| {
|
skill::handle_skill_tool(ctx, &cmd_name, &json_data)
|
||||||
let error_msg = format!("Skill tool failed: {e}");
|
.await
|
||||||
eprintln!("{}", warning_text(&format!("⚠️ {error_msg} ⚠️")));
|
.unwrap_or_else(|e| {
|
||||||
json!({"tool_call_error": error_msg})
|
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) => {
|
_ if cmd_name.starts_with(SUPERVISOR_FUNCTION_PREFIX) => {
|
||||||
supervisor::handle_supervisor_tool(ctx, &cmd_name, &json_data)
|
supervisor::handle_supervisor_tool(ctx, &cmd_name, &json_data)
|
||||||
|
|||||||
+32
-18
@@ -1,5 +1,6 @@
|
|||||||
use super::{FunctionDeclaration, JsonSchema};
|
use super::{FunctionDeclaration, JsonSchema};
|
||||||
use crate::config::{RequestContext, Skill, SkillPolicy, paths};
|
use crate::config::{RequestContext, Skill, SkillPolicy, paths};
|
||||||
|
use crate::utils::create_abort_signal;
|
||||||
|
|
||||||
use anyhow::{Result, bail};
|
use anyhow::{Result, bail};
|
||||||
use indexmap::IndexMap;
|
use indexmap::IndexMap;
|
||||||
@@ -71,7 +72,7 @@ pub fn skill_function_declarations() -> Vec<FunctionDeclaration> {
|
|||||||
]
|
]
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn handle_skill_tool(
|
pub async fn handle_skill_tool(
|
||||||
ctx: &mut RequestContext,
|
ctx: &mut RequestContext,
|
||||||
cmd_name: &str,
|
cmd_name: &str,
|
||||||
args: &Value,
|
args: &Value,
|
||||||
@@ -95,8 +96,8 @@ pub fn handle_skill_tool(
|
|||||||
|
|
||||||
match action {
|
match action {
|
||||||
"list" => handle_list(ctx, &policy),
|
"list" => handle_list(ctx, &policy),
|
||||||
"load" => handle_load(ctx, args, &policy),
|
"load" => handle_load(ctx, args, &policy).await,
|
||||||
"unload" => handle_unload(ctx, args),
|
"unload" => handle_unload(ctx, args).await,
|
||||||
_ => bail!("Unknown skill action: {action}"),
|
_ => bail!("Unknown skill action: {action}"),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -137,7 +138,7 @@ fn handle_list(ctx: &RequestContext, policy: &SkillPolicy) -> Result<Value> {
|
|||||||
Ok(json!({"skills": entries}))
|
Ok(json!({"skills": entries}))
|
||||||
}
|
}
|
||||||
|
|
||||||
fn handle_load(
|
async fn handle_load(
|
||||||
ctx: &mut RequestContext,
|
ctx: &mut RequestContext,
|
||||||
args: &Value,
|
args: &Value,
|
||||||
policy: &SkillPolicy,
|
policy: &SkillPolicy,
|
||||||
@@ -189,29 +190,42 @@ fn handle_load(
|
|||||||
}));
|
}));
|
||||||
}
|
}
|
||||||
|
|
||||||
match ctx.skill_registry.insert(skill) {
|
if let Err(e) = ctx.skill_registry.insert(skill) {
|
||||||
Ok(()) => Ok(json!({
|
return Ok(json!({"error": e.to_string()}));
|
||||||
"status": "ok",
|
|
||||||
"loaded": name,
|
|
||||||
"message": format!("Skill '{name}' loaded")
|
|
||||||
})),
|
|
||||||
Err(e) => 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<Value> {
|
async fn handle_unload(ctx: &mut RequestContext, args: &Value) -> Result<Value> {
|
||||||
let name = match args.get("name").and_then(Value::as_str) {
|
let name = match args.get("name").and_then(Value::as_str) {
|
||||||
Some(n) if !n.is_empty() => n,
|
Some(n) if !n.is_empty() => n,
|
||||||
_ => return Ok(json!({"error": "name is required"})),
|
_ => return Ok(json!({"error": "name is required"})),
|
||||||
};
|
};
|
||||||
|
|
||||||
match ctx.skill_registry.unload(name) {
|
if let Err(e) = ctx.skill_registry.unload(name) {
|
||||||
Ok(()) => Ok(json!({
|
return Ok(json!({"error": e.to_string()}));
|
||||||
"status": "ok",
|
|
||||||
"unloaded": name
|
|
||||||
})),
|
|
||||||
Err(e) => 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<String> {
|
fn csv_to_vec(csv: Option<&str>) -> Vec<String> {
|
||||||
|
|||||||
Reference in New Issue
Block a user