feat: generalize supervisor registry to TaskHandle enum with job scaffolding, kill discipline, and max_concurrent_jobs config
Implements T1 of plans/background-jobs-design.md (§6, R7/R8/R9): - Supervisor.handles is now HashMap<String, TaskHandle> where TaskHandle = Agent(AgentHandle) | Job(JobHandle); agent-facing accessors (active_count, effective_active_count, is_finished, take, inbox, abort_signal_for, list_agents) match only Agent variants, preserving all existing external behavior byte-for-byte. - New JobHandle/JobState/JobStatus/JobResult types with pgid-guarded process-group kill discipline: Drop and cancel_all/cancel_recursive kill the group only while state.pgid is still set (pid-reuse guard), via libc::killpg on unix and JoinHandle::abort elsewhere. - Per-kind job capacity: Supervisor carries max_concurrent_jobs (builder-set, default 0); job registration rejects at capacity. - Cross-kind teaching errors at the four agent-lookup miss sites (agent__check/collect/cancel/send_message) when the id is a registered job or job_-prefixed; genuinely-unknown ids keep their existing messages. - Supervisor init condition is now can_spawn_agents || jobs_enabled in use_agent and both child-agent spawn paths, with agent capacity 0 in jobs-only contexts; use_agent cancels the old supervisor recursively before replacing it. - max_concurrent_jobs config plumbing: global Config field, AgentConfig override + accessor, all four AppConfig touch points including the COYOTE_MAX_CONCURRENT_JOBS env override; shared effective_max_concurrent_jobs/jobs_enabled predicates (agent override -> global -> default 5; 0 disables). - Stage dependency-free RingBuf (64 KiB default) in src/function/jobs.rs for the upcoming job output pump. - New sanctioned dependency: libc 0.2 under cfg(unix).
This commit is contained in:
Generated
+1
@@ -1703,6 +1703,7 @@ dependencies = [
|
||||
"inquire",
|
||||
"is-terminal",
|
||||
"json-patch",
|
||||
"libc",
|
||||
"log",
|
||||
"log4rs",
|
||||
"nu-ansi-term",
|
||||
|
||||
@@ -139,6 +139,9 @@ arboard = { version = "3.3.0", default-features = false, features = [
|
||||
[target.'cfg(not(any(target_os = "linux", target_os = "android", target_os = "emscripten")))'.dependencies]
|
||||
arboard = { version = "3.3.0", default-features = false }
|
||||
|
||||
[target.'cfg(unix)'.dependencies]
|
||||
libc = "0.2"
|
||||
|
||||
[dev-dependencies]
|
||||
pretty_assertions = "1.4.0"
|
||||
rmcp = { version = "3.1.2", features = ["server"] }
|
||||
|
||||
@@ -561,6 +561,10 @@ impl Agent {
|
||||
self.config.max_tool_result_chars
|
||||
}
|
||||
|
||||
pub fn max_concurrent_jobs(&self) -> Option<usize> {
|
||||
self.config.max_concurrent_jobs
|
||||
}
|
||||
|
||||
pub fn compression_keep_last(&self) -> Option<usize> {
|
||||
self.config.compression_keep_last
|
||||
}
|
||||
@@ -735,6 +739,8 @@ pub struct AgentConfig {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub max_tool_result_chars: Option<usize>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub max_concurrent_jobs: Option<usize>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub compression_keep_last: Option<usize>,
|
||||
#[serde(default)]
|
||||
pub description: String,
|
||||
|
||||
@@ -68,6 +68,7 @@ pub struct AppConfig {
|
||||
pub summarization_prompt: Option<String>,
|
||||
pub summary_context_prompt: Option<String>,
|
||||
pub max_tool_result_chars: Option<usize>,
|
||||
pub max_concurrent_jobs: Option<usize>,
|
||||
|
||||
pub memory: Option<bool>,
|
||||
pub memory_cap_with_tools: Option<usize>,
|
||||
@@ -153,6 +154,7 @@ impl Default for AppConfig {
|
||||
summarization_prompt: None,
|
||||
summary_context_prompt: None,
|
||||
max_tool_result_chars: None,
|
||||
max_concurrent_jobs: None,
|
||||
|
||||
memory: None,
|
||||
memory_cap_with_tools: None,
|
||||
@@ -239,6 +241,7 @@ impl AppConfig {
|
||||
summarization_prompt: config.summarization_prompt,
|
||||
summary_context_prompt: config.summary_context_prompt,
|
||||
max_tool_result_chars: config.max_tool_result_chars,
|
||||
max_concurrent_jobs: config.max_concurrent_jobs,
|
||||
|
||||
memory: config.memory,
|
||||
memory_cap_with_tools: config.memory_cap_with_tools,
|
||||
@@ -574,6 +577,9 @@ impl AppConfig {
|
||||
{
|
||||
self.compression_threshold = v;
|
||||
}
|
||||
if let Some(v) = super::read_env_value::<usize>(&get_env_name("max_concurrent_jobs")) {
|
||||
self.max_concurrent_jobs = v;
|
||||
}
|
||||
if let Some(v) = super::read_env_value::<String>(&get_env_name("summarization_prompt")) {
|
||||
self.summarization_prompt = v;
|
||||
}
|
||||
@@ -844,6 +850,49 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_config_copies_max_concurrent_jobs() {
|
||||
let cfg = Config {
|
||||
model_id: "test-model".to_string(),
|
||||
max_concurrent_jobs: Some(3),
|
||||
clients: vec![ClientConfig::default()],
|
||||
..Config::default()
|
||||
};
|
||||
|
||||
let app = AppConfig::from_config(cfg).unwrap();
|
||||
|
||||
assert_eq!(app.max_concurrent_jobs, Some(3));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial_test::serial]
|
||||
fn load_envs_overrides_max_concurrent_jobs() {
|
||||
let env_name = get_env_name("max_concurrent_jobs");
|
||||
let prev = std::env::var_os(&env_name);
|
||||
|
||||
let mut app = AppConfig::default();
|
||||
|
||||
unsafe { std::env::set_var(&env_name, "7") };
|
||||
app.load_envs();
|
||||
assert_eq!(app.max_concurrent_jobs, Some(7));
|
||||
|
||||
unsafe { std::env::set_var(&env_name, "0") };
|
||||
app.load_envs();
|
||||
assert_eq!(app.max_concurrent_jobs, Some(0));
|
||||
|
||||
unsafe { std::env::remove_var(&env_name) };
|
||||
app.max_concurrent_jobs = Some(2);
|
||||
app.load_envs();
|
||||
assert_eq!(app.max_concurrent_jobs, Some(2));
|
||||
|
||||
unsafe {
|
||||
match prev {
|
||||
Some(v) => std::env::set_var(&env_name, v),
|
||||
None => std::env::remove_var(&env_name),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn editor_returns_configured_value() {
|
||||
let configured = cached_editor()
|
||||
|
||||
+6
-1
@@ -42,7 +42,10 @@ pub use self::macro_policy::{
|
||||
MacroAllowlistLevel, MacroPolicy, MacroSource, MacroState, RESERVED_MACRO_NAMES, ResolvedMacro,
|
||||
};
|
||||
#[allow(unused_imports)]
|
||||
pub use self::request_context::{RenderMode, RequestContext, should_inject_skill_instructions};
|
||||
pub use self::request_context::{
|
||||
RenderMode, RequestContext, effective_max_concurrent_jobs, jobs_enabled,
|
||||
should_inject_skill_instructions,
|
||||
};
|
||||
pub use self::role::{
|
||||
CODE_ROLE, CREATE_TITLE_ROLE, EXPLAIN_SHELL_ROLE, Role, RoleLike, SHELL_ROLE,
|
||||
};
|
||||
@@ -264,6 +267,7 @@ pub struct Config {
|
||||
pub summarization_prompt: Option<String>,
|
||||
pub summary_context_prompt: Option<String>,
|
||||
pub max_tool_result_chars: Option<usize>,
|
||||
pub max_concurrent_jobs: Option<usize>,
|
||||
|
||||
pub memory: Option<bool>,
|
||||
pub memory_cap_with_tools: Option<usize>,
|
||||
@@ -346,6 +350,7 @@ impl Default for Config {
|
||||
summarization_prompt: None,
|
||||
summary_context_prompt: None,
|
||||
max_tool_result_chars: None,
|
||||
max_concurrent_jobs: None,
|
||||
|
||||
memory: None,
|
||||
memory_cap_with_tools: None,
|
||||
|
||||
@@ -138,6 +138,17 @@ pub fn should_inject_skill_instructions(app: &AppConfig, policy: &SkillPolicy) -
|
||||
app.function_calling_support && policy.skills_enabled && !policy.compatible_enabled.is_empty()
|
||||
}
|
||||
|
||||
pub fn effective_max_concurrent_jobs(agent: Option<&Agent>, app: &AppConfig) -> usize {
|
||||
agent
|
||||
.and_then(|a| a.max_concurrent_jobs())
|
||||
.or(app.max_concurrent_jobs)
|
||||
.unwrap_or(5)
|
||||
}
|
||||
|
||||
pub fn jobs_enabled(agent: Option<&Agent>, app: &AppConfig) -> bool {
|
||||
app.function_calling_support && effective_max_concurrent_jobs(agent, app) > 0
|
||||
}
|
||||
|
||||
fn print_asset_names(kind: &str, names: &[String]) -> Result<()> {
|
||||
if names.is_empty() {
|
||||
println!("No {kind} found.");
|
||||
@@ -4046,6 +4057,11 @@ impl RequestContext {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn jobs_enabled(&self) -> bool {
|
||||
jobs_enabled(self.agent.as_ref(), &self.app.config)
|
||||
}
|
||||
|
||||
pub async fn use_agent(
|
||||
&mut self,
|
||||
app: &AppConfig,
|
||||
@@ -4136,11 +4152,20 @@ impl RequestContext {
|
||||
);
|
||||
}
|
||||
|
||||
let should_init_supervisor = agent.can_spawn_agents();
|
||||
let max_concurrent = agent.max_concurrent_agents();
|
||||
let jobs_enabled = jobs_enabled(Some(&agent), app);
|
||||
let should_init_supervisor = agent.can_spawn_agents() || jobs_enabled;
|
||||
let max_concurrent = if agent.can_spawn_agents() {
|
||||
agent.max_concurrent_agents()
|
||||
} else {
|
||||
0
|
||||
};
|
||||
let max_depth = agent.max_agent_depth();
|
||||
let supervisor = should_init_supervisor
|
||||
.then(|| Arc::new(RwLock::new(Supervisor::new(max_concurrent, max_depth))));
|
||||
let max_jobs = effective_max_concurrent_jobs(Some(&agent), app);
|
||||
let supervisor = should_init_supervisor.then(|| {
|
||||
Arc::new(RwLock::new(
|
||||
Supervisor::new(max_concurrent, max_depth).with_max_concurrent_jobs(max_jobs),
|
||||
))
|
||||
});
|
||||
|
||||
self.rag = agent.rag();
|
||||
// Keep `rag_key` in lockstep with `rag`. Agent RAGs are cached under
|
||||
@@ -4152,6 +4177,9 @@ impl RequestContext {
|
||||
.is_some()
|
||||
.then(|| RagKey::Agent(agent.name().to_string()));
|
||||
self.agent = Some(agent);
|
||||
if let Some(old) = self.supervisor.as_ref() {
|
||||
old.read().cancel_recursive();
|
||||
}
|
||||
self.supervisor = supervisor;
|
||||
self.inbox = None;
|
||||
self.escalation_queue = None;
|
||||
@@ -5217,6 +5245,171 @@ mod tests {
|
||||
assert_eq!(ctx.rag_key, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn effective_max_concurrent_jobs_resolution_precedence() {
|
||||
let mut app = AppConfig::default();
|
||||
assert_eq!(effective_max_concurrent_jobs(None, &app), 5);
|
||||
|
||||
app.max_concurrent_jobs = Some(9);
|
||||
assert_eq!(effective_max_concurrent_jobs(None, &app), 9);
|
||||
|
||||
let agent = Agent::test_new(AgentConfig {
|
||||
max_concurrent_jobs: Some(2),
|
||||
..AgentConfig::default()
|
||||
});
|
||||
assert_eq!(effective_max_concurrent_jobs(Some(&agent), &app), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn jobs_enabled_requires_function_calling_and_nonzero_capacity() {
|
||||
let mut app = AppConfig::default();
|
||||
assert!(jobs_enabled(None, &app));
|
||||
|
||||
app.max_concurrent_jobs = Some(0);
|
||||
assert!(!jobs_enabled(None, &app));
|
||||
|
||||
app.max_concurrent_jobs = None;
|
||||
app.function_calling_support = false;
|
||||
assert!(!jobs_enabled(None, &app));
|
||||
|
||||
app.function_calling_support = true;
|
||||
let agent = Agent::test_new(AgentConfig {
|
||||
max_concurrent_jobs: Some(0),
|
||||
..AgentConfig::default()
|
||||
});
|
||||
assert!(!jobs_enabled(Some(&agent), &app));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial]
|
||||
fn use_agent_cancels_previous_supervisor() {
|
||||
let _guard = TestConfigDirGuard::new();
|
||||
let mut ctx = create_test_ctx();
|
||||
let app = ctx.app.config.clone();
|
||||
let agent_name = format!(
|
||||
"test_agent_{}",
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_nanos()
|
||||
);
|
||||
let agent_dir = paths::agent_data_dir(&agent_name);
|
||||
create_dir_all(&agent_dir).unwrap();
|
||||
write(
|
||||
agent_dir.join("config.yaml"),
|
||||
format!("name: {agent_name}\ninstructions: hi\n"),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let old_sig = utils::create_abort_signal();
|
||||
|
||||
tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.unwrap()
|
||||
.block_on(async {
|
||||
let join_handle = tokio::spawn(async {
|
||||
Ok(crate::supervisor::AgentResult {
|
||||
id: "a1".into(),
|
||||
agent_name: "explore".into(),
|
||||
output: String::new(),
|
||||
exit_status: crate::supervisor::AgentExitStatus::Completed,
|
||||
})
|
||||
});
|
||||
let handle = crate::supervisor::AgentHandle {
|
||||
id: "a1".to_string(),
|
||||
agent_name: "explore".to_string(),
|
||||
depth: 1,
|
||||
inbox: Arc::new(Inbox::new()),
|
||||
abort_signal: old_sig.clone(),
|
||||
join_handle,
|
||||
child_supervisor: None,
|
||||
};
|
||||
let old_sup = Arc::new(RwLock::new(Supervisor::new(4, 3)));
|
||||
old_sup.write().register(handle).unwrap();
|
||||
ctx.supervisor = Some(old_sup);
|
||||
|
||||
ctx.use_agent(&app, &agent_name, None, utils::create_abort_signal())
|
||||
.await
|
||||
.unwrap();
|
||||
});
|
||||
|
||||
assert!(old_sig.aborted());
|
||||
assert!(ctx.supervisor.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial]
|
||||
fn use_agent_inits_job_capable_supervisor_without_spawning() {
|
||||
let _guard = TestConfigDirGuard::new();
|
||||
let mut ctx = create_test_ctx();
|
||||
let app = ctx.app.config.clone();
|
||||
let agent_name = format!(
|
||||
"test_agent_{}",
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_nanos()
|
||||
);
|
||||
let agent_dir = paths::agent_data_dir(&agent_name);
|
||||
create_dir_all(&agent_dir).unwrap();
|
||||
write(
|
||||
agent_dir.join("config.yaml"),
|
||||
format!("name: {agent_name}\ninstructions: hi\n"),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.unwrap()
|
||||
.block_on(async {
|
||||
ctx.use_agent(&app, &agent_name, None, utils::create_abort_signal())
|
||||
.await
|
||||
.unwrap();
|
||||
});
|
||||
|
||||
let supervisor = ctx.supervisor.as_ref().expect("supervisor for jobs");
|
||||
let supervisor = supervisor.read();
|
||||
assert_eq!(supervisor.max_concurrent(), 0);
|
||||
assert_eq!(supervisor.max_concurrent_jobs(), 5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial]
|
||||
fn use_agent_skips_supervisor_when_jobs_disabled() {
|
||||
let _guard = TestConfigDirGuard::new();
|
||||
let mut ctx = create_test_ctx();
|
||||
let mut app = ctx.app.config.as_ref().clone();
|
||||
app.max_concurrent_jobs = Some(0);
|
||||
let agent_name = format!(
|
||||
"test_agent_{}",
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_nanos()
|
||||
);
|
||||
let agent_dir = paths::agent_data_dir(&agent_name);
|
||||
create_dir_all(&agent_dir).unwrap();
|
||||
write(
|
||||
agent_dir.join("config.yaml"),
|
||||
format!("name: {agent_name}\ninstructions: hi\n"),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.unwrap()
|
||||
.block_on(async {
|
||||
ctx.use_agent(&app, &agent_name, None, utils::create_abort_signal())
|
||||
.await
|
||||
.unwrap();
|
||||
});
|
||||
|
||||
assert!(ctx.supervisor.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn current_depth_default_is_zero() {
|
||||
let ctx = create_test_ctx();
|
||||
|
||||
@@ -0,0 +1,161 @@
|
||||
use crate::supervisor::Supervisor;
|
||||
|
||||
use parking_lot::RwLock;
|
||||
use std::sync::Arc;
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn is_agent_task(supervisor: Option<&Arc<RwLock<Supervisor>>>, id: &str) -> bool {
|
||||
id.starts_with("agent_")
|
||||
|| id.starts_with("graph_agent_")
|
||||
|| supervisor.is_some_and(|sup| sup.read().has_agent(id))
|
||||
}
|
||||
|
||||
pub struct RingBuf {
|
||||
buf: Vec<u8>,
|
||||
capacity: usize,
|
||||
write_pos: usize,
|
||||
total_written: u64,
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
impl RingBuf {
|
||||
pub fn new(capacity: usize) -> Self {
|
||||
Self {
|
||||
buf: Vec::new(),
|
||||
capacity,
|
||||
write_pos: 0,
|
||||
total_written: 0,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn push(&mut self, bytes: &[u8]) {
|
||||
self.total_written += bytes.len() as u64;
|
||||
if self.capacity == 0 {
|
||||
return;
|
||||
}
|
||||
let src = if bytes.len() > self.capacity {
|
||||
&bytes[bytes.len() - self.capacity..]
|
||||
} else {
|
||||
bytes
|
||||
};
|
||||
for &byte in src {
|
||||
if self.buf.len() < self.capacity {
|
||||
self.buf.push(byte);
|
||||
} else {
|
||||
self.buf[self.write_pos] = byte;
|
||||
}
|
||||
self.write_pos = (self.write_pos + 1) % self.capacity;
|
||||
}
|
||||
}
|
||||
|
||||
pub fn total_written(&self) -> u64 {
|
||||
self.total_written
|
||||
}
|
||||
|
||||
pub fn tail(&self) -> Vec<u8> {
|
||||
if self.buf.len() < self.capacity {
|
||||
return self.buf.clone();
|
||||
}
|
||||
let mut out = Vec::with_capacity(self.capacity);
|
||||
out.extend_from_slice(&self.buf[self.write_pos..]);
|
||||
out.extend_from_slice(&self.buf[..self.write_pos]);
|
||||
out
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for RingBuf {
|
||||
fn default() -> Self {
|
||||
Self::new(64 * 1024)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn ring_buf_returns_contents_below_capacity() {
|
||||
let mut buf = RingBuf::new(8);
|
||||
buf.push(b"abc");
|
||||
buf.push(b"de");
|
||||
assert_eq!(buf.tail(), b"abcde");
|
||||
assert_eq!(buf.total_written(), 5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ring_buf_exact_fit_keeps_everything() {
|
||||
let mut buf = RingBuf::new(5);
|
||||
buf.push(b"abcde");
|
||||
assert_eq!(buf.tail(), b"abcde");
|
||||
assert_eq!(buf.total_written(), 5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ring_buf_wrap_around_keeps_newest_bytes() {
|
||||
let mut buf = RingBuf::new(5);
|
||||
buf.push(b"abcde");
|
||||
buf.push(b"fg");
|
||||
assert_eq!(buf.tail(), b"cdefg");
|
||||
assert_eq!(buf.total_written(), 7);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ring_buf_oversize_push_keeps_last_capacity_bytes() {
|
||||
let mut buf = RingBuf::new(4);
|
||||
buf.push(b"abcdefghij");
|
||||
assert_eq!(buf.tail(), b"ghij");
|
||||
assert_eq!(buf.total_written(), 10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ring_buf_default_capacity_is_64_kib() {
|
||||
let mut buf = RingBuf::default();
|
||||
let payload = vec![b'x'; 64 * 1024 + 1];
|
||||
buf.push(&payload);
|
||||
assert_eq!(buf.tail().len(), 64 * 1024);
|
||||
assert_eq!(buf.total_written(), 64 * 1024 + 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_agent_task_matches_agent_prefixes() {
|
||||
assert!(is_agent_task(None, "agent_explore_a1b2c3d4"));
|
||||
assert!(is_agent_task(None, "graph_agent_explore_a1b2c3d4"));
|
||||
assert!(!is_agent_task(None, "job_deadbeef"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_agent_task_matches_registered_agents() {
|
||||
use crate::supervisor::mailbox::Inbox;
|
||||
use crate::supervisor::{AgentExitStatus, AgentHandle, AgentResult};
|
||||
use crate::utils::create_abort_signal;
|
||||
|
||||
let rt = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.unwrap();
|
||||
let join_handle = rt.spawn(async {
|
||||
Ok(AgentResult {
|
||||
id: "a1".into(),
|
||||
agent_name: "explore".into(),
|
||||
output: String::new(),
|
||||
exit_status: AgentExitStatus::Completed,
|
||||
})
|
||||
});
|
||||
std::mem::forget(rt);
|
||||
let handle = AgentHandle {
|
||||
id: "a1".to_string(),
|
||||
agent_name: "explore".to_string(),
|
||||
depth: 1,
|
||||
inbox: Arc::new(Inbox::new()),
|
||||
abort_signal: create_abort_signal(),
|
||||
join_handle,
|
||||
child_supervisor: None,
|
||||
};
|
||||
let mut sup = Supervisor::new(4, 3);
|
||||
sup.register(handle).unwrap();
|
||||
let sup = Arc::new(RwLock::new(sup));
|
||||
|
||||
assert!(is_agent_task(Some(&sup), "a1"));
|
||||
assert!(!is_agent_task(Some(&sup), "missing"));
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
pub(crate) mod jobs;
|
||||
pub(crate) mod memory;
|
||||
pub(crate) mod rag_query;
|
||||
pub(crate) mod skill;
|
||||
|
||||
+223
-25
@@ -1,7 +1,8 @@
|
||||
use super::{FunctionDeclaration, JsonSchema};
|
||||
use crate::client::{Model, ModelType, call_chat_completions};
|
||||
use crate::config::{
|
||||
Agent, AppState, Input, RequestContext, Role, RoleLike, list_agents_with_descriptions,
|
||||
Agent, AppState, Input, RequestContext, Role, RoleLike, effective_max_concurrent_jobs,
|
||||
jobs_enabled, list_agents_with_descriptions,
|
||||
};
|
||||
use crate::supervisor::mailbox::{Envelope, EnvelopePayload, Inbox};
|
||||
use crate::supervisor::{AgentExitStatus, AgentHandle, AgentResult, Supervisor};
|
||||
@@ -32,6 +33,19 @@ fn agent_permitted(whitelist: Option<&[String]>, target: &str) -> bool {
|
||||
}
|
||||
}
|
||||
|
||||
fn is_job_task(supervisor: Option<&Arc<RwLock<Supervisor>>>, id: &str) -> bool {
|
||||
id.starts_with("job_") || supervisor.is_some_and(|sup| sup.read().has_job(id))
|
||||
}
|
||||
|
||||
fn job_id_teaching_error(id: &str) -> Value {
|
||||
json!({
|
||||
"status": "error",
|
||||
"message": format!(
|
||||
"'{id}' is a background job, not an agent — use job__check / job__collect / job__cancel"
|
||||
),
|
||||
})
|
||||
}
|
||||
|
||||
pub enum GuardrailAction {
|
||||
NoAction,
|
||||
Inject(String),
|
||||
@@ -557,9 +571,15 @@ pub async fn run_agent_for_graph(
|
||||
|
||||
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 agent_max_concurrent = agent.max_concurrent_agents();
|
||||
let child_jobs_enabled = jobs_enabled(Some(&agent), app_config.as_ref());
|
||||
let should_init_supervisor = agent.can_spawn_agents() || child_jobs_enabled;
|
||||
let agent_max_concurrent = if agent.can_spawn_agents() {
|
||||
agent.max_concurrent_agents()
|
||||
} else {
|
||||
0
|
||||
};
|
||||
let agent_max_depth = agent.max_agent_depth();
|
||||
let agent_max_jobs = effective_max_concurrent_jobs(Some(&agent), app_config.as_ref());
|
||||
|
||||
let mut child_ctx = RequestContext::new_for_child(
|
||||
Arc::clone(&child_app_state),
|
||||
@@ -571,10 +591,10 @@ pub async fn run_agent_for_graph(
|
||||
child_ctx.rag = agent.rag();
|
||||
child_ctx.agent = Some(agent);
|
||||
if should_init_supervisor {
|
||||
child_ctx.supervisor = Some(Arc::new(RwLock::new(Supervisor::new(
|
||||
agent_max_concurrent,
|
||||
agent_max_depth,
|
||||
))));
|
||||
child_ctx.supervisor = Some(Arc::new(RwLock::new(
|
||||
Supervisor::new(agent_max_concurrent, agent_max_depth)
|
||||
.with_max_concurrent_jobs(agent_max_jobs),
|
||||
)));
|
||||
}
|
||||
|
||||
if let Some(session) = session {
|
||||
@@ -736,9 +756,15 @@ async fn handle_spawn(ctx: &mut RequestContext, args: &Value) -> Result<Value> {
|
||||
|
||||
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 child_jobs_enabled = jobs_enabled(Some(&agent), app_config.as_ref());
|
||||
let should_init_supervisor = agent.can_spawn_agents() || child_jobs_enabled;
|
||||
let max_concurrent = if agent.can_spawn_agents() {
|
||||
agent.max_concurrent_agents()
|
||||
} else {
|
||||
0
|
||||
};
|
||||
let max_depth = agent.max_agent_depth();
|
||||
let max_jobs = effective_max_concurrent_jobs(Some(&agent), app_config.as_ref());
|
||||
let mut child_ctx = RequestContext::new_for_child(
|
||||
Arc::clone(&child_app_state),
|
||||
ctx,
|
||||
@@ -749,10 +775,9 @@ async fn handle_spawn(ctx: &mut RequestContext, args: &Value) -> Result<Value> {
|
||||
child_ctx.rag = agent.rag();
|
||||
child_ctx.agent = Some(agent);
|
||||
if should_init_supervisor {
|
||||
child_ctx.supervisor = Some(Arc::new(RwLock::new(Supervisor::new(
|
||||
max_concurrent,
|
||||
max_depth,
|
||||
))));
|
||||
child_ctx.supervisor = Some(Arc::new(RwLock::new(
|
||||
Supervisor::new(max_concurrent, max_depth).with_max_concurrent_jobs(max_jobs),
|
||||
)));
|
||||
}
|
||||
|
||||
if let Some(session) = session {
|
||||
@@ -861,10 +886,15 @@ async fn handle_check(ctx: &mut RequestContext, args: &Value) -> Result<Value> {
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
None => Ok(json!({
|
||||
"status": "error",
|
||||
"message": format!("No agent found with id '{id}'")
|
||||
})),
|
||||
None => {
|
||||
if is_job_task(ctx.supervisor.as_ref(), id) {
|
||||
return Ok(job_id_teaching_error(id));
|
||||
}
|
||||
Ok(json!({
|
||||
"status": "error",
|
||||
"message": format!("No agent found with id '{id}'")
|
||||
}))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -883,6 +913,9 @@ async fn handle_collect(ctx: &mut RequestContext, args: &Value) -> Result<Value>
|
||||
let target_abort = {
|
||||
let sup = supervisor.read();
|
||||
if sup.is_finished(id).is_none() {
|
||||
if id.starts_with("job_") || sup.has_job(id) {
|
||||
return Ok(job_id_teaching_error(id));
|
||||
}
|
||||
return Ok(json!({
|
||||
"status": "error",
|
||||
"message": format!("Agent '{id}' not found. Use agent__check to verify it exists and is finished.")
|
||||
@@ -1065,10 +1098,15 @@ async fn handle_cancel(ctx: &mut RequestContext, args: &Value) -> Result<Value>
|
||||
"message": message,
|
||||
}))
|
||||
}
|
||||
None => Ok(json!({
|
||||
"status": "error",
|
||||
"message": format!("No agent found with id '{id}'"),
|
||||
})),
|
||||
None => {
|
||||
if is_job_task(ctx.supervisor.as_ref(), id) {
|
||||
return Ok(job_id_teaching_error(id));
|
||||
}
|
||||
Ok(json!({
|
||||
"status": "error",
|
||||
"message": format!("No agent found with id '{id}'"),
|
||||
}))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1115,10 +1153,17 @@ fn handle_send_message(ctx: &mut RequestContext, args: &Value) -> Result<Value>
|
||||
"message": format!("Message delivered to agent '{id}'"),
|
||||
}))
|
||||
}
|
||||
None => Ok(json!({
|
||||
"status": "error",
|
||||
"message": format!("No agent found with id '{id}'. Agent may not exist or may have already completed."),
|
||||
})),
|
||||
None => {
|
||||
if is_job_task(ctx.supervisor.as_ref(), id)
|
||||
|| is_job_task(ctx.parent_supervisor.as_ref(), id)
|
||||
{
|
||||
return Ok(job_id_teaching_error(id));
|
||||
}
|
||||
Ok(json!({
|
||||
"status": "error",
|
||||
"message": format!("No agent found with id '{id}'. Agent may not exist or may have already completed."),
|
||||
}))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1455,7 +1500,10 @@ mod tests {
|
||||
use super::*;
|
||||
use crate::config::test_fixtures::{FixtureServer, fixture_runtime};
|
||||
use crate::config::{AgentConfig, AppState, WorkingMode};
|
||||
use crate::function::jobs::RingBuf;
|
||||
use crate::supervisor::escalation::{EscalationQueue, EscalationRequest};
|
||||
use crate::supervisor::{JobHandle, JobResult, JobState, JobStatus};
|
||||
use parking_lot::Mutex;
|
||||
use serde_json::json;
|
||||
use serial_test::serial;
|
||||
|
||||
@@ -1472,6 +1520,59 @@ mod tests {
|
||||
ctx
|
||||
}
|
||||
|
||||
fn ctx_with_job_capable_supervisor() -> RequestContext {
|
||||
let mut ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd);
|
||||
ctx.supervisor = Some(Arc::new(RwLock::new(
|
||||
Supervisor::new(4, 3).with_max_concurrent_jobs(4),
|
||||
)));
|
||||
ctx
|
||||
}
|
||||
|
||||
fn make_fake_job(id: &str) -> JobHandle {
|
||||
let rt = tokio::runtime::Runtime::new().unwrap();
|
||||
let join_handle = rt.spawn(async {
|
||||
Ok(JobResult {
|
||||
output: json!(null),
|
||||
exit_code: Some(0),
|
||||
output_bytes_captured: 0,
|
||||
})
|
||||
});
|
||||
std::mem::forget(rt);
|
||||
JobHandle {
|
||||
id: id.to_string(),
|
||||
tool: "execute_command".to_string(),
|
||||
started_at: std::time::Instant::now(),
|
||||
join_handle,
|
||||
abort_signal: create_abort_signal(),
|
||||
state: Arc::new(Mutex::new(JobState {
|
||||
status: JobStatus::Running,
|
||||
pgid: None,
|
||||
})),
|
||||
output_buf: Arc::new(Mutex::new(RingBuf::default())),
|
||||
no_change_checks: 0,
|
||||
}
|
||||
}
|
||||
|
||||
fn register_fake_job(ctx: &mut RequestContext, id: &str) {
|
||||
ctx.supervisor
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.write()
|
||||
.register(make_fake_job(id))
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
fn assert_job_teaching_error(result: &Value, id: &str) {
|
||||
assert_eq!(result["status"], "error");
|
||||
let message = result["message"].as_str().unwrap();
|
||||
assert_eq!(
|
||||
message,
|
||||
format!(
|
||||
"'{id}' is a background job, not an agent — use job__check / job__collect / job__cancel"
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
fn register_fake_agent(ctx: &mut RequestContext, id: &str, name: &str) {
|
||||
register_fake_agent_with_output(ctx, id, name, "fake output");
|
||||
}
|
||||
@@ -2533,4 +2634,101 @@ mod tests {
|
||||
assert_eq!(ctx.pending_agents_guardrail_count, 0);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn handle_check_registered_job_id_teaches_job_tools() {
|
||||
run_async(async {
|
||||
let mut ctx = ctx_with_job_capable_supervisor();
|
||||
register_fake_job(&mut ctx, "bg_1");
|
||||
|
||||
let result = handle_check(&mut ctx, &json!({"id": "bg_1"}))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_job_teaching_error(&result, "bg_1");
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn handle_check_job_prefixed_id_teaches_job_tools() {
|
||||
run_async(async {
|
||||
let mut ctx = ctx_with_supervisor(4, 3);
|
||||
|
||||
let result = handle_check(&mut ctx, &json!({"id": "job_deadbeef"}))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_job_teaching_error(&result, "job_deadbeef");
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn handle_collect_registered_job_id_teaches_job_tools() {
|
||||
run_async(async {
|
||||
let mut ctx = ctx_with_job_capable_supervisor();
|
||||
register_fake_job(&mut ctx, "bg_1");
|
||||
|
||||
let result = handle_collect(&mut ctx, &json!({"id": "bg_1"}))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_job_teaching_error(&result, "bg_1");
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn handle_collect_job_prefixed_id_teaches_job_tools() {
|
||||
run_async(async {
|
||||
let mut ctx = ctx_with_supervisor(4, 3);
|
||||
|
||||
let result = handle_collect(&mut ctx, &json!({"id": "job_deadbeef"}))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_job_teaching_error(&result, "job_deadbeef");
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn handle_cancel_registered_job_id_teaches_job_tools_and_keeps_job() {
|
||||
run_async(async {
|
||||
let mut ctx = ctx_with_job_capable_supervisor();
|
||||
register_fake_job(&mut ctx, "bg_1");
|
||||
|
||||
let result = handle_cancel(&mut ctx, &json!({"id": "bg_1"}))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_job_teaching_error(&result, "bg_1");
|
||||
assert!(ctx.supervisor.as_ref().unwrap().read().has_job("bg_1"));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn handle_send_message_registered_job_id_teaches_job_tools() {
|
||||
run_async(async {
|
||||
let mut ctx = ctx_with_job_capable_supervisor();
|
||||
register_fake_job(&mut ctx, "bg_1");
|
||||
|
||||
let result =
|
||||
handle_send_message(&mut ctx, &json!({"id": "bg_1", "message": "hi"})).unwrap();
|
||||
|
||||
assert_job_teaching_error(&result, "bg_1");
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn handle_send_message_job_in_parent_supervisor_teaches_job_tools() {
|
||||
run_async(async {
|
||||
let mut ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd);
|
||||
let mut parent_sup = Supervisor::new(4, 3).with_max_concurrent_jobs(4);
|
||||
parent_sup.register(make_fake_job("bg_p")).unwrap();
|
||||
ctx.parent_supervisor = Some(Arc::new(RwLock::new(parent_sup)));
|
||||
|
||||
let result =
|
||||
handle_send_message(&mut ctx, &json!({"id": "bg_p", "message": "hi"})).unwrap();
|
||||
|
||||
assert_job_teaching_error(&result, "bg_p");
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
+295
-30
@@ -2,16 +2,19 @@ pub mod escalation;
|
||||
pub mod mailbox;
|
||||
pub mod taskqueue;
|
||||
|
||||
use crate::function::jobs::RingBuf;
|
||||
use crate::utils::AbortSignal;
|
||||
use fmt::{Debug, Formatter};
|
||||
use mailbox::Inbox;
|
||||
use parking_lot::RwLock;
|
||||
use parking_lot::{Mutex, RwLock};
|
||||
use taskqueue::TaskQueue;
|
||||
|
||||
use anyhow::{Result, bail};
|
||||
use serde_json::Value;
|
||||
use std::collections::HashMap;
|
||||
use std::fmt;
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
use tokio::task::JoinHandle;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
@@ -37,11 +40,85 @@ pub struct AgentHandle {
|
||||
pub child_supervisor: Option<Arc<RwLock<Supervisor>>>,
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum JobStatus {
|
||||
Running,
|
||||
Completed,
|
||||
Failed,
|
||||
}
|
||||
|
||||
pub struct JobState {
|
||||
#[allow(dead_code)]
|
||||
pub status: JobStatus,
|
||||
pub pgid: Option<i32>,
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub struct JobResult {
|
||||
pub output: Value,
|
||||
pub exit_code: Option<i32>,
|
||||
pub output_bytes_captured: u64,
|
||||
}
|
||||
|
||||
pub struct JobHandle {
|
||||
pub id: String,
|
||||
#[allow(dead_code)]
|
||||
pub tool: String,
|
||||
#[allow(dead_code)]
|
||||
pub started_at: Instant,
|
||||
pub join_handle: JoinHandle<Result<JobResult>>,
|
||||
pub abort_signal: AbortSignal,
|
||||
pub state: Arc<Mutex<JobState>>,
|
||||
#[allow(dead_code)]
|
||||
pub output_buf: Arc<Mutex<RingBuf>>,
|
||||
#[allow(dead_code)]
|
||||
pub no_change_checks: u32,
|
||||
}
|
||||
|
||||
impl JobHandle {
|
||||
// pgid == child pid under process_group(0); after wait() reaps the child
|
||||
// the pid can be recycled, so never kill unless pgid is still set.
|
||||
fn kill_process_group(&self) {
|
||||
#[cfg(unix)]
|
||||
if let Some(pgid) = self.state.lock().pgid {
|
||||
unsafe {
|
||||
libc::killpg(pgid, libc::SIGTERM);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for JobHandle {
|
||||
fn drop(&mut self) {
|
||||
self.kill_process_group();
|
||||
self.join_handle.abort();
|
||||
}
|
||||
}
|
||||
|
||||
pub enum TaskHandle {
|
||||
Agent(AgentHandle),
|
||||
Job(JobHandle),
|
||||
}
|
||||
|
||||
impl From<AgentHandle> for TaskHandle {
|
||||
fn from(handle: AgentHandle) -> Self {
|
||||
Self::Agent(handle)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<JobHandle> for TaskHandle {
|
||||
fn from(handle: JobHandle) -> Self {
|
||||
Self::Job(handle)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct Supervisor {
|
||||
handles: HashMap<String, AgentHandle>,
|
||||
handles: HashMap<String, TaskHandle>,
|
||||
task_queue: TaskQueue,
|
||||
max_concurrent: usize,
|
||||
max_depth: usize,
|
||||
max_concurrent_jobs: usize,
|
||||
}
|
||||
|
||||
impl Supervisor {
|
||||
@@ -51,17 +128,43 @@ impl Supervisor {
|
||||
task_queue: TaskQueue::new(),
|
||||
max_concurrent,
|
||||
max_depth,
|
||||
max_concurrent_jobs: 0,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_max_concurrent_jobs(mut self, max_concurrent_jobs: usize) -> Self {
|
||||
self.max_concurrent_jobs = max_concurrent_jobs;
|
||||
self
|
||||
}
|
||||
|
||||
fn agent(&self, id: &str) -> Option<&AgentHandle> {
|
||||
match self.handles.get(id) {
|
||||
Some(TaskHandle::Agent(handle)) => Some(handle),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn agents(&self) -> impl Iterator<Item = &AgentHandle> {
|
||||
self.handles.values().filter_map(|handle| match handle {
|
||||
TaskHandle::Agent(handle) => Some(handle),
|
||||
TaskHandle::Job(_) => None,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn active_count(&self) -> usize {
|
||||
self.handles.len()
|
||||
self.agents().count()
|
||||
}
|
||||
|
||||
pub fn effective_active_count(&self) -> usize {
|
||||
self.agents()
|
||||
.filter(|h| !h.join_handle.is_finished())
|
||||
.count()
|
||||
}
|
||||
|
||||
pub fn active_job_count(&self) -> usize {
|
||||
self.handles
|
||||
.values()
|
||||
.filter(|h| !h.join_handle.is_finished())
|
||||
.filter(|h| matches!(h, TaskHandle::Job(job) if !job.join_handle.is_finished()))
|
||||
.count()
|
||||
}
|
||||
|
||||
@@ -73,6 +176,11 @@ impl Supervisor {
|
||||
self.max_depth
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn max_concurrent_jobs(&self) -> usize {
|
||||
self.max_concurrent_jobs
|
||||
}
|
||||
|
||||
pub fn task_queue(&self) -> &TaskQueue {
|
||||
&self.task_queue
|
||||
}
|
||||
@@ -81,59 +189,111 @@ impl Supervisor {
|
||||
&mut self.task_queue
|
||||
}
|
||||
|
||||
pub fn register(&mut self, handle: AgentHandle) -> Result<()> {
|
||||
if self.effective_active_count() >= self.max_concurrent {
|
||||
bail!(
|
||||
"Cannot spawn agent: at capacity ({}/{})",
|
||||
self.effective_active_count(),
|
||||
self.max_concurrent
|
||||
);
|
||||
pub fn register(&mut self, handle: impl Into<TaskHandle>) -> Result<()> {
|
||||
match handle.into() {
|
||||
TaskHandle::Agent(handle) => {
|
||||
if self.effective_active_count() >= self.max_concurrent {
|
||||
bail!(
|
||||
"Cannot spawn agent: at capacity ({}/{})",
|
||||
self.effective_active_count(),
|
||||
self.max_concurrent
|
||||
);
|
||||
}
|
||||
if handle.depth > self.max_depth {
|
||||
bail!(
|
||||
"Cannot spawn agent: max depth exceeded ({}/{})",
|
||||
handle.depth,
|
||||
self.max_depth
|
||||
);
|
||||
}
|
||||
self.handles
|
||||
.insert(handle.id.clone(), TaskHandle::Agent(handle));
|
||||
}
|
||||
TaskHandle::Job(handle) => {
|
||||
if self.active_job_count() >= self.max_concurrent_jobs {
|
||||
bail!(
|
||||
"Cannot start job: at capacity ({}/{})",
|
||||
self.active_job_count(),
|
||||
self.max_concurrent_jobs
|
||||
);
|
||||
}
|
||||
self.handles
|
||||
.insert(handle.id.clone(), TaskHandle::Job(handle));
|
||||
}
|
||||
}
|
||||
if handle.depth > self.max_depth {
|
||||
bail!(
|
||||
"Cannot spawn agent: max depth exceeded ({}/{})",
|
||||
handle.depth,
|
||||
self.max_depth
|
||||
);
|
||||
}
|
||||
self.handles.insert(handle.id.clone(), handle);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn is_finished(&self, id: &str) -> Option<bool> {
|
||||
self.handles.get(id).map(|h| h.join_handle.is_finished())
|
||||
self.agent(id).map(|h| h.join_handle.is_finished())
|
||||
}
|
||||
|
||||
pub fn take(&mut self, id: &str) -> Option<AgentHandle> {
|
||||
self.handles.remove(id)
|
||||
self.agent(id)?;
|
||||
match self.handles.remove(id) {
|
||||
Some(TaskHandle::Agent(handle)) => Some(handle),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn take_job(&mut self, id: &str) -> Option<JobHandle> {
|
||||
if !self.has_job(id) {
|
||||
return None;
|
||||
}
|
||||
match self.handles.remove(id) {
|
||||
Some(TaskHandle::Job(handle)) => Some(handle),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn has_job(&self, id: &str) -> bool {
|
||||
matches!(self.handles.get(id), Some(TaskHandle::Job(_)))
|
||||
}
|
||||
|
||||
pub fn has_agent(&self, id: &str) -> bool {
|
||||
self.agent(id).is_some()
|
||||
}
|
||||
|
||||
pub fn inbox(&self, id: &str) -> Option<&Arc<Inbox>> {
|
||||
self.handles.get(id).map(|h| &h.inbox)
|
||||
self.agent(id).map(|h| &h.inbox)
|
||||
}
|
||||
|
||||
pub fn abort_signal_for(&self, id: &str) -> Option<AbortSignal> {
|
||||
self.handles.get(id).map(|h| h.abort_signal.clone())
|
||||
self.agent(id).map(|h| h.abort_signal.clone())
|
||||
}
|
||||
|
||||
pub fn list_agents(&self) -> Vec<(&str, &str)> {
|
||||
self.handles
|
||||
.values()
|
||||
self.agents()
|
||||
.map(|h| (h.id.as_str(), h.agent_name.as_str()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn cancel_all(&self) {
|
||||
for handle in self.handles.values() {
|
||||
handle.abort_signal.set_ctrlc();
|
||||
match handle {
|
||||
TaskHandle::Agent(agent) => agent.abort_signal.set_ctrlc(),
|
||||
TaskHandle::Job(job) => {
|
||||
job.abort_signal.set_ctrlc();
|
||||
job.kill_process_group();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn cancel_recursive(&self) {
|
||||
for handle in self.handles.values() {
|
||||
handle.abort_signal.set_ctrlc();
|
||||
if let Some(child_sup) = handle.child_supervisor.as_ref() {
|
||||
child_sup.read().cancel_recursive();
|
||||
match handle {
|
||||
TaskHandle::Agent(agent) => {
|
||||
agent.abort_signal.set_ctrlc();
|
||||
if let Some(child_sup) = agent.child_supervisor.as_ref() {
|
||||
child_sup.read().cancel_recursive();
|
||||
}
|
||||
}
|
||||
TaskHandle::Job(job) => {
|
||||
job.abort_signal.set_ctrlc();
|
||||
job.kill_process_group();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -142,7 +302,7 @@ impl Supervisor {
|
||||
impl Debug for Supervisor {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("Supervisor")
|
||||
.field("active_agents", &self.handles.len())
|
||||
.field("active_agents", &self.active_count())
|
||||
.field("max_concurrent", &self.max_concurrent)
|
||||
.field("max_depth", &self.max_depth)
|
||||
.finish()
|
||||
@@ -177,6 +337,33 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
fn make_job(id: &str, abort_signal: AbortSignal) -> JobHandle {
|
||||
// Keep the runtime alive so the spawned task is never polled and the
|
||||
// job counts as running for capacity checks.
|
||||
let rt = Builder::new_current_thread().enable_all().build().unwrap();
|
||||
let join_handle = rt.spawn(async {
|
||||
Ok(JobResult {
|
||||
output: Value::Null,
|
||||
exit_code: Some(0),
|
||||
output_bytes_captured: 0,
|
||||
})
|
||||
});
|
||||
std::mem::forget(rt);
|
||||
JobHandle {
|
||||
id: id.to_string(),
|
||||
tool: "execute_command".to_string(),
|
||||
started_at: Instant::now(),
|
||||
join_handle,
|
||||
abort_signal,
|
||||
state: Arc::new(Mutex::new(JobState {
|
||||
status: JobStatus::Running,
|
||||
pgid: None,
|
||||
})),
|
||||
output_buf: Arc::new(Mutex::new(RingBuf::default())),
|
||||
no_change_checks: 0,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn supervisor_new_empty() {
|
||||
let sup = Supervisor::new(4, 3);
|
||||
@@ -315,4 +502,82 @@ mod tests {
|
||||
assert!(parent_sig.aborted());
|
||||
assert!(child_sig.aborted());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn job_registration_rejects_when_job_capacity_zero() {
|
||||
let mut sup = Supervisor::new(4, 3);
|
||||
let result = sup.register(make_job("j1", create_abort_signal()));
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().to_string().contains("at capacity"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn job_registration_rejects_at_job_capacity() {
|
||||
let mut sup = Supervisor::new(4, 3).with_max_concurrent_jobs(1);
|
||||
sup.register(make_job("j1", create_abort_signal())).unwrap();
|
||||
let result = sup.register(make_job("j2", create_abort_signal()));
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().to_string().contains("at capacity"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn job_capacity_is_independent_of_agent_capacity() {
|
||||
let mut sup = Supervisor::new(1, 3).with_max_concurrent_jobs(1);
|
||||
sup.register(make_job("j1", create_abort_signal())).unwrap();
|
||||
sup.register(make_handle("a1", "explore", 1)).unwrap();
|
||||
assert_eq!(sup.active_job_count(), 1);
|
||||
assert_eq!(sup.active_count(), 1);
|
||||
assert_eq!(sup.max_concurrent_jobs(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn agent_accessors_ignore_jobs() {
|
||||
let mut sup = Supervisor::new(4, 3).with_max_concurrent_jobs(2);
|
||||
sup.register(make_job("j1", create_abort_signal())).unwrap();
|
||||
|
||||
assert_eq!(sup.active_count(), 0);
|
||||
assert_eq!(sup.effective_active_count(), 0);
|
||||
assert!(sup.list_agents().is_empty());
|
||||
assert_eq!(sup.is_finished("j1"), None);
|
||||
assert!(sup.inbox("j1").is_none());
|
||||
assert!(sup.abort_signal_for("j1").is_none());
|
||||
assert!(sup.take("j1").is_none());
|
||||
assert!(sup.has_job("j1"));
|
||||
assert!(!sup.has_agent("j1"));
|
||||
assert_eq!(sup.active_job_count(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn take_job_removes_job_but_not_agents() {
|
||||
let mut sup = Supervisor::new(4, 3).with_max_concurrent_jobs(2);
|
||||
sup.register(make_job("j1", create_abort_signal())).unwrap();
|
||||
sup.register(make_handle("a1", "explore", 1)).unwrap();
|
||||
|
||||
assert!(sup.take_job("a1").is_none());
|
||||
assert!(sup.has_agent("a1"));
|
||||
assert!(sup.take_job("j1").is_some());
|
||||
assert_eq!(sup.active_job_count(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cancel_recursive_aborts_jobs() {
|
||||
let sig = create_abort_signal();
|
||||
let mut sup = Supervisor::new(4, 3).with_max_concurrent_jobs(1);
|
||||
sup.register(make_job("j1", sig.clone())).unwrap();
|
||||
|
||||
sup.cancel_recursive();
|
||||
|
||||
assert!(sig.aborted());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cancel_all_aborts_jobs() {
|
||||
let sig = create_abort_signal();
|
||||
let mut sup = Supervisor::new(4, 3).with_max_concurrent_jobs(1);
|
||||
sup.register(make_job("j1", sig.clone())).unwrap();
|
||||
|
||||
sup.cancel_all();
|
||||
|
||||
assert!(sig.aborted());
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user