From 47edaa35ffd967f22678c5339f839862caf28a37 Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Wed, 9 Sep 2026 16:32:59 -0600 Subject: [PATCH] feat: teammates flag enabling peer messaging between concurrent graph agent nodes Agent nodes marked teammates: true whose sub-agents run concurrently (map fan-out branches or multiple flagged agent nodes in one super-step) now form a team: each member gets a pre-provisioned inbox in a messaging-only PeerRegistry before any branch starts, a teammate roster injected into its prompt, peer resolution in agent__send_message, and a peers section in agent__list_running. Registries are scoped to a single fan-out and carry no lifecycle semantics: outputs still flow exclusively through collect_into/state reducers, finished peers reject new messages, and nothing registers in a supervisor. Also hardens the agent-recursion depth guard to work without a supervisor and adds validator warnings for zero-valued orchestration limits. --- graph.example.yaml | 10 + src/config/mod.rs | 1 + src/config/prompts.rs | 2 +- src/config/request_context.rs | 51 +++- src/function/agents.rs | 459 ++++++++++++++++++++++++++++++---- src/graph/agent.rs | 1 + src/graph/executor.rs | 173 +++++++++++++ src/graph/map.rs | 87 +++++++ src/graph/types.rs | 38 +++ src/graph/validator.rs | 73 ++++++ src/supervisor/mailbox.rs | 125 +++++++++ 11 files changed, 975 insertions(+), 45 deletions(-) diff --git a/graph.example.yaml b/graph.example.yaml index 029766e..95b6514 100644 --- a/graph.example.yaml +++ b/graph.example.yaml @@ -334,6 +334,16 @@ nodes: Research {{topic}} in depth. Existing context: {{context}} timeout: 600 # Optional wall-clock cap, seconds. Default 300. + teammates: false # Optional. When true and this node runs concurrently + # with other flagged agent nodes (map fan-out branches + # or the same super-step), the siblings can message + # each other via agent__send_message / agent__check_inbox. + # Each fan-out gets its OWN registry: a `teammates` map + # fan-out and a `teammates` agent node in the same + # super-step frontier form two separate teams and cannot + # message each other. Peers only span a single map + # node's branches or a single super-step's flagged + # agent nodes. output_schema: # Optional. Same extraction as llm nodes type: object properties: diff --git a/src/config/mod.rs b/src/config/mod.rs index 10d29ea..3c71f1f 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -25,6 +25,7 @@ mod update; #[cfg(test)] pub(crate) use self::agent::AgentConfig; +pub(crate) use self::agent::default_max_agent_depth; pub use self::agent::{ Agent, AgentVariable, AgentVariables, complete_agent_variables, list_agents, list_agents_with_descriptions, diff --git a/src/config/prompts.rs b/src/config/prompts.rs index 4baa732..bc9818e 100644 --- a/src/config/prompts.rs +++ b/src/config/prompts.rs @@ -215,7 +215,7 @@ pub(in crate::config) const DEFAULT_TEAMMATE_INSTRUCTIONS: &str = indoc! {" ## Teammate Messaging You have tools to communicate with other agents running alongside you: - - `agent__send_message --id --message \"...\"`: Send a message to a child or sibling agent by its id, or to your parent with the reserved id `parent`. + - `agent__send_message --id --message \"...\"`: Send a message to a child, sibling, or concurrent teammate agent by its id, or to your parent with the reserved id `parent`. - `agent__check_inbox`: Check for messages sent to you by other agents. If you are working alongside other agents (e.g. reviewing different files, exploring different areas): diff --git a/src/config/request_context.rs b/src/config/request_context.rs index 48a2de4..e596ad7 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -35,7 +35,7 @@ use crate::mcp::{ use crate::rag::Rag; use crate::supervisor::Supervisor; use crate::supervisor::escalation::EscalationQueue; -use crate::supervisor::mailbox::Inbox; +use crate::supervisor::mailbox::{Inbox, PeerRegistry}; use crate::supervisor::notification::NotificationQueue; use crate::utils::{ AbortSignal, abortable_run_with_spinner, edit_file, fuzzy_filter, get_env_name, @@ -328,6 +328,13 @@ pub struct RequestContext { pub self_agent_id: Option, pub inbox: Option>, pub parent_inbox: Option>, + /// Directory of concurrent teammates for `teammates: true` graph fan-outs. + /// Never inherited: fork sites and `run_agent_for_graph` propagate it + /// explicitly so it cannot leak into nested fan-outs or spawned children. + pub peer_registry: Option>, + /// Pre-provisioned (agent id, inbox) for the sub-agent this branch context + /// will run, consumed by `run_agent_for_graph`. + pub peer_assignment: Option<(String, Arc)>, pub escalation_queue: Option>, pub notification_queue: Arc, pub current_depth: usize, @@ -369,6 +376,8 @@ impl RequestContext { self_agent_id: None, inbox: None, parent_inbox: None, + peer_registry: None, + peer_assignment: None, escalation_queue: None, notification_queue: Arc::new(NotificationQueue::new()), current_depth: 0, @@ -435,6 +444,8 @@ impl RequestContext { self_agent_id: None, inbox: None, parent_inbox: None, + peer_registry: None, + peer_assignment: None, escalation_queue: None, notification_queue: Arc::new(NotificationQueue::new()), current_depth: 0, @@ -490,6 +501,8 @@ impl RequestContext { self_agent_id: self.self_agent_id.clone(), inbox: self.inbox.clone(), parent_inbox: self.parent_inbox.clone(), + peer_registry: None, + peer_assignment: None, escalation_queue: self.escalation_queue.clone(), notification_queue: self.notification_queue.clone(), current_depth: self.current_depth, @@ -541,6 +554,8 @@ impl RequestContext { self_agent_id: Some(self_agent_id), inbox: Some(inbox), parent_inbox: parent.inbox.clone(), + peer_registry: None, + peer_assignment: None, escalation_queue: parent.escalation_queue.clone(), notification_queue: Arc::new(NotificationQueue::new()), current_depth, @@ -6040,6 +6055,40 @@ mod tests { assert!(Arc::ptr_eq(&parent_inbox, &child_parent_inbox)); } + #[test] + fn fork_for_branch_resets_peer_fields() { + let mut ctx = create_test_ctx(); + ctx.peer_registry = Some(Arc::new(PeerRegistry::new())); + ctx.peer_assignment = Some(("graph_agent_a_1".to_string(), Arc::new(Inbox::new()))); + + let branch = ctx.fork_for_branch(); + + assert!( + branch.peer_registry.is_none() && branch.peer_assignment.is_none(), + "branches must not inherit peer identity; the executor provisions per-branch assignments" + ); + } + + #[test] + fn new_for_child_resets_peer_fields() { + let mut parent = create_test_ctx(); + parent.peer_registry = Some(Arc::new(PeerRegistry::new())); + parent.peer_assignment = Some(("graph_agent_a_1".to_string(), Arc::new(Inbox::new()))); + + let child = RequestContext::new_for_child( + Arc::clone(&parent.app), + &parent, + 1, + Arc::new(Inbox::new()), + "agent_test_1".to_string(), + ); + + assert!( + child.peer_registry.is_none() && child.peer_assignment.is_none(), + "children must not inherit the parent frontier's peer roster" + ); + } + #[test] fn escalation_queue_defaults_to_none() { let ctx = create_test_ctx(); diff --git a/src/function/agents.rs b/src/function/agents.rs index 1b753c9..a62b2c1 100644 --- a/src/function/agents.rs +++ b/src/function/agents.rs @@ -2,10 +2,10 @@ use super::todo::TODO_FUNCTION_PREFIX; use super::{FunctionDeclaration, JsonSchema}; use crate::client::{Model, ModelType, call_chat_completions}; use crate::config::{ - Agent, AppState, Input, RequestContext, Role, RoleLike, effective_max_concurrent_jobs, - jobs_enabled, list_agents_with_descriptions, + Agent, AppState, Input, RequestContext, Role, RoleLike, default_max_agent_depth, + effective_max_concurrent_jobs, jobs_enabled, list_agents_with_descriptions, }; -use crate::supervisor::mailbox::{Envelope, EnvelopePayload, Inbox}; +use crate::supervisor::mailbox::{Envelope, EnvelopePayload, Inbox, PeerRegistry, graph_agent_id}; use crate::supervisor::notification::agent_notification; use crate::supervisor::{AgentExitStatus, AgentHandle, AgentResult, Supervisor, TaskKind}; use crate::utils::{AbortSignal, create_abort_signal, wait_abort_signal, wait_user_interrupt}; @@ -314,8 +314,10 @@ pub fn agent_function_declarations() -> Vec { FunctionDeclaration { name: format!("{AGENT_FUNCTION_PREFIX}list_running"), description: "List all subagents YOU have spawned that are still tracked by the supervisor, with their \ - status. Use this to see which of your background agents are still active. To discover which \ - agent types you can spawn in the first place, use `agent__list_available` instead.".to_string(), + status. Use this to see which of your background agents are still active. When you run with \ + concurrent teammates, they appear in a separate `peers` section (message-only, not \ + collectible). To discover which agent types you can spawn in the first place, use \ + `agent__list_available` instead.".to_string(), parameters: JsonSchema { type_value: Some("object".to_string()), properties: Some(IndexMap::new()), @@ -466,7 +468,7 @@ pub fn teammate_function_declarations() -> Vec { vec![ FunctionDeclaration { name: format!("{AGENT_FUNCTION_PREFIX}send_message"), - description: "Send a text message to another agent's inbox: a child or sibling by its agent id, or your parent via the reserved id 'parent'. Use to share cross-cutting findings or coordinate with teammates.".to_string(), + description: "Send a text message to another agent's inbox: a child, sibling, or concurrent teammate by its agent id, or your parent via the reserved id 'parent'. Use to share cross-cutting findings or coordinate with teammates.".to_string(), parameters: JsonSchema { type_value: Some("object".to_string()), properties: Some(IndexMap::from([ @@ -474,7 +476,7 @@ pub fn teammate_function_declarations() -> Vec { "id".to_string(), JsonSchema { type_value: Some("string".to_string()), - description: Some("The target agent ID, or the reserved id 'parent' to message the agent that spawned you".into()), + description: Some("The target agent ID (including teammate peer ids), or the reserved id 'parent' to message the agent that spawned you".into()), ..Default::default() }, ), @@ -494,7 +496,7 @@ pub fn teammate_function_declarations() -> Vec { }, FunctionDeclaration { name: format!("{AGENT_FUNCTION_PREFIX}check_inbox"), - description: "Check for and drain all pending messages in your inbox from sibling agents or your parent.".to_string(), + description: "Check for and drain all pending messages in your inbox from sibling agents, concurrent teammates, or your parent.".to_string(), parameters: JsonSchema { type_value: Some("object".to_string()), properties: Some(IndexMap::new()), @@ -636,6 +638,49 @@ fn todo_continuation_prompt(ctx: &mut RequestContext) -> Option { Some(format!("{prompt}\n\n{}", ctx.todo_list.render_for_model())) } +fn render_teammate_roster(self_id: &str, registry: &PeerRegistry) -> String { + let mut own_label = None; + let mut lines = Vec::new(); + + for (id, label) in registry.roster() { + if id == self_id { + own_label = Some(label); + } else { + lines.push(format!("- `{id}` — {label}")); + } + } + + let own = own_label.map(|l| format!(" ({l})")).unwrap_or_default(); + + format!( + "## Teammate roster\nYou are `{self_id}`{own}. Concurrent teammates:\n{}\n\nMessage a teammate directly via `agent__send_message` when you find something relevant to their slice; check `agent__check_inbox` after completing work phases and before finalizing. You can rediscover peers anytime via `agent__list_running`.", + lines.join("\n") + ) +} + +fn augment_prompt_with_roster( + prompt: &str, + registry: Option<&PeerRegistry>, + self_id: &str, +) -> String { + match registry { + Some(registry) => format!("{prompt}\n\n{}", render_teammate_roster(self_id, registry)), + None => prompt.to_string(), + } +} + +fn effective_max_agent_depth(parent_ctx: &RequestContext) -> usize { + if let Some(supervisor) = parent_ctx.supervisor.as_ref() { + return supervisor.read().max_depth(); + } + + parent_ctx + .agent + .as_ref() + .map(|agent| agent.max_agent_depth()) + .unwrap_or_else(default_max_agent_depth) +} + /// Spawn an agent synchronously from a graph node and return its accumulated /// output. This is similar to `handle_spawn` but runs the child agent in the /// current task (no tokio::spawn, no supervisor handle registration) so the @@ -645,22 +690,26 @@ pub async fn run_agent_for_graph( agent_name: &str, prompt: &str, ) -> Result { - let short_uuid = &Uuid::new_v4().to_string()[..8]; - let agent_id = format!("graph_agent_{agent_name}_{short_uuid}"); + let peer_assignment = parent_ctx.peer_assignment.take(); + let assigned_peer = peer_assignment.is_some(); + let peer_registry = if assigned_peer { + parent_ctx.peer_registry.take() + } else { + None + }; + let (agent_id, child_inbox) = + peer_assignment.unwrap_or_else(|| (graph_agent_id(agent_name), Arc::new(Inbox::new()))); let current_depth = parent_ctx.current_depth + 1; - if let Some(supervisor) = parent_ctx.supervisor.as_ref().cloned() { - let max_depth = supervisor.read().max_depth(); - if current_depth > max_depth { - bail!("Max agent depth exceeded ({current_depth}/{max_depth})"); - } + let max_depth = effective_max_agent_depth(parent_ctx); + if current_depth > max_depth { + bail!("Max agent depth exceeded ({current_depth}/{max_depth})"); } if !parent_ctx.app.config.function_calling_support { bail!("Function calling support must be enabled to spawn agents."); } - let child_inbox = Arc::new(Inbox::new()); parent_ctx.ensure_root_escalation_queue(); parent_ctx.ensure_inbox(); let child_abort = create_abort_signal(); @@ -710,6 +759,9 @@ pub async fn run_agent_for_graph( ); child_ctx.rag = agent.rag(); child_ctx.agent = Some(agent); + if assigned_peer { + child_ctx.peer_registry = peer_registry.clone(); + } if should_init_supervisor { child_ctx.supervisor = Some(Arc::new(RwLock::new( Supervisor::new(agent_max_concurrent_subagents, agent_max_depth) @@ -729,11 +781,18 @@ pub async fn run_agent_for_graph( child_ctx.init_agent_shared_variables()?; } - let input = Input::from_str(&child_ctx, prompt, None)?; + let prompt = augment_prompt_with_roster(prompt, child_ctx.peer_registry.as_deref(), &agent_id); + let input = Input::from_str(&child_ctx, &prompt, None)?; debug!("Spawning agent '{agent_name}' for graph node as '{agent_id}'"); - run_child_agent(child_ctx, input, child_abort).await + let result = run_child_agent(child_ctx, input, child_abort).await; + + if let Some(registry) = &peer_registry { + registry.mark_finished(&agent_id); + } + + result } async fn populate_agent_mcp_runtime(ctx: &mut RequestContext, server_ids: &[String]) -> Result<()> { @@ -1171,31 +1230,59 @@ async fn handle_collect(ctx: &mut RequestContext, args: &Value) -> Result } fn handle_list_running(ctx: &mut RequestContext) -> Result { - let supervisor = ctx - .supervisor - .as_ref() - .cloned() - .ok_or_else(|| anyhow!("No supervisor active"))?; - let sup = supervisor.read(); - - let agents: Vec = sup - .list_agents() - .into_iter() - .map(|(id, name)| { - let finished = sup.is_finished(id).unwrap_or(false); - json!({ - "id": id, - "agent": name, - "status": if finished { "finished" } else { "running" }, + let peers: Option> = ctx.peer_registry.as_ref().map(|registry| { + registry + .roster() + .into_iter() + .filter(|(id, _)| Some(id.as_str()) != ctx.self_agent_id.as_deref()) + .map(|(id, label)| { + let status = if registry.is_finished(&id) { + "finished" + } else { + "running" + }; + json!({ "id": id, "label": label, "kind": "peer", "status": status }) }) - }) - .collect(); + .collect() + }); - Ok(json!({ - "active_count": sup.active_count(), - "max_concurrent": sup.max_concurrent(), - "agents": agents, - })) + let supervisor = ctx.supervisor.as_ref().cloned(); + if supervisor.is_none() && peers.is_none() { + return Err(anyhow!("No supervisor active")); + } + + let mut result = match supervisor { + Some(supervisor) => { + let sup = supervisor.read(); + let agents: Vec = sup + .list_agents() + .into_iter() + .map(|(id, name)| { + let finished = sup.is_finished(id).unwrap_or(false); + json!({ + "id": id, + "agent": name, + "status": if finished { "finished" } else { "running" }, + }) + }) + .collect(); + json!({ + "active_count": sup.active_count(), + "max_concurrent": sup.max_concurrent(), + "agents": agents, + }) + } + None => json!({}), + }; + + if let Some(peers) = peers { + result["peers"] = Value::Array(peers); + result["peers_note"] = json!( + "Peers are concurrent teammates reachable via agent__send_message; they cannot be checked, collected, or cancelled." + ); + } + + Ok(result) } fn handle_list_available(ctx: &RequestContext) -> Result { @@ -1351,6 +1438,25 @@ fn handle_send_message(ctx: &mut RequestContext, args: &Value) -> Result })); } + if let Some(registry) = ctx.peer_registry.as_ref() + && let Some(inbox) = registry.get(id) + { + if registry.is_finished(id) { + return Ok(json!({ + "status": "error", + "message": format!( + "Teammate '{id}' has already finished; message not delivered. \ + Its results flow through the graph's collected outputs." + ), + })); + } + deliver(&inbox, self_label); + return Ok(json!({ + "status": "ok", + "message": format!("Message delivered to teammate '{id}'"), + })); + } + 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)); } @@ -1693,20 +1799,66 @@ async fn summarize_output(ctx: &RequestContext, agent_name: &str, output: &str) mod tests { use super::*; use crate::config::test_fixtures::{FixtureServer, fixture_runtime}; - use crate::config::{AgentConfig, AppConfig, AppState, WorkingMode}; + use crate::config::{AgentConfig, AppConfig, AppState, WorkingMode, paths}; use crate::function::jobs::RingBuf; use crate::mcp::{McpServer, McpServersConfig, McpTransportType}; use crate::supervisor::escalation::{EscalationQueue, EscalationRequest}; use crate::supervisor::{JobHandle, JobResult, JobState, JobStatus}; + use crate::utils::get_env_name; use parking_lot::Mutex; use serde_json::json; use serial_test::serial; - use std::mem; + use std::fs::{create_dir_all, remove_dir_all, write}; + use std::path::PathBuf; + use std::time::{SystemTime, UNIX_EPOCH}; + use std::{env, mem}; fn default_app_state() -> Arc { Arc::new(AppState::test_default()) } + struct TestConfigDirGuard { + key: String, + previous: Option, + path: PathBuf, + } + + impl TestConfigDirGuard { + fn new() -> Self { + let key = get_env_name("config_dir"); + let previous = env::var_os(&key); + let unique = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos(); + let path = env::temp_dir().join(format!("coyote-agents-tests-{unique}")); + create_dir_all(&path).unwrap(); + unsafe { + env::set_var(&key, &path); + } + Self { + key, + previous, + path, + } + } + } + + impl Drop for TestConfigDirGuard { + fn drop(&mut self) { + if let Some(previous) = &self.previous { + unsafe { + env::set_var(&self.key, previous); + } + } else { + unsafe { + env::remove_var(&self.key); + } + } + let _ = remove_dir_all(&self.path); + } + } + fn ctx_with_supervisor(max_concurrent: usize, max_depth: usize) -> RequestContext { let mut ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd); ctx.supervisor = Some(Arc::new(RwLock::new(Supervisor::new( @@ -1716,6 +1868,57 @@ mod tests { ctx } + #[test] + fn effective_max_agent_depth_prefers_supervisor() { + let ctx = ctx_with_supervisor(4, 7); + assert_eq!(effective_max_agent_depth(&ctx), 7); + } + + #[test] + fn effective_max_agent_depth_falls_back_without_supervisor() { + let ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd); + assert!(ctx.supervisor.is_none()); + assert!(ctx.agent.is_none()); + assert_eq!(effective_max_agent_depth(&ctx), default_max_agent_depth()); + } + + #[test] + #[serial] + fn effective_max_agent_depth_uses_agent_limit_without_supervisor() { + let _guard = TestConfigDirGuard::new(); + let agent_name = format!( + "test_depth_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\nmax_agent_depth: 5\n"), + ) + .unwrap(); + + let mut ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd); + let app_config = Arc::clone(&ctx.app.config); + let model = ctx.current_model().clone(); + let agent = run_async(Agent::init( + app_config.as_ref(), + ctx.app.as_ref(), + &model, + false, + &agent_name, + create_abort_signal(), + )) + .unwrap(); + ctx.agent = Some(agent); + + assert!(ctx.supervisor.is_none()); + assert_eq!(effective_max_agent_depth(&ctx), 5); + } + fn auto_continue_ctx() -> RequestContext { let app = AppState { config: Arc::new(AppConfig { @@ -2042,6 +2245,68 @@ mod tests { assert!(result.is_err()); } + #[test] + fn handle_list_running_includes_peers_and_excludes_self() { + let mut ctx = ctx_with_supervisor(4, 3); + ctx.self_agent_id = Some("p_self".to_string()); + let registry = Arc::new(PeerRegistry::new()); + registry.insert("p_self".into(), "shards[0]".into(), Arc::new(Inbox::new())); + registry.insert("p_other".into(), "shards[1]".into(), Arc::new(Inbox::new())); + ctx.peer_registry = Some(registry); + + let result = handle_list_running(&mut ctx).unwrap(); + + assert_eq!(result["active_count"], 0); + let peers = result["peers"].as_array().unwrap(); + assert_eq!(peers.len(), 1); + assert_eq!(peers[0]["id"], "p_other"); + assert_eq!(peers[0]["label"], "shards[1]"); + assert_eq!(peers[0]["kind"], "peer"); + } + + #[test] + fn handle_list_running_peer_status_reflects_finished() { + let mut ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd); + ctx.self_agent_id = Some("p1".to_string()); + let registry = Arc::new(PeerRegistry::new()); + registry.insert("p1".into(), "shards[0]".into(), Arc::new(Inbox::new())); + registry.insert("p2".into(), "shards[1]".into(), Arc::new(Inbox::new())); + registry.insert("p3".into(), "shards[2]".into(), Arc::new(Inbox::new())); + registry.mark_finished("p3"); + ctx.peer_registry = Some(registry); + + let result = handle_list_running(&mut ctx).unwrap(); + + let peers = result["peers"].as_array().unwrap(); + assert_eq!(peers.len(), 2); + assert_eq!(peers[0]["id"], "p2"); + assert_eq!(peers[0]["status"], "running"); + assert_eq!(peers[1]["id"], "p3"); + assert_eq!(peers[1]["status"], "finished"); + } + + #[test] + fn handle_list_running_peers_without_supervisor() { + let mut ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd); + ctx.self_agent_id = Some("p_self".to_string()); + let registry = Arc::new(PeerRegistry::new()); + registry.insert("p_self".into(), "a".into(), Arc::new(Inbox::new())); + registry.insert("p_other".into(), "b".into(), Arc::new(Inbox::new())); + ctx.peer_registry = Some(registry); + + let result = handle_list_running(&mut ctx).unwrap(); + + assert_eq!(result["peers"].as_array().unwrap().len(), 1); + assert!(result.get("agents").is_none()); + } + + #[test] + fn handle_list_running_no_peers_key_without_registry() { + let mut ctx = ctx_with_supervisor(4, 3); + let result = handle_list_running(&mut ctx).unwrap(); + assert!(result.get("peers").is_none()); + } + #[test] fn handle_list_available_returns_shape() { let ctx = ctx_with_supervisor(4, 3); @@ -2266,6 +2531,114 @@ mod tests { assert_eq!(msgs[0].from, "agent_explore_sender"); } + #[test] + fn handle_send_message_resolves_teammate_peer_id() { + let mut ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd); + ctx.self_agent_id = Some("graph_agent_worker_aaaa1111".to_string()); + let peer_inbox = Arc::new(Inbox::new()); + let registry = Arc::new(PeerRegistry::new()); + registry.insert( + "graph_agent_worker_bbbb2222".into(), + "shards[1]".into(), + Arc::clone(&peer_inbox), + ); + ctx.peer_registry = Some(registry); + + let result = handle_send_message( + &mut ctx, + &json!({"id": "graph_agent_worker_bbbb2222", "message": "hi peer"}), + ) + .unwrap(); + assert_eq!(result["status"], "ok"); + + let msgs = peer_inbox.drain(); + assert_eq!(msgs.len(), 1); + // The sender's self_agent_id IS its peer id, so the recipient can + // reply symmetrically through the same registry. + assert_eq!(msgs[0].from, "graph_agent_worker_aaaa1111"); + match &msgs[0].payload { + EnvelopePayload::Text { content } => assert_eq!(content, "hi peer"), + _ => panic!("expected text payload"), + } + } + + #[test] + fn handle_send_message_to_finished_teammate_errors() { + let mut ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd); + ctx.self_agent_id = Some("graph_agent_worker_aaaa1111".to_string()); + let peer_inbox = Arc::new(Inbox::new()); + let registry = Arc::new(PeerRegistry::new()); + registry.insert( + "graph_agent_worker_bbbb2222".into(), + "shards[1]".into(), + Arc::clone(&peer_inbox), + ); + registry.mark_finished("graph_agent_worker_bbbb2222"); + ctx.peer_registry = Some(registry); + + let result = handle_send_message( + &mut ctx, + &json!({"id": "graph_agent_worker_bbbb2222", "message": "hi peer"}), + ) + .unwrap(); + + assert_eq!(result["status"], "error"); + assert!( + result["message"] + .as_str() + .unwrap() + .contains("has already finished; message not delivered"), + "{result}" + ); + assert!(peer_inbox.drain().is_empty()); + } + + #[test] + fn handle_send_message_unknown_id_with_peer_registry_errors() { + let mut ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd); + ctx.peer_registry = Some(Arc::new(PeerRegistry::new())); + let result = + handle_send_message(&mut ctx, &json!({"id": "missing", "message": "hi"})).unwrap(); + assert_eq!(result["status"], "error"); + } + + #[test] + fn teammate_roster_excludes_self_and_includes_labels() { + let registry = PeerRegistry::new(); + registry.insert("p1".into(), "shards[0]".into(), Arc::new(Inbox::new())); + registry.insert("p2".into(), "shards[1]".into(), Arc::new(Inbox::new())); + registry.insert("p3".into(), "shards[2]".into(), Arc::new(Inbox::new())); + + let roster = render_teammate_roster("p2", ®istry); + + assert!(roster.contains("You are `p2` (shards[1])"), "{roster}"); + assert!(roster.contains("- `p1` — shards[0]"), "{roster}"); + assert!(roster.contains("- `p3` — shards[2]"), "{roster}"); + assert!(!roster.contains("- `p2`"), "{roster}"); + } + + #[test] + fn augment_prompt_with_roster_appends_non_self_peers() { + let registry = PeerRegistry::new(); + registry.insert("p1".into(), "shards[0]".into(), Arc::new(Inbox::new())); + registry.insert("p2".into(), "shards[1]".into(), Arc::new(Inbox::new())); + + let prompt = augment_prompt_with_roster("do the work", Some(®istry), "p1"); + + assert!(prompt.starts_with("do the work\n\n"), "{prompt}"); + assert!(prompt.contains("You are `p1` (shards[0])"), "{prompt}"); + assert!(prompt.contains("- `p2` — shards[1]"), "{prompt}"); + assert!(!prompt.contains("- `p1`"), "{prompt}"); + } + + #[test] + fn augment_prompt_without_registry_is_unchanged() { + assert_eq!( + augment_prompt_with_roster("do the work", None, "p1"), + "do the work" + ); + } + #[test] fn handle_check_inbox_with_messages() { let mut ctx = ctx_with_supervisor(4, 3); diff --git a/src/graph/agent.rs b/src/graph/agent.rs index 9bde501..e5eb7d1 100644 --- a/src/graph/agent.rs +++ b/src/graph/agent.rs @@ -113,6 +113,7 @@ mod tests { state_updates: updates, output_schema: None, timeout: None, + teammates: false, } } diff --git a/src/graph/executor.rs b/src/graph/executor.rs index 9a178f9..aa854b5 100644 --- a/src/graph/executor.rs +++ b/src/graph/executor.rs @@ -10,6 +10,7 @@ use super::types::{EndNode, Graph, Node, NodeType}; use super::user_interaction::{ApprovalNodeExecutor, InputNodeExecutor}; use super::validator::{AgentValidationContext, GraphValidator}; use crate::config::{RenderMode, RequestContext}; +use crate::supervisor::mailbox::{Inbox, PeerAssignment, PeerRegistry, graph_agent_id}; use crate::utils::AbortSignal; use anyhow::{Context, Result, anyhow, bail}; use futures_util::future::join_all; @@ -160,6 +161,13 @@ impl GraphExecutor { logger.super_step_start(&branches); } + // Pre-provision teammate identities for `teammates: true` agent + // nodes running concurrently in this super-step, before any branch + // starts, so an early finisher can message one still waiting on + // the semaphore. A lone flagged node has no peers. + let flagged = teammate_flagged_nodes(&graph, &frontier); + let (peer_registry, mut peer_assignments) = provision_frontier_peers(flagged); + let mut branch_tasks = Vec::with_capacity(frontier_size); for node_id in &frontier { let node = graph @@ -171,6 +179,10 @@ impl GraphExecutor { logger.node_start(&node, in_super_step); let branch_state = state.fork_for_branch_state(); let mut branch_ctx = ctx.fork_for_branch(); + if let Some(assignment) = peer_assignments.remove(node_id) { + branch_ctx.peer_registry = peer_registry.clone(); + branch_ctx.peer_assignment = Some(assignment); + } if in_super_step { branch_ctx.render_mode = RenderMode::Silent; } @@ -355,6 +367,36 @@ fn sorted_frontier(frontier: &HashSet) -> Vec { v } +fn teammate_flagged_nodes(graph: &Graph, frontier: &HashSet) -> Vec<(String, String)> { + frontier + .iter() + .filter_map( + |node_id| match graph.get_node(node_id).map(|n| &n.node_type) { + Some(NodeType::Agent(n)) if n.teammates => Some((node_id.clone(), n.agent.clone())), + _ => None, + }, + ) + .collect() +} + +fn provision_frontier_peers( + flagged: Vec<(String, String)>, +) -> (Option>, HashMap) { + if flagged.len() < 2 { + return (None, HashMap::new()); + } + + let registry = Arc::new(PeerRegistry::new()); + let mut assignments = HashMap::new(); + for (node_id, agent_name) in flagged { + let id = graph_agent_id(&agent_name); + let inbox = Arc::new(Inbox::new()); + registry.insert(id.clone(), node_id.clone(), Arc::clone(&inbox)); + assignments.insert(node_id, (id, inbox)); + } + (Some(registry), assignments) +} + pub(super) struct StepContext<'a> { pub graph: &'a Graph, pub script_executor: &'a ScriptExecutor, @@ -475,7 +517,9 @@ fn apply_simple_state_updates(updates: Option<&HashMap>, state: #[cfg(test)] mod tests { + use super::super::types::{AgentNode, GraphSettings}; use super::*; + use indexmap::IndexMap; use serde_json::json; fn state_with(pairs: &[(&str, Value)]) -> StateManager { @@ -557,6 +601,135 @@ mod tests { assert_eq!(state.state().get("k"), Some(&json!("new-old"))); } + + fn agent_node(id: &str, teammates: bool) -> Node { + Node { + id: id.into(), + description: String::new(), + node_type: NodeType::Agent(AgentNode { + agent: format!("{id}-agent"), + prompt: "p".into(), + state_updates: None, + output_schema: None, + timeout: None, + teammates, + }), + next: None, + } + } + + fn terminal_node(id: &str) -> Node { + Node { + id: id.into(), + description: String::new(), + node_type: NodeType::End(EndNode { + output: String::new(), + state_updates: None, + }), + next: None, + } + } + + fn graph_of(nodes: Vec) -> Graph { + let start = nodes[0].id.clone(); + let mut map: IndexMap = IndexMap::new(); + for node in nodes { + map.insert(node.id.clone(), node); + } + + Graph { + name: "t".into(), + description: String::new(), + version: "1.0".into(), + model: None, + temperature: None, + top_p: None, + reasoning_effort: None, + max_concurrent_jobs: None, + can_spawn_agents: None, + max_concurrent_agents: None, + max_agent_depth: None, + global_tools: Vec::new(), + mcp_servers: Vec::new(), + mcp_tools: None, + skills_enabled: None, + enabled_skills: None, + inject_skill_instructions: None, + skill_instructions: None, + conversation_starters: Vec::new(), + variables: Vec::new(), + settings: GraphSettings::default(), + initial_state: HashMap::new(), + reducers: HashMap::new(), + start, + nodes: map, + } + } + + #[test] + fn teammate_flagged_nodes_selects_only_flagged_agent_nodes() { + let graph = graph_of(vec![ + agent_node("a", true), + agent_node("b", true), + agent_node("c", false), + terminal_node("done"), + ]); + let frontier: HashSet = ["a", "b", "c", "done"] + .iter() + .map(|s| s.to_string()) + .collect(); + + let mut flagged = teammate_flagged_nodes(&graph, &frontier); + flagged.sort(); + + assert_eq!( + flagged, + vec![ + ("a".to_string(), "a-agent".to_string()), + ("b".to_string(), "b-agent".to_string()), + ] + ); + } + + #[test] + fn provision_frontier_peers_two_flagged_share_one_registry() { + let flagged = vec![ + ("a".to_string(), "worker".to_string()), + ("b".to_string(), "worker".to_string()), + ]; + + let (registry, assignments) = provision_frontier_peers(flagged); + + let registry = registry.expect("two flagged nodes should get a registry"); + assert_eq!(assignments.len(), 2); + let roster = registry.roster(); + assert_eq!(roster.len(), 2); + + let (id_a, inbox_a) = &assignments["a"]; + let (id_b, _) = &assignments["b"]; + assert_ne!(id_a, id_b); + let resolved = registry.get(id_a).expect("assigned id should resolve"); + assert!(Arc::ptr_eq(&resolved, inbox_a)); + assert!(roster.iter().any(|(id, label)| id == id_a && label == "a")); + assert!(roster.iter().any(|(id, label)| id == id_b && label == "b")); + } + + #[test] + fn provision_frontier_peers_single_flagged_gets_no_registry() { + let (registry, assignments) = + provision_frontier_peers(vec![("a".to_string(), "worker".to_string())]); + + assert!(registry.is_none()); + assert!(assignments.is_empty()); + } + + #[test] + fn provision_frontier_peers_empty_frontier_gets_no_registry() { + let (registry, assignments) = provision_frontier_peers(Vec::new()); + + assert!(registry.is_none()); + assert!(assignments.is_empty()); + } } #[cfg(test)] diff --git a/src/graph/map.rs b/src/graph/map.rs index 1222887..e541950 100644 --- a/src/graph/map.rs +++ b/src/graph/map.rs @@ -6,6 +6,7 @@ use super::state::StateManager; use super::types::{MapNode, NodeType}; use crate::config::{RenderMode, RequestContext}; use crate::graph::type_name; +use crate::supervisor::mailbox::{Inbox, PeerAssignment, PeerRegistry, graph_agent_id}; use anyhow::{Context, Result, anyhow}; use futures_util::future::join_all; use serde_json::Value; @@ -55,6 +56,8 @@ impl MapNodeExecutor { let semaphore = Arc::new(Semaphore::new(max_conc)); let mut sub_tasks = Vec::with_capacity(items.len()); + let peers = provision_map_peers(&branch_node.node_type, &node.branch, items.len()); + for (idx, item) in items.iter().enumerate() { let item = item.clone(); let as_name = node.as_name.clone(); @@ -62,6 +65,10 @@ impl MapNodeExecutor { let mut sub_state = state.fork_for_branch_state(); let mut sub_ctx = ctx.fork_for_branch(); sub_ctx.render_mode = RenderMode::Silent; + if let Some((registry, assignments)) = &peers { + sub_ctx.peer_registry = Some(Arc::clone(registry)); + sub_ctx.peer_assignment = Some(assignments[idx].clone()); + } let script_clone = step_ctx.script_executor.clone(); let sub_branch_id = node.branch.clone(); let sem = semaphore.clone(); @@ -152,3 +159,83 @@ impl MapNodeExecutor { Ok(()) } } + +/// Pre-provision teammate identities before any branch starts, so a fast +/// sibling can message one still waiting on the semaphore (the message +/// queues in the pre-created inbox). A single-item fan-out has no peers, so +/// it gets no registry. +fn provision_map_peers( + branch: &NodeType, + branch_id: &str, + item_count: usize, +) -> Option<(Arc, Vec)> { + match branch { + NodeType::Agent(n) if n.teammates && item_count >= 2 => { + let registry = Arc::new(PeerRegistry::new()); + let assignments = (0..item_count) + .map(|i| { + let id = graph_agent_id(&n.agent); + let inbox = Arc::new(Inbox::new()); + registry.insert(id.clone(), format!("{branch_id}[{i}]"), inbox.clone()); + (id, inbox) + }) + .collect(); + Some((registry, assignments)) + } + _ => None, + } +} + +#[cfg(test)] +mod tests { + use super::super::types::{AgentNode, EndNode}; + use super::*; + + fn agent_branch(teammates: bool) -> NodeType { + NodeType::Agent(AgentNode { + agent: "worker".into(), + prompt: "p".into(), + state_updates: None, + output_schema: None, + timeout: None, + teammates, + }) + } + + #[test] + fn provision_map_peers_registers_every_branch() { + let (registry, assignments) = + provision_map_peers(&agent_branch(true), "shards", 3).expect("should provision"); + + assert_eq!(assignments.len(), 3); + let roster = registry.roster(); + assert_eq!(roster.len(), 3); + for (i, (id, inbox)) in assignments.iter().enumerate() { + assert_eq!(roster[i].0, *id); + assert_eq!(roster[i].1, format!("shards[{i}]")); + let resolved = registry.get(id).expect("assigned id should resolve"); + assert!(Arc::ptr_eq(&resolved, inbox)); + } + assert_ne!(assignments[0].0, assignments[1].0); + assert_ne!(assignments[1].0, assignments[2].0); + } + + #[test] + fn provision_map_peers_single_item_gets_no_registry() { + assert!(provision_map_peers(&agent_branch(true), "shards", 1).is_none()); + } + + #[test] + fn provision_map_peers_unflagged_gets_no_registry() { + assert!(provision_map_peers(&agent_branch(false), "shards", 3).is_none()); + } + + #[test] + fn provision_map_peers_non_agent_branch_gets_no_registry() { + let branch = NodeType::End(EndNode { + output: String::new(), + state_updates: None, + }); + assert!(provision_map_peers(&branch, "shards", 3).is_none()); + } +} diff --git a/src/graph/types.rs b/src/graph/types.rs index 01c7220..7af4483 100644 --- a/src/graph/types.rs +++ b/src/graph/types.rs @@ -239,6 +239,9 @@ pub struct AgentNode { #[serde(default, skip_serializing_if = "Option::is_none")] pub timeout: Option, + + #[serde(default)] + pub teammates: bool, } #[derive(Debug, Clone, Deserialize, Serialize)] @@ -545,6 +548,41 @@ nodes: } } + #[test] + fn agent_node_teammates_parses_and_defaults_false() { + let yaml = r#" +name: g +start: a +nodes: + a: + id: a + type: agent + agent: helper + prompt: hi + teammates: true + next: b + b: + id: b + type: agent + agent: helper + prompt: hi + next: e + e: + id: e + type: end + output: done +"#; + let graph: Graph = serde_yaml::from_str(yaml).unwrap(); + match &graph.get_node("a").unwrap().node_type { + NodeType::Agent(n) => assert!(n.teammates), + _ => panic!("expected Agent variant"), + } + match &graph.get_node("b").unwrap().node_type { + NodeType::Agent(n) => assert!(!n.teammates), + _ => panic!("expected Agent variant"), + } + } + #[test] fn deserializes_mcp_tools_at_graph_and_node_level() { let yaml = r#" diff --git a/src/graph/validator.rs b/src/graph/validator.rs index f7abb34..b900e14 100644 --- a/src/graph/validator.rs +++ b/src/graph/validator.rs @@ -176,6 +176,7 @@ impl GraphValidator { self.validate_llm_nodes(graph, &mut result); self.validate_llm_skills(graph, &mut result); self.validate_max_concurrency(graph, &mut result); + self.validate_orchestration_limits(graph, &mut result); self.validate_map_branches(graph, &mut result); self.validate_parallel_user_interaction(graph, &mut result); self.validate_parallel_writes(graph, &mut result); @@ -529,6 +530,23 @@ impl GraphValidator { } } + fn validate_orchestration_limits(&self, graph: &Graph, result: &mut ValidationResult) { + if graph.max_agent_depth == Some(0) { + result.warning(ValidationError::new( + "`max_agent_depth: 0` forbids spawning any sub-agent, so every \ + `agent` node and dynamic spawn will fail; remove the key or \ + set it >= 1", + )); + } + if graph.max_concurrent_agents == Some(0) { + result.warning(ValidationError::new( + "`max_concurrent_agents: 0` leaves no slot for any spawn to \ + acquire, so agent spawning deadlocks; remove the key or set \ + it >= 1", + )); + } + } + fn validate_map_branches(&self, graph: &Graph, result: &mut ValidationResult) { for (map_id, node) in &graph.nodes { let NodeType::Map(m) = &node.node_type else { @@ -1027,6 +1045,9 @@ mod tests { top_p: None, reasoning_effort: None, max_concurrent_jobs: None, + can_spawn_agents: None, + max_concurrent_agents: None, + max_agent_depth: None, global_tools: Vec::new(), mcp_servers: Vec::new(), mcp_tools: None, @@ -1637,6 +1658,7 @@ mod tests { state_updates: None, output_schema: None, timeout: None, + teammates: false, }), next: next.map(NextTargets::from), } @@ -2596,6 +2618,57 @@ mod tests { ); } + #[test] + fn zero_max_agent_depth_warns() { + let mut graph = graph_with(vec![("e", end_node("e"))], "e"); + graph.max_agent_depth = Some(0); + + let result = validator().validate(&graph); + + assert!(result.is_valid()); + assert!( + result + .warnings + .iter() + .any(|w| w.message.contains("`max_agent_depth: 0`")), + "expected max_agent_depth=0 warning: {:?}", + result.warnings + ); + } + + #[test] + fn zero_max_concurrent_agents_warns() { + let mut graph = graph_with(vec![("e", end_node("e"))], "e"); + graph.max_concurrent_agents = Some(0); + + let result = validator().validate(&graph); + + assert!(result.is_valid()); + assert!( + result + .warnings + .iter() + .any(|w| w.message.contains("`max_concurrent_agents: 0`")), + "expected max_concurrent_agents=0 warning: {:?}", + result.warnings + ); + } + + #[test] + fn unset_orchestration_limits_do_not_warn() { + let graph = graph_with(vec![("e", end_node("e"))], "e"); + + let result = validator().validate(&graph); + + assert!( + !result.warnings.iter().any(|w| { + w.message.contains("max_agent_depth") || w.message.contains("max_concurrent_agents") + }), + "unset limits should not warn: {:?}", + result.warnings + ); + } + #[test] fn map_max_concurrency_zero_errors() { let mut map = map_node_basic("m", "br", Some("end")); diff --git a/src/supervisor/mailbox.rs b/src/supervisor/mailbox.rs index 396f4cb..a570c24 100644 --- a/src/supervisor/mailbox.rs +++ b/src/supervisor/mailbox.rs @@ -1,4 +1,6 @@ +use indexmap::IndexMap; use serde::{Deserialize, Serialize}; +use std::sync::Arc; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Envelope { @@ -59,6 +61,79 @@ impl Clone for Inbox { } } +pub struct PeerEntry { + pub label: String, + pub inbox: Arc, + pub finished: bool, +} + +/// A pre-provisioned teammate identity: the routable peer id and the inbox +/// registered for it in a `PeerRegistry`. +pub type PeerAssignment = (String, Arc); + +/// Messaging-only directory of concurrent sibling agents ("teammates"). +/// Entries carry no lifecycle semantics: peers are never supervised, +/// checked, or collected. The registry only routes `agent__send_message` +/// and answers roster queries, and it dies when the fan-out's Arcs drop. +#[derive(Default)] +pub struct PeerRegistry { + entries: parking_lot::RwLock>, +} + +impl PeerRegistry { + pub fn new() -> Self { + Self::default() + } + + pub fn insert(&self, id: String, label: String, inbox: Arc) { + let replaced = self.entries.write().insert( + id.clone(), + PeerEntry { + label, + inbox, + finished: false, + }, + ); + debug_assert!( + replaced.is_none(), + "peer id '{id}' registered twice; short-uuid collision or duplicate insert" + ); + } + + pub fn get(&self, id: &str) -> Option> { + self.entries + .read() + .get(id) + .map(|entry| Arc::clone(&entry.inbox)) + } + + pub fn mark_finished(&self, id: &str) { + if let Some(entry) = self.entries.write().get_mut(id) { + entry.finished = true; + } + } + + pub fn is_finished(&self, id: &str) -> bool { + self.entries + .read() + .get(id) + .is_some_and(|entry| entry.finished) + } + + pub fn roster(&self) -> Vec<(String, String)> { + self.entries + .read() + .iter() + .map(|(id, entry)| (id.clone(), entry.label.clone())) + .collect() + } +} + +pub fn graph_agent_id(agent_name: &str) -> String { + let short_uuid = &uuid::Uuid::new_v4().to_string()[..8]; + format!("graph_agent_{agent_name}_{short_uuid}") +} + #[cfg(test)] mod tests { use super::*; @@ -176,4 +251,54 @@ mod tests { } assert_eq!(inbox.drain().len(), 5); } + + #[test] + fn peer_registry_get_returns_inserted_inbox() { + let registry = PeerRegistry::new(); + let inbox = Arc::new(Inbox::new()); + registry.insert("p1".into(), "branch[0]".into(), Arc::clone(&inbox)); + + let found = registry.get("p1").expect("inserted peer should resolve"); + assert!(Arc::ptr_eq(&found, &inbox)); + } + + #[test] + fn peer_registry_get_unknown_id_is_none() { + let registry = PeerRegistry::new(); + registry.insert("p1".into(), "branch[0]".into(), Arc::new(Inbox::new())); + + assert!(registry.get("missing").is_none()); + } + + #[test] + fn peer_registry_mark_finished_flips_state() { + let registry = PeerRegistry::new(); + registry.insert("p1".into(), "branch[0]".into(), Arc::new(Inbox::new())); + + assert!(!registry.is_finished("p1")); + registry.mark_finished("p1"); + assert!(registry.is_finished("p1")); + } + + #[test] + fn peer_registry_mark_finished_unknown_id_is_noop() { + let registry = PeerRegistry::new(); + registry.mark_finished("missing"); + assert!(!registry.is_finished("missing")); + } + + #[test] + fn peer_registry_roster_preserves_insertion_order() { + let registry = PeerRegistry::new(); + registry.insert("p1".into(), "branch[0]".into(), Arc::new(Inbox::new())); + registry.insert("p2".into(), "branch[1]".into(), Arc::new(Inbox::new())); + + assert_eq!( + registry.roster(), + vec![ + ("p1".to_string(), "branch[0]".to_string()), + ("p2".to_string(), "branch[1]".to_string()), + ] + ); + } }