diff --git a/Cargo.lock b/Cargo.lock index 05028a1..d856e4a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1703,6 +1703,7 @@ dependencies = [ "inquire", "is-terminal", "json-patch", + "libc", "log", "log4rs", "nu-ansi-term", diff --git a/Cargo.toml b/Cargo.toml index d8473a6..fbfa196 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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"] } diff --git a/src/config/agent.rs b/src/config/agent.rs index 40771fa..acddb70 100644 --- a/src/config/agent.rs +++ b/src/config/agent.rs @@ -561,6 +561,10 @@ impl Agent { self.config.max_tool_result_chars } + pub fn max_concurrent_jobs(&self) -> Option { + self.config.max_concurrent_jobs + } + pub fn compression_keep_last(&self) -> Option { 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, #[serde(default, skip_serializing_if = "Option::is_none")] + pub max_concurrent_jobs: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] pub compression_keep_last: Option, #[serde(default)] pub description: String, diff --git a/src/config/app_config.rs b/src/config/app_config.rs index 5dd6b2d..a050234 100644 --- a/src/config/app_config.rs +++ b/src/config/app_config.rs @@ -68,6 +68,7 @@ pub struct AppConfig { pub summarization_prompt: Option, pub summary_context_prompt: Option, pub max_tool_result_chars: Option, + pub max_concurrent_jobs: Option, pub memory: Option, pub memory_cap_with_tools: Option, @@ -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::(&get_env_name("max_concurrent_jobs")) { + self.max_concurrent_jobs = v; + } if let Some(v) = super::read_env_value::(&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() diff --git a/src/config/mod.rs b/src/config/mod.rs index 08992fd..6f070b4 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -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, pub summary_context_prompt: Option, pub max_tool_result_chars: Option, + pub max_concurrent_jobs: Option, pub memory: Option, pub memory_cap_with_tools: Option, @@ -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, diff --git a/src/config/request_context.rs b/src/config/request_context.rs index 2be3d37..3dc350c 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -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(); diff --git a/src/function/jobs.rs b/src/function/jobs.rs new file mode 100644 index 0000000..05ac697 --- /dev/null +++ b/src/function/jobs.rs @@ -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>>, 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, + 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 { + 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")); + } +} diff --git a/src/function/mod.rs b/src/function/mod.rs index f552854..9d7d3d6 100644 --- a/src/function/mod.rs +++ b/src/function/mod.rs @@ -1,3 +1,4 @@ +pub(crate) mod jobs; pub(crate) mod memory; pub(crate) mod rag_query; pub(crate) mod skill; diff --git a/src/function/supervisor.rs b/src/function/supervisor.rs index 87dfbe5..e69ef6c 100644 --- a/src/function/supervisor.rs +++ b/src/function/supervisor.rs @@ -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>>, 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 { 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 { 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 { 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 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 "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 "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"); + }); + } } diff --git a/src/supervisor/mod.rs b/src/supervisor/mod.rs index 0739404..b63f608 100644 --- a/src/supervisor/mod.rs +++ b/src/supervisor/mod.rs @@ -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>>, } +#[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, +} + +#[allow(dead_code)] +pub struct JobResult { + pub output: Value, + pub exit_code: Option, + 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>, + pub abort_signal: AbortSignal, + pub state: Arc>, + #[allow(dead_code)] + pub output_buf: Arc>, + #[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 for TaskHandle { + fn from(handle: AgentHandle) -> Self { + Self::Agent(handle) + } +} + +impl From for TaskHandle { + fn from(handle: JobHandle) -> Self { + Self::Job(handle) + } +} + pub struct Supervisor { - handles: HashMap, + handles: HashMap, 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 { + 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) -> 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 { - 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 { - 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 { + 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> { - self.handles.get(id).map(|h| &h.inbox) + self.agent(id).map(|h| &h.inbox) } pub fn abort_signal_for(&self, id: &str) -> Option { - 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()); + } }