Dark-Alex-17/loki
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

testing

Browse Source
This commit is contained in:
Alex Clarke
2026-04-16 10:17:03 -06:00
parent ff3419a714
commit d906713d7d
82 changed files with 14886 additions and 3310 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/function/mod.rs
+91 -106
View File
@@ -3,11 +3,12 @@ pub(crate) mod todo;
pub(crate) mod user_interaction;
use crate::{
config::{Agent, Config, GlobalConfig},
config::{Agent, RequestContext},
utils::*,
};
use crate::config::ensure_parent_exists;
use crate::config::paths;
use crate::mcp::{
MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, MCP_INVOKE_META_FUNCTION_NAME_PREFIX,
MCP_SEARCH_META_FUNCTION_NAME_PREFIX,
@@ -110,7 +111,7 @@ fn extract_shebang_runtime(path: &Path) -> Option<String> {
}
pub async fn eval_tool_calls(
config: &GlobalConfig,
ctx: &mut RequestContext,
mut calls: Vec<ToolCall>,
) -> Result<Vec<ToolResult>> {
let mut output = vec![];
@@ -123,9 +124,7 @@ pub async fn eval_tool_calls(
}
let mut is_all_null = true;
for call in calls {
if let Some(checker) = &config.read().tool_call_tracker
&& let Some(msg) = checker.check_loop(&call.clone())
{
if let Some(msg) = ctx.tool_scope.tool_tracker.check_loop(&call.clone()) {
let dup_msg = format!("{{\"tool_call_loop_alert\":{}}}", &msg.trim());
println!(
"{}",
@@ -136,7 +135,7 @@ pub async fn eval_tool_calls(
is_all_null = false;
continue;
}
let mut result = call.eval(config).await?;
let mut result = call.eval(ctx).await?;
if result.is_null() {
result = json!("DONE");
} else {
@@ -149,16 +148,13 @@ pub async fn eval_tool_calls(
}
if !output.is_empty() {
let (has_escalations, summary) = {
let cfg = config.read();
if cfg.current_depth == 0
&& let Some(ref queue) = cfg.root_escalation_queue
&& queue.has_pending()
{
(true, queue.pending_summary())
} else {
(false, vec![])
}
let (has_escalations, summary) = if ctx.current_depth() == 0
&& let Some(queue) = ctx.root_escalation_queue()
&& queue.has_pending()
{
(true, queue.pending_summary())
} else {
(false, vec![])
};
if has_escalations {
@@ -199,7 +195,7 @@ impl Functions {
fn install_global_tools() -> Result<()> {
info!(
"Installing global built-in functions in {}",
Config::functions_dir().display()
paths::functions_dir().display()
);
for file in FunctionAssets::iter() {
@@ -213,7 +209,7 @@ impl Functions {
anyhow!("Failed to load embedded function file: {}", file.as_ref())
})?;
let content = unsafe { std::str::from_utf8_unchecked(&embedded_file.data) };
let file_path = Config::functions_dir().join(file.as_ref());
let file_path = paths::functions_dir().join(file.as_ref());
let file_extension = file_path
.extension()
.and_then(OsStr::to_str)
@@ -254,7 +250,7 @@ impl Functions {
info!(
"Building global function binaries in {}",
Config::functions_bin_dir().display()
paths::functions_bin_dir().display()
);
Self::build_global_function_binaries(visible_tools, None)?;
@@ -271,7 +267,7 @@ impl Functions {
info!(
"Building global function binaries required by agent: {name} in {}",
Config::functions_bin_dir().display()
paths::functions_bin_dir().display()
);
Self::build_global_function_binaries(global_tools, Some(name))?;
tools_declarations
@@ -279,7 +275,7 @@ impl Functions {
debug!("No global tools found for agent: {}", name);
Vec::new()
};
let agent_script_declarations = match Config::agent_functions_file(name) {
let agent_script_declarations = match paths::agent_functions_file(name) {
Ok(path) if path.exists() => {
info!(
"Loading functions script for agent: {name} from {}",
@@ -290,7 +286,7 @@ impl Functions {
info!(
"Building function binary for agent: {name} in {}",
Config::agent_bin_dir(name).display()
paths::agent_bin_dir(name).display()
);
Self::build_agent_tool_binaries(name)?;
script_declarations
@@ -342,14 +338,6 @@ impl Functions {
.extend(user_interaction::user_interaction_function_declarations());
}
pub fn clear_mcp_meta_functions(&mut self) {
self.declarations.retain(|d| {
!d.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
&& !d.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX)
&& !d.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX)
});
}
pub fn append_mcp_meta_functions(&mut self, mcp_servers: Vec<String>) {
let mut invoke_function_properties = IndexMap::new();
invoke_function_properties.insert(
@@ -453,7 +441,7 @@ impl Functions {
fn build_global_tool_declarations(
enabled_tools: &[String],
) -> Result<Vec<FunctionDeclaration>> {
let global_tools_directory = Config::global_tools_dir();
let global_tools_directory = paths::global_tools_dir();
let mut function_declarations = Vec::new();
for tool in enabled_tools {
@@ -542,7 +530,7 @@ impl Functions {
bail!("Unsupported tool file extension: {}", language.as_ref());
}
let tool_path = Config::global_tools_dir().join(tool);
let tool_path = paths::global_tools_dir().join(tool);
let custom_runtime = extract_shebang_runtime(&tool_path);
Self::build_binaries(
binary_name,
@@ -556,7 +544,7 @@ impl Functions {
}
fn clear_agent_bin_dir(name: &str) -> Result<()> {
let agent_bin_directory = Config::agent_bin_dir(name);
let agent_bin_directory = paths::agent_bin_dir(name);
if !agent_bin_directory.exists() {
debug!(
"Creating agent bin directory: {}",
@@ -575,7 +563,7 @@ impl Functions {
}
fn clear_global_functions_bin_dir() -> Result<()> {
let bin_dir = Config::functions_bin_dir();
let bin_dir = paths::functions_bin_dir();
if !bin_dir.exists() {
fs::create_dir_all(&bin_dir)?;
}
@@ -590,7 +578,7 @@ impl Functions {
}
fn build_agent_tool_binaries(name: &str) -> Result<()> {
let tools_file = Config::agent_functions_file(name)?;
let tools_file = paths::agent_functions_file(name)?;
let language = Language::from(
&tools_file
.extension()
@@ -619,18 +607,18 @@ impl Functions {
use native::runtime;
let (binary_file, binary_script_file) = match binary_type {
BinaryType::Tool(None) => (
Config::functions_bin_dir().join(format!("{binary_name}.cmd")),
Config::functions_bin_dir()
paths::functions_bin_dir().join(format!("{binary_name}.cmd")),
paths::functions_bin_dir()
.join(format!("run-{binary_name}.{}", language.to_extension())),
),
BinaryType::Tool(Some(agent_name)) => (
Config::agent_bin_dir(agent_name).join(format!("{binary_name}.cmd")),
Config::agent_bin_dir(agent_name)
paths::agent_bin_dir(agent_name).join(format!("{binary_name}.cmd")),
paths::agent_bin_dir(agent_name)
.join(format!("run-{binary_name}.{}", language.to_extension())),
),
BinaryType::Agent => (
Config::agent_bin_dir(binary_name).join(format!("{binary_name}.cmd")),
Config::agent_bin_dir(binary_name)
paths::agent_bin_dir(binary_name).join(format!("{binary_name}.cmd")),
paths::agent_bin_dir(binary_name)
.join(format!("run-{binary_name}.{}", language.to_extension())),
),
};
@@ -655,7 +643,7 @@ impl Functions {
let to_script_path = |p: &str| -> String { p.replace('\\', "/") };
let content = match binary_type {
BinaryType::Tool(None) => {
let root_dir = Config::functions_dir();
let root_dir = paths::functions_dir();
let tool_path = format!(
"{}/{binary_name}",
&Config::global_tools_dir().to_string_lossy()
@@ -666,7 +654,7 @@ impl Functions {
.replace("{tool_path}", &to_script_path(&tool_path))
}
BinaryType::Tool(Some(agent_name)) => {
let root_dir = Config::agent_data_dir(agent_name);
let root_dir = paths::agent_data_dir(agent_name);
let tool_path = format!(
"{}/{binary_name}",
&Config::global_tools_dir().to_string_lossy()
@@ -680,12 +668,12 @@ impl Functions {
.replace("{agent_name}", binary_name)
.replace(
"{config_dir}",
&to_script_path(&Config::config_dir().to_string_lossy()),
&to_script_path(&paths::config_dir().to_string_lossy()),
),
}
.replace(
"{prompt_utils_file}",
&to_script_path(&Config::bash_prompt_utils_file().to_string_lossy()),
&to_script_path(&paths::bash_prompt_utils_file().to_string_lossy()),
);
if binary_script_file.exists() {
fs::remove_file(&binary_script_file)?;
@@ -769,11 +757,11 @@ impl Functions {
use std::os::unix::prelude::PermissionsExt;
let binary_file = match binary_type {
BinaryType::Tool(None) => Config::functions_bin_dir().join(binary_name),
BinaryType::Tool(None) => paths::functions_bin_dir().join(binary_name),
BinaryType::Tool(Some(agent_name)) => {
Config::agent_bin_dir(agent_name).join(binary_name)
paths::agent_bin_dir(agent_name).join(binary_name)
}
BinaryType::Agent => Config::agent_bin_dir(binary_name).join(binary_name),
BinaryType::Agent => paths::agent_bin_dir(binary_name).join(binary_name),
};
info!(
"Building binary for function: {} ({})",
@@ -795,10 +783,10 @@ impl Functions {
let content_template = unsafe { std::str::from_utf8_unchecked(&embedded_file.data) };
let mut content = match binary_type {
BinaryType::Tool(None) => {
let root_dir = Config::functions_dir();
let root_dir = paths::functions_dir();
let tool_path = format!(
"{}/{binary_name}",
&Config::global_tools_dir().to_string_lossy()
&paths::global_tools_dir().to_string_lossy()
);
content_template
.replace("{function_name}", binary_name)
@@ -806,10 +794,10 @@ impl Functions {
.replace("{tool_path}", &tool_path)
}
BinaryType::Tool(Some(agent_name)) => {
let root_dir = Config::agent_data_dir(agent_name);
let root_dir = paths::agent_data_dir(agent_name);
let tool_path = format!(
"{}/{binary_name}",
&Config::global_tools_dir().to_string_lossy()
&paths::global_tools_dir().to_string_lossy()
);
content_template
.replace("{function_name}", binary_name)
@@ -818,11 +806,11 @@ impl Functions {
}
BinaryType::Agent => content_template
.replace("{agent_name}", binary_name)
.replace("{config_dir}", &Config::config_dir().to_string_lossy()),
.replace("{config_dir}", &paths::config_dir().to_string_lossy()),
}
.replace(
"{prompt_utils_file}",
&Config::bash_prompt_utils_file().to_string_lossy(),
&paths::bash_prompt_utils_file().to_string_lossy(),
);
if let Some(rt) = custom_runtime
@@ -952,16 +940,15 @@ impl ToolCall {
self
}
pub async fn eval(&self, config: &GlobalConfig) -> Result<Value> {
let (call_name, cmd_name, mut cmd_args, envs) = match &config.read().agent {
Some(agent) => self.extract_call_config_from_agent(config, agent)?,
None => self.extract_call_config_from_config(config)?,
pub async fn eval(&self, ctx: &mut RequestContext) -> Result<Value> {
let agent = ctx.agent.clone();
let functions = ctx.tool_scope.functions.clone();
let current_depth = ctx.current_depth();
let agent_name = agent.as_ref().map(|agent| agent.name().to_owned());
let (call_name, cmd_name, mut cmd_args, envs) = match agent.as_ref() {
Some(agent) => self.extract_call_config_from_agent(&functions, agent)?,
None => self.extract_call_config_from_ctx(&functions)?,
};
let agent_name = config
.read()
.agent
.as_ref()
.map(|agent| agent.name().to_owned());
let json_data = if self.arguments.is_object() {
self.arguments.clone()
@@ -981,20 +968,22 @@ impl ToolCall {
let prompt = format!("Call {cmd_name} {}", cmd_args.join(" "));
if *IS_STDOUT_TERMINAL && config.read().current_depth == 0 {
if *IS_STDOUT_TERMINAL && current_depth == 0 {
println!("{}", dimmed_text(&prompt));
}
let output = match cmd_name.as_str() {
_ if cmd_name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX) => {
Self::search_mcp_tools(config, &cmd_name, &json_data).unwrap_or_else(|e| {
let error_msg = format!("MCP search failed: {e}");
eprintln!("{}", warning_text(&format!("⚠️ {error_msg} ⚠️")));
json!({"tool_call_error": error_msg})
})
Self::search_mcp_tools(ctx, &cmd_name, &json_data)
.await
.unwrap_or_else(|e| {
let error_msg = format!("MCP search failed: {e}");
eprintln!("{}", warning_text(&format!("⚠️ {error_msg} ⚠️")));
json!({"tool_call_error": error_msg})
})
}
_ if cmd_name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX) => {
Self::describe_mcp_tool(config, &cmd_name, json_data)
Self::describe_mcp_tool(ctx, &cmd_name, json_data)
.await
.unwrap_or_else(|e| {
let error_msg = format!("MCP describe failed: {e}");
@@ -1003,7 +992,7 @@ impl ToolCall {
})
}
_ if cmd_name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX) => {
Self::invoke_mcp_tool(config, &cmd_name, &json_data)
Self::invoke_mcp_tool(ctx, &cmd_name, &json_data)
.await
.unwrap_or_else(|e| {
let error_msg = format!("MCP tool invocation failed: {e}");
@@ -1012,14 +1001,14 @@ impl ToolCall {
})
}
_ if cmd_name.starts_with(TODO_FUNCTION_PREFIX) => {
todo::handle_todo_tool(config, &cmd_name, &json_data).unwrap_or_else(|e| {
todo::handle_todo_tool(ctx, &cmd_name, &json_data).unwrap_or_else(|e| {
let error_msg = format!("Todo 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(config, &cmd_name, &json_data)
supervisor::handle_supervisor_tool(ctx, &cmd_name, &json_data)
.await
.unwrap_or_else(|e| {
let error_msg = format!("Supervisor tool failed: {e}");
@@ -1028,7 +1017,7 @@ impl ToolCall {
})
}
_ if cmd_name.starts_with(USER_FUNCTION_PREFIX) => {
user_interaction::handle_user_tool(config, &cmd_name, &json_data)
user_interaction::handle_user_tool(ctx, &cmd_name, &json_data)
.await
.unwrap_or_else(|e| {
let error_msg = format!("User interaction failed: {e}");
@@ -1051,7 +1040,7 @@ impl ToolCall {
}
async fn describe_mcp_tool(
config: &GlobalConfig,
ctx: &RequestContext,
cmd_name: &str,
json_data: Value,
) -> Result<Value> {
@@ -1061,18 +1050,19 @@ impl ToolCall {
.ok_or_else(|| anyhow!("Missing 'tool' in arguments"))?
.as_str()
.ok_or_else(|| anyhow!("Invalid 'tool' in arguments"))?;
let registry_arc = {
let cfg = config.read();
cfg.mcp_registry
.clone()
.with_context(|| "MCP is not configured")?
};
let result = registry_arc.describe(&server_id, tool).await?;
let result = ctx
.tool_scope
.mcp_runtime
.describe(&server_id, tool)
.await?;
Ok(serde_json::to_value(result)?)
}
fn search_mcp_tools(config: &GlobalConfig, cmd_name: &str, json_data: &Value) -> Result<Value> {
async fn search_mcp_tools(
ctx: &RequestContext,
cmd_name: &str,
json_data: &Value,
) -> Result<Value> {
let server = cmd_name.replace(&format!("{MCP_SEARCH_META_FUNCTION_NAME_PREFIX}_"), "");
let query = json_data
.get("query")
@@ -1085,15 +1075,12 @@ impl ToolCall {
.unwrap_or_else(|| Value::from(8u64))
.as_u64()
.ok_or_else(|| anyhow!("Invalid 'top_k' in arguments"))? as usize;
let registry_arc = {
let cfg = config.read();
cfg.mcp_registry
.clone()
.with_context(|| "MCP is not configured")?
};
let catalog_items = registry_arc
.search_tools_server(&server, query, top_k)
let catalog_items = ctx
.tool_scope
.mcp_runtime
.search(&server, query, top_k)
.await?
.into_iter()
.map(|it| serde_json::to_value(&it).unwrap_or_default())
.collect();
@@ -1101,7 +1088,7 @@ impl ToolCall {
}
async fn invoke_mcp_tool(
config: &GlobalConfig,
ctx: &RequestContext,
cmd_name: &str,
json_data: &Value,
) -> Result<Value> {
@@ -1115,20 +1102,18 @@ impl ToolCall {
.get("arguments")
.cloned()
.unwrap_or_else(|| json!({}));
let registry_arc = {
let cfg = config.read();
cfg.mcp_registry
.clone()
.with_context(|| "MCP is not configured")?
};
let result = registry_arc.invoke(&server, tool, arguments).await?;
let result = ctx
.tool_scope
.mcp_runtime
.invoke(&server, tool, arguments)
.await?;
Ok(serde_json::to_value(result)?)
}
fn extract_call_config_from_agent(
&self,
config: &GlobalConfig,
functions: &Functions,
agent: &Agent,
) -> Result<CallConfig> {
let function_name = self.name.clone();
@@ -1151,13 +1136,13 @@ impl ToolCall {
))
}
}
None => self.extract_call_config_from_config(config),
None => self.extract_call_config_from_ctx(functions),
}
}
fn extract_call_config_from_config(&self, config: &GlobalConfig) -> Result<CallConfig> {
fn extract_call_config_from_ctx(&self, functions: &Functions) -> Result<CallConfig> {
let function_name = self.name.clone();
match config.read().functions.contains(&function_name) {
match functions.contains(&function_name) {
true => Ok((
function_name.clone(),
function_name,
@@ -1179,12 +1164,12 @@ pub fn run_llm_function(
let mut command_name = cmd_name.clone();
if let Some(agent_name) = agent_name {
command_name = cmd_args[0].clone();
let dir = Config::agent_bin_dir(&agent_name);
let dir = paths::agent_bin_dir(&agent_name);
if dir.exists() {
bin_dirs.push(dir);
}
} else {
bin_dirs.push(Config::functions_bin_dir());
bin_dirs.push(paths::functions_bin_dir());
}
let current_path = env::var("PATH").context("No PATH environment variable")?;
let prepend_path = bin_dirs
src/function/supervisor.rs
+221 -154
View File
@@ -1,12 +1,11 @@
use super::{FunctionDeclaration, JsonSchema};
use crate::client::{Model, ModelType, call_chat_completions};
use crate::config::{Config, GlobalConfig, Input, Role, RoleLike};
use crate::supervisor::escalation::EscalationQueue;
use crate::config::{Agent, AppState, Input, RequestContext, Role, RoleLike};
use crate::supervisor::mailbox::{Envelope, EnvelopePayload, Inbox};
use crate::supervisor::{AgentExitStatus, AgentHandle, AgentResult};
use crate::supervisor::{AgentExitStatus, AgentHandle, AgentResult, Supervisor};
use crate::utils::{AbortSignal, create_abort_signal};
use anyhow::{Result, anyhow, bail};
use anyhow::{Context, Result, anyhow, bail};
use chrono::Utc;
use indexmap::IndexMap;
use log::debug;
@@ -300,7 +299,7 @@ pub fn teammate_function_declarations() -> Vec<FunctionDeclaration> {
}
pub async fn handle_supervisor_tool(
config: &GlobalConfig,
ctx: &mut RequestContext,
cmd_name: &str,
args: &Value,
) -> Result<Value> {
@@ -309,42 +308,47 @@ pub async fn handle_supervisor_tool(
.unwrap_or(cmd_name);
match action {
"spawn" => handle_spawn(config, args).await,
"check" => handle_check(config, args).await,
"collect" => handle_collect(config, args).await,
"list" => handle_list(config),
"cancel" => handle_cancel(config, args),
"send_message" => handle_send_message(config, args),
"check_inbox" => handle_check_inbox(config),
"task_create" => handle_task_create(config, args),
"task_list" => handle_task_list(config),
"task_complete" => handle_task_complete(config, args).await,
"task_fail" => handle_task_fail(config, args),
"reply_escalation" => handle_reply_escalation(config, args),
"spawn" => handle_spawn(ctx, args).await,
"check" => handle_check(ctx, args).await,
"collect" => handle_collect(ctx, args).await,
"list" => handle_list(ctx),
"cancel" => handle_cancel(ctx, args),
"send_message" => handle_send_message(ctx, args),
"check_inbox" => handle_check_inbox(ctx),
"task_create" => handle_task_create(ctx, args),
"task_list" => handle_task_list(ctx),
"task_complete" => handle_task_complete(ctx, args).await,
"task_fail" => handle_task_fail(ctx, args),
"reply_escalation" => handle_reply_escalation(ctx, args),
_ => bail!("Unknown supervisor action: {action}"),
}
}
fn run_child_agent(
child_config: GlobalConfig,
mut child_ctx: RequestContext,
initial_input: Input,
abort_signal: AbortSignal,
) -> Pin<Box<dyn Future<Output = Result<String>> + Send>> {
Box::pin(async move {
let mut accumulated_output = String::new();
let mut input = initial_input;
let app = Arc::clone(&child_ctx.app.config);
loop {
let client = input.create_client()?;
child_config.write().before_chat_completion(&input)?;
child_ctx.before_chat_completion(&input)?;
let (output, tool_results) =
call_chat_completions(&input, false, false, client.as_ref(), abort_signal.clone())
.await?;
let (output, tool_results) = call_chat_completions(
&input,
false,
false,
client.as_ref(),
&mut child_ctx,
abort_signal.clone(),
)
.await?;
child_config
.write()
.after_chat_completion(&input, &output, &tool_results)?;
child_ctx.after_chat_completion(app.as_ref(), &input, &output, &tool_results)?;
if !output.is_empty() {
if !accumulated_output.is_empty() {
@@ -360,7 +364,7 @@ fn run_child_agent(
input = input.merge_tool_results(output, tool_results);
}
if let Some(ref supervisor) = child_config.read().supervisor {
if let Some(supervisor) = child_ctx.supervisor().cloned() {
supervisor.read().cancel_all();
}
@@ -368,7 +372,58 @@ fn run_child_agent(
})
}
async fn handle_spawn(config: &GlobalConfig, args: &Value) -> Result<Value> {
async fn populate_agent_mcp_runtime(ctx: &mut RequestContext, server_ids: &[String]) -> Result<()> {
if !ctx.app.config.mcp_server_support {
return Ok(());
}
let app = Arc::clone(&ctx.app);
let server_specs = app
.mcp_config
.as_ref()
.map(|mcp_config| {
server_ids
.iter()
.filter_map(|id| {
mcp_config
.mcp_servers
.get(id)
.cloned()
.map(|spec| (id.clone(), spec))
})
.collect::<Vec<_>>()
})
.unwrap_or_default();
for (id, spec) in server_specs {
let handle = app
.mcp_factory
.acquire(&id, &spec, app.mcp_log_path.as_deref())
.await?;
ctx.tool_scope.mcp_runtime.insert(id, handle);
}
Ok(())
}
fn sync_agent_functions_to_ctx(ctx: &mut RequestContext) -> Result<()> {
let server_names = ctx.tool_scope.mcp_runtime.server_names();
let functions = {
let agent = ctx
.agent
.as_mut()
.with_context(|| "Agent should be initialized")?;
if !server_names.is_empty() {
agent.append_mcp_meta_functions(server_names);
}
agent.functions().clone()
};
ctx.tool_scope.functions = functions;
Ok(())
}
async fn handle_spawn(ctx: &mut RequestContext, args: &Value) -> Result<Value> {
let agent_name = args
.get("agent")
.and_then(Value::as_str)
@@ -385,10 +440,9 @@ async fn handle_spawn(config: &GlobalConfig, args: &Value) -> Result<Value> {
let agent_id = format!("agent_{agent_name}_{short_uuid}");
let (max_depth, current_depth) = {
let cfg = config.read();
let supervisor = cfg
.supervisor
.as_ref()
let supervisor = ctx
.supervisor()
.cloned()
.ok_or_else(|| anyhow!("No supervisor active; Agent spawning not enabled"))?;
let sup = supervisor.read();
if sup.active_count() >= sup.max_concurrent() {
@@ -401,7 +455,7 @@ async fn handle_spawn(config: &GlobalConfig, args: &Value) -> Result<Value> {
),
}));
}
(sup.max_depth(), cfg.current_depth + 1)
(sup.max_depth(), ctx.current_depth() + 1)
};
if current_depth > max_depth {
@@ -413,37 +467,77 @@ async fn handle_spawn(config: &GlobalConfig, args: &Value) -> Result<Value> {
let child_inbox = Arc::new(Inbox::new());
{
let mut cfg = config.write();
if cfg.root_escalation_queue.is_none() {
cfg.root_escalation_queue = Some(Arc::new(EscalationQueue::new()));
}
}
let child_config: GlobalConfig = {
let mut child_cfg = config.read().clone();
child_cfg.parent_supervisor = child_cfg.supervisor.clone();
child_cfg.agent = None;
child_cfg.session = None;
child_cfg.rag = None;
child_cfg.supervisor = None;
child_cfg.last_message = None;
child_cfg.tool_call_tracker = None;
child_cfg.stream = false;
child_cfg.save = false;
child_cfg.current_depth = current_depth;
child_cfg.inbox = Some(Arc::clone(&child_inbox));
child_cfg.self_agent_id = Some(agent_id.clone());
Arc::new(RwLock::new(child_cfg))
};
ctx.ensure_root_escalation_queue();
let child_abort = create_abort_signal();
Config::use_agent(&child_config, &agent_name, None, child_abort.clone()).await?;
let input = Input::from_str(&child_config, &prompt, None);
if !ctx.app.config.function_calling_support {
bail!("Please enable function calling support before using the agent.");
}
let app_config = Arc::clone(&ctx.app.config);
let current_model = ctx.current_model().clone();
let info_flag = ctx.info_flag;
let child_app_state = Arc::new(AppState {
config: Arc::new(app_config.as_ref().clone()),
vault: ctx.app.vault.clone(),
mcp_factory: Default::default(),
rag_cache: Default::default(),
mcp_config: ctx.app.mcp_config.clone(),
mcp_log_path: ctx.app.mcp_log_path.clone(),
});
let agent = Agent::init(
app_config.as_ref(),
child_app_state.as_ref(),
&current_model,
info_flag,
&agent_name,
child_abort.clone(),
)
.await?;
let agent_mcp_servers = agent.mcp_server_names().to_vec();
let session = agent.agent_session().map(|v| v.to_string());
let should_init_supervisor = agent.can_spawn_agents();
let max_concurrent = agent.max_concurrent_agents();
let max_depth = agent.max_agent_depth();
let mut child_ctx = RequestContext::new_for_child(
Arc::clone(&child_app_state),
ctx,
current_depth,
Arc::clone(&child_inbox),
agent_id.clone(),
);
child_ctx.rag = agent.rag();
child_ctx
.agent_runtime
.as_mut()
.expect("child agent runtime should be initialized")
.rag = child_ctx.rag.clone();
child_ctx.agent = Some(agent);
if should_init_supervisor {
child_ctx
.agent_runtime
.as_mut()
.expect("child agent runtime should be initialized")
.supervisor = Some(Arc::new(RwLock::new(Supervisor::new(
max_concurrent,
max_depth,
))));
}
if let Some(session) = session {
child_ctx
.use_session(app_config.as_ref(), Some(&session), child_abort.clone())
.await?;
sync_agent_functions_to_ctx(&mut child_ctx)?;
} else {
populate_agent_mcp_runtime(&mut child_ctx, &agent_mcp_servers).await?;
sync_agent_functions_to_ctx(&mut child_ctx)?;
child_ctx.init_agent_shared_variables()?;
}
let input = Input::from_str(&child_ctx, &prompt, None);
debug!("Spawning child agent '{agent_name}' as '{agent_id}'");
@@ -452,7 +546,7 @@ async fn handle_spawn(config: &GlobalConfig, args: &Value) -> Result<Value> {
let spawn_abort = child_abort.clone();
let join_handle = tokio::spawn(async move {
let result = run_child_agent(child_config, input, spawn_abort).await;
let result = run_child_agent(child_ctx, input, spawn_abort).await;
match result {
Ok(output) => Ok(AgentResult {
@@ -479,15 +573,12 @@ async fn handle_spawn(config: &GlobalConfig, args: &Value) -> Result<Value> {
join_handle,
};
{
let cfg = config.read();
let supervisor = cfg
.supervisor
.as_ref()
.ok_or_else(|| anyhow!("No supervisor active"))?;
let mut sup = supervisor.write();
sup.register(handle)?;
}
let supervisor = ctx
.supervisor()
.cloned()
.ok_or_else(|| anyhow!("No supervisor active"))?;
let mut sup = supervisor.write();
sup.register(handle)?;
Ok(json!({
"status": "ok",
@@ -497,24 +588,23 @@ async fn handle_spawn(config: &GlobalConfig, args: &Value) -> Result<Value> {
}))
}
async fn handle_check(config: &GlobalConfig, args: &Value) -> Result<Value> {
async fn handle_check(ctx: &mut RequestContext, args: &Value) -> Result<Value> {
let id = args
.get("id")
.and_then(Value::as_str)
.ok_or_else(|| anyhow!("'id' is required"))?;
let is_finished = {
let cfg = config.read();
let supervisor = cfg
.supervisor
.as_ref()
let supervisor = ctx
.supervisor()
.cloned()
.ok_or_else(|| anyhow!("No supervisor active"))?;
let sup = supervisor.read();
sup.is_finished(id)
};
match is_finished {
Some(true) => handle_collect(config, args).await,
Some(true) => handle_collect(ctx, args).await,
Some(false) => Ok(json!({
"status": "pending",
"id": id,
@@ -527,17 +617,16 @@ async fn handle_check(config: &GlobalConfig, args: &Value) -> Result<Value> {
}
}
async fn handle_collect(config: &GlobalConfig, args: &Value) -> Result<Value> {
async fn handle_collect(ctx: &mut RequestContext, args: &Value) -> Result<Value> {
let id = args
.get("id")
.and_then(Value::as_str)
.ok_or_else(|| anyhow!("'id' is required"))?;
let handle = {
let cfg = config.read();
let supervisor = cfg
.supervisor
.as_ref()
let supervisor = ctx
.supervisor()
.cloned()
.ok_or_else(|| anyhow!("No supervisor active"))?;
let mut sup = supervisor.write();
sup.take(id)
@@ -551,7 +640,7 @@ async fn handle_collect(config: &GlobalConfig, args: &Value) -> Result<Value> {
.map_err(|e| anyhow!("Agent task panicked: {e}"))?
.map_err(|e| anyhow!("Agent failed: {e}"))?;
let output = summarize_output(config, &result.agent_name, &result.output).await?;
let output = summarize_output(ctx, &result.agent_name, &result.output).await?;
Ok(json!({
"status": "completed",
@@ -568,11 +657,10 @@ async fn handle_collect(config: &GlobalConfig, args: &Value) -> Result<Value> {
}
}
fn handle_list(config: &GlobalConfig) -> Result<Value> {
let cfg = config.read();
let supervisor = cfg
.supervisor
.as_ref()
fn handle_list(ctx: &mut RequestContext) -> Result<Value> {
let supervisor = ctx
.supervisor()
.cloned()
.ok_or_else(|| anyhow!("No supervisor active"))?;
let sup = supervisor.read();
@@ -596,16 +684,15 @@ fn handle_list(config: &GlobalConfig) -> Result<Value> {
}))
}
fn handle_cancel(config: &GlobalConfig, args: &Value) -> Result<Value> {
fn handle_cancel(ctx: &mut RequestContext, args: &Value) -> Result<Value> {
let id = args
.get("id")
.and_then(Value::as_str)
.ok_or_else(|| anyhow!("'id' is required"))?;
let cfg = config.read();
let supervisor = cfg
.supervisor
.as_ref()
let supervisor = ctx
.supervisor()
.cloned()
.ok_or_else(|| anyhow!("No supervisor active"))?;
let mut sup = supervisor.write();
@@ -624,7 +711,7 @@ fn handle_cancel(config: &GlobalConfig, args: &Value) -> Result<Value> {
}
}
fn handle_send_message(config: &GlobalConfig, args: &Value) -> Result<Value> {
fn handle_send_message(ctx: &mut RequestContext, args: &Value) -> Result<Value> {
let id = args
.get("id")
.and_then(Value::as_str)
@@ -634,25 +721,18 @@ fn handle_send_message(config: &GlobalConfig, args: &Value) -> Result<Value> {
.and_then(Value::as_str)
.ok_or_else(|| anyhow!("'message' is required"))?;
let cfg = config.read();
// Determine sender identity: self_agent_id (child), agent name (parent), or "parent"
let sender = cfg
.self_agent_id
.clone()
.or_else(|| cfg.agent.as_ref().map(|a| a.name().to_string()))
let sender = ctx
.self_agent_id()
.map(str::to_owned)
.or_else(|| ctx.agent.as_ref().map(|a| a.name().to_string()))
.unwrap_or_else(|| "parent".to_string());
// Try local supervisor first (parent -> child routing)
let inbox = cfg
.supervisor
.as_ref()
let inbox = ctx
.supervisor()
.and_then(|sup| sup.read().inbox(id).cloned());
// Fall back to parent_supervisor (sibling -> sibling routing)
let inbox = inbox.or_else(|| {
cfg.parent_supervisor
.as_ref()
ctx.parent_supervisor()
.and_then(|sup| sup.read().inbox(id).cloned())
});
@@ -679,9 +759,8 @@ fn handle_send_message(config: &GlobalConfig, args: &Value) -> Result<Value> {
}
}
fn handle_check_inbox(config: &GlobalConfig) -> Result<Value> {
let cfg = config.read();
match &cfg.inbox {
fn handle_check_inbox(ctx: &mut RequestContext) -> Result<Value> {
match ctx.inbox() {
Some(inbox) => {
let messages: Vec<Value> = inbox
.drain()
@@ -707,7 +786,7 @@ fn handle_check_inbox(config: &GlobalConfig) -> Result<Value> {
}
}
fn handle_reply_escalation(config: &GlobalConfig, args: &Value) -> Result<Value> {
fn handle_reply_escalation(ctx: &mut RequestContext, args: &Value) -> Result<Value> {
let escalation_id = args
.get("escalation_id")
.and_then(Value::as_str)
@@ -717,12 +796,10 @@ fn handle_reply_escalation(config: &GlobalConfig, args: &Value) -> Result<Value>
.and_then(Value::as_str)
.ok_or_else(|| anyhow!("'reply' is required"))?;
let queue = {
let cfg = config.read();
cfg.root_escalation_queue
.clone()
.ok_or_else(|| anyhow!("No escalation queue available"))?
};
let queue = ctx
.root_escalation_queue()
.cloned()
.ok_or_else(|| anyhow!("No escalation queue available"))?;
match queue.take(escalation_id) {
Some(request) => {
@@ -742,7 +819,7 @@ fn handle_reply_escalation(config: &GlobalConfig, args: &Value) -> Result<Value>
}
}
fn handle_task_create(config: &GlobalConfig, args: &Value) -> Result<Value> {
fn handle_task_create(ctx: &mut RequestContext, args: &Value) -> Result<Value> {
let subject = args
.get("subject")
.and_then(Value::as_str)
@@ -768,10 +845,9 @@ fn handle_task_create(config: &GlobalConfig, args: &Value) -> Result<Value> {
bail!("'prompt' is required when 'agent' is set");
}
let cfg = config.read();
let supervisor = cfg
.supervisor
.as_ref()
let supervisor = ctx
.supervisor()
.cloned()
.ok_or_else(|| anyhow!("No supervisor active"))?;
let mut sup = supervisor.write();
@@ -805,11 +881,10 @@ fn handle_task_create(config: &GlobalConfig, args: &Value) -> Result<Value> {
Ok(result)
}
fn handle_task_list(config: &GlobalConfig) -> Result<Value> {
let cfg = config.read();
let supervisor = cfg
.supervisor
.as_ref()
fn handle_task_list(ctx: &mut RequestContext) -> Result<Value> {
let supervisor = ctx
.supervisor()
.cloned()
.ok_or_else(|| anyhow!("No supervisor active"))?;
let sup = supervisor.read();
@@ -834,17 +909,16 @@ fn handle_task_list(config: &GlobalConfig) -> Result<Value> {
Ok(json!({ "tasks": tasks }))
}
async fn handle_task_complete(config: &GlobalConfig, args: &Value) -> Result<Value> {
async fn handle_task_complete(ctx: &mut RequestContext, args: &Value) -> Result<Value> {
let task_id = args
.get("task_id")
.and_then(Value::as_str)
.ok_or_else(|| anyhow!("'task_id' is required"))?;
let (newly_runnable, dispatchable) = {
let cfg = config.read();
let supervisor = cfg
.supervisor
.as_ref()
let supervisor = ctx
.supervisor()
.cloned()
.ok_or_else(|| anyhow!("No supervisor active"))?;
let mut sup = supervisor.write();
@@ -884,7 +958,7 @@ async fn handle_task_complete(config: &GlobalConfig, args: &Value) -> Result<Val
"agent": agent,
"prompt": prompt,
});
match handle_spawn(config, &spawn_args).await {
match handle_spawn(ctx, &spawn_args).await {
Ok(result) => {
let agent_id = result
.get("id")
@@ -916,16 +990,15 @@ async fn handle_task_complete(config: &GlobalConfig, args: &Value) -> Result<Val
Ok(result)
}
fn handle_task_fail(config: &GlobalConfig, args: &Value) -> Result<Value> {
fn handle_task_fail(ctx: &mut RequestContext, args: &Value) -> Result<Value> {
let task_id = args
.get("task_id")
.and_then(Value::as_str)
.ok_or_else(|| anyhow!("'task_id' is required"))?;
let cfg = config.read();
let supervisor = cfg
.supervisor
.as_ref()
let supervisor = ctx
.supervisor()
.cloned()
.ok_or_else(|| anyhow!("No supervisor active"))?;
let mut sup = supervisor.write();
@@ -958,17 +1031,12 @@ Rules:
- Use bullet points for multiple findings
- If the output contains a final answer or conclusion, lead with it"#;
async fn summarize_output(config: &GlobalConfig, agent_name: &str, output: &str) -> Result<String> {
let (threshold, summarization_model_id) = {
let cfg = config.read();
match cfg.agent.as_ref() {
Some(agent) => (
agent.summarization_threshold(),
agent.summarization_model().map(|s| s.to_string()),
),
None => return Ok(output.to_string()),
}
async fn summarize_output(ctx: &RequestContext, agent_name: &str, output: &str) -> Result<String> {
let Some(agent) = ctx.agent.as_ref() else {
return Ok(output.to_string());
};
let threshold = agent.summarization_threshold();
let summarization_model_id = agent.summarization_model().map(|s| s.to_string());
if output.len() < threshold {
debug!(
@@ -987,12 +1055,11 @@ async fn summarize_output(config: &GlobalConfig, agent_name: &str, output: &str)
threshold
);
let model = {
let cfg = config.read();
match summarization_model_id {
Some(ref model_id) => Model::retrieve_model(&cfg, model_id, ModelType::Chat)?,
None => cfg.current_model().clone(),
let model = match summarization_model_id {
Some(ref model_id) => {
Model::retrieve_model(ctx.app.config.as_ref(), model_id, ModelType::Chat)?
}
None => ctx.current_model().clone(),
};
let mut role = Role::new("summarizer", SUMMARIZATION_PROMPT);
@@ -1002,7 +1069,7 @@ async fn summarize_output(config: &GlobalConfig, agent_name: &str, output: &str)
"Summarize the following sub-agent output from '{}':\n\n{}",
agent_name, output
);
let input = Input::from_str(config, &user_message, Some(role));
let input = Input::from_str(ctx, &user_message, Some(role));
let summary = input.fetch_chat_text().await?;
src/function/todo.rs
+7 -12
View File
@@ -1,5 +1,5 @@
use super::{FunctionDeclaration, JsonSchema};
use crate::config::GlobalConfig;
use crate::config::RequestContext;
use anyhow::{Result, bail};
use indexmap::IndexMap;
@@ -89,7 +89,7 @@ pub fn todo_function_declarations() -> Vec<FunctionDeclaration> {
]
}
pub fn handle_todo_tool(config: &GlobalConfig, cmd_name: &str, args: &Value) -> Result<Value> {
pub fn handle_todo_tool(ctx: &mut RequestContext, cmd_name: &str, args: &Value) -> Result<Value> {
let action = cmd_name
.strip_prefix(TODO_FUNCTION_PREFIX)
.unwrap_or(cmd_name);
@@ -97,8 +97,7 @@ pub fn handle_todo_tool(config: &GlobalConfig, cmd_name: &str, args: &Value) ->
match action {
"init" => {
let goal = args.get("goal").and_then(Value::as_str).unwrap_or_default();
let mut cfg = config.write();
let agent = cfg.agent.as_mut();
let agent = ctx.agent.as_mut();
match agent {
Some(agent) => {
agent.init_todo_list(goal);
@@ -112,8 +111,7 @@ pub fn handle_todo_tool(config: &GlobalConfig, cmd_name: &str, args: &Value) ->
if task.is_empty() {
return Ok(json!({"error": "task description is required"}));
}
let mut cfg = config.write();
let agent = cfg.agent.as_mut();
let agent = ctx.agent.as_mut();
match agent {
Some(agent) => {
let id = agent.add_todo(task);
@@ -132,8 +130,7 @@ pub fn handle_todo_tool(config: &GlobalConfig, cmd_name: &str, args: &Value) ->
.map(|v| v as usize);
match id {
Some(id) => {
let mut cfg = config.write();
let agent = cfg.agent.as_mut();
let agent = ctx.agent.as_mut();
match agent {
Some(agent) => {
if agent.mark_todo_done(id) {
@@ -151,8 +148,7 @@ pub fn handle_todo_tool(config: &GlobalConfig, cmd_name: &str, args: &Value) ->
}
}
"list" => {
let cfg = config.read();
let agent = cfg.agent.as_ref();
let agent = ctx.agent.as_ref();
match agent {
Some(agent) => {
let list = agent.todo_list();
@@ -167,8 +163,7 @@ pub fn handle_todo_tool(config: &GlobalConfig, cmd_name: &str, args: &Value) ->
}
}
"clear" => {
let mut cfg = config.write();
let agent = cfg.agent.as_mut();
let agent = ctx.agent.as_mut();
match agent {
Some(agent) => {
agent.clear_todo_list();
src/function/user_interaction.rs
+23 -27
View File
@@ -1,5 +1,5 @@
use super::{FunctionDeclaration, JsonSchema};
use crate::config::GlobalConfig;
use crate::config::RequestContext;
use crate::supervisor::escalation::{EscalationRequest, new_escalation_id};
use anyhow::{Result, anyhow};
@@ -120,7 +120,7 @@ pub fn user_interaction_function_declarations() -> Vec<FunctionDeclaration> {
}
pub async fn handle_user_tool(
config: &GlobalConfig,
ctx: &mut RequestContext,
cmd_name: &str,
args: &Value,
) -> Result<Value> {
@@ -128,12 +128,12 @@ pub async fn handle_user_tool(
.strip_prefix(USER_FUNCTION_PREFIX)
.unwrap_or(cmd_name);
let depth = config.read().current_depth;
let depth = ctx.current_depth();
if depth == 0 {
handle_direct(action, args)
} else {
handle_escalated(config, action, args).await
handle_escalated(ctx, action, args).await
}
}
@@ -198,7 +198,7 @@ fn handle_direct_checkbox(args: &Value) -> Result<Value> {
Ok(json!({ "answers": answers }))
}
async fn handle_escalated(config: &GlobalConfig, action: &str, args: &Value) -> Result<Value> {
async fn handle_escalated(ctx: &RequestContext, action: &str, args: &Value) -> Result<Value> {
let question = args
.get("question")
.and_then(Value::as_str)
@@ -212,28 +212,24 @@ async fn handle_escalated(config: &GlobalConfig, action: &str, args: &Value) ->
.collect()
});
let (from_agent_id, from_agent_name, root_queue, timeout_secs) = {
let cfg = config.read();
let agent_id = cfg
.self_agent_id
.clone()
.unwrap_or_else(|| "unknown".to_string());
let agent_name = cfg
.agent
.as_ref()
.map(|a| a.name().to_string())
.unwrap_or_else(|| "unknown".to_string());
let queue = cfg
.root_escalation_queue
.clone()
.ok_or_else(|| anyhow!("No escalation queue available; cannot reach parent agent"))?;
let timeout = cfg
.agent
.as_ref()
.map(|a| a.escalation_timeout())
.unwrap_or(DEFAULT_ESCALATION_TIMEOUT_SECS);
(agent_id, agent_name, queue, timeout)
};
let from_agent_id = ctx
.self_agent_id()
.map(ToOwned::to_owned)
.unwrap_or_else(|| "unknown".to_string());
let from_agent_name = ctx
.agent
.as_ref()
.map(|a| a.name().to_string())
.unwrap_or_else(|| "unknown".to_string());
let root_queue = ctx
.root_escalation_queue()
.cloned()
.ok_or_else(|| anyhow!("No escalation queue available; cannot reach parent agent"))?;
let timeout_secs = ctx
.agent
.as_ref()
.map(|a| a.escalation_timeout())
.unwrap_or(DEFAULT_ESCALATION_TIMEOUT_SECS);
let escalation_id = new_escalation_id();
let (tx, rx) = oneshot::channel();