diff --git a/Cargo.lock b/Cargo.lock index 699e40136..030ec62a3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -857,6 +857,7 @@ dependencies = [ "serde", "serde_json", "sqlx", + "tempfile", "thiserror 2.0.18", "tokio", "tracing", diff --git a/crates/aionui-ai-agent/src/agent_task.rs b/crates/aionui-ai-agent/src/agent_task.rs index 10f72387f..753581288 100644 --- a/crates/aionui-ai-agent/src/agent_task.rs +++ b/crates/aionui-ai-agent/src/agent_task.rs @@ -529,9 +529,11 @@ mod aionrs_config_option_tests { compat_overrides: Default::default(), session_directory: std::env::temp_dir().join("aionrs-agent-task-test-sessions"), session_mode: None, + skills: Vec::new(), extra_mcp_servers: std::collections::HashMap::new(), bedrock_config: None, runtime_env: Vec::new(), + prompt_dump_dir: None, } } diff --git a/crates/aionui-ai-agent/src/dev_prompt_dump.rs b/crates/aionui-ai-agent/src/dev_prompt_dump.rs index c03cd595c..226ea965e 100644 --- a/crates/aionui-ai-agent/src/dev_prompt_dump.rs +++ b/crates/aionui-ai-agent/src/dev_prompt_dump.rs @@ -3,6 +3,8 @@ use std::io; use std::path::{Path, PathBuf}; use std::time::{SystemTime, UNIX_EPOCH}; +use serde_json::Value; + pub(crate) struct PromptDump<'a> { pub kind: &'static str, pub backend: Option<&'a str>, @@ -13,6 +15,17 @@ pub(crate) struct PromptDump<'a> { pub prompt: &'a str, } +pub(crate) struct AgentFinalInputDump<'a> { + pub kind: &'static str, + pub backend: &'static str, + pub conversation_id: &'a str, + pub session_id: Option<&'a str>, + pub msg_id: Option<&'a str>, + pub turn_id: Option<&'a str>, + pub input: Value, + pub resolved_context: Value, +} + pub(crate) fn dump_dir_for_data_dir(data_dir: &Path, enabled: bool) -> Option { if enabled { Some(data_dir.join("prompt-dumps")) @@ -51,6 +64,36 @@ pub(crate) fn dump_prompt(dump_dir: &Path, dump: PromptDump<'_>) -> io::Result

) -> io::Result { + fs::create_dir_all(dump_dir)?; + + let created_at_ms = current_time_ms(); + let stamp = current_time_nanos(); + let discriminator = dump.msg_id.or(dump.session_id).or(dump.turn_id).unwrap_or("none"); + let file_name = format!( + "{}-{}-{}-{}.json", + stamp, + sanitize_segment(dump.kind), + sanitize_segment(dump.conversation_id), + sanitize_segment(discriminator) + ); + let path = dump_dir.join(file_name); + + let body = serde_json::to_vec_pretty(&serde_json::json!({ + "kind": dump.kind, + "backend": dump.backend, + "conversation_id": dump.conversation_id, + "session_id": dump.session_id.unwrap_or("none"), + "msg_id": dump.msg_id.unwrap_or("none"), + "turn_id": dump.turn_id.unwrap_or("none"), + "created_at_ms": created_at_ms, + "input": dump.input, + "resolved_context": dump.resolved_context, + }))?; + fs::write(&path, body)?; + Ok(path) +} + fn current_time_ms() -> u128 { SystemTime::now() .duration_since(UNIX_EPOCH) @@ -91,7 +134,7 @@ mod tests { let path = dump_prompt( temp.path(), PromptDump { - kind: "acp-first-prompt", + kind: "test-prompt", backend: Some("claude"), conversation_id: "conversation/123", session_id: Some("session-1"), @@ -104,12 +147,12 @@ mod tests { assert_eq!(path.parent(), Some(temp.path())); let file_name = path.file_name().unwrap().to_string_lossy(); - assert!(file_name.contains("acp-first-prompt")); + assert!(file_name.contains("test-prompt")); assert!(file_name.contains("conversation_123")); assert!(file_name.contains("msg-1")); let content = std::fs::read_to_string(path).unwrap(); - assert!(content.contains("kind: acp-first-prompt\n")); + assert!(content.contains("kind: test-prompt\n")); assert!(content.contains("backend: claude\n")); assert!(content.contains("conversation_id: conversation/123\n")); assert!(content.contains("session_id: session-1\n")); @@ -118,6 +161,54 @@ mod tests { assert!(content.ends_with("---- prompt ----\nfinal prompt body\n")); } + #[test] + fn dump_agent_final_input_writes_json_with_raw_values() { + let temp = tempfile::tempdir().unwrap(); + + let path = dump_agent_final_input( + temp.path(), + AgentFinalInputDump { + kind: "acp-final-input", + backend: "acp", + conversation_id: "conversation/123", + session_id: Some("session-1"), + msg_id: Some("msg-1"), + turn_id: Some("turn-1"), + input: serde_json::json!({ + "content": "[Assistant Rules]\nraw secret\n[/Assistant Rules]\n\nhello" + }), + resolved_context: serde_json::json!({ + "mcp_servers": [{ + "name": "raw-mcp", + "env": { "TOKEN": "raw-token-value" } + }] + }), + }, + ) + .unwrap(); + + let file_name = path.file_name().unwrap().to_string_lossy(); + assert!(file_name.contains("acp-final-input")); + assert!(file_name.contains("conversation_123")); + assert!(file_name.contains("msg-1")); + + let dump: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(path).unwrap()).unwrap(); + assert_eq!(dump["kind"], "acp-final-input"); + assert_eq!(dump["backend"], "acp"); + assert_eq!(dump["conversation_id"], "conversation/123"); + assert_eq!(dump["session_id"], "session-1"); + assert_eq!(dump["msg_id"], "msg-1"); + assert_eq!(dump["turn_id"], "turn-1"); + assert_eq!( + dump["input"]["content"], + "[Assistant Rules]\nraw secret\n[/Assistant Rules]\n\nhello" + ); + assert_eq!( + dump["resolved_context"]["mcp_servers"][0]["env"]["TOKEN"], + "raw-token-value" + ); + } + #[test] fn disabled_dump_prompts_returns_no_dump_dir() { let data_dir = Path::new("/Users/alice/.aionui-dev"); diff --git a/crates/aionui-ai-agent/src/factory/aionrs.rs b/crates/aionui-ai-agent/src/factory/aionrs.rs index 9414f4294..758f0749e 100644 --- a/crates/aionui-ai-agent/src/factory/aionrs.rs +++ b/crates/aionui-ai-agent/src/factory/aionrs.rs @@ -29,6 +29,7 @@ pub(super) async fn build( ctx: FactoryContext, ) -> Result { let mut overrides = build_context.config; + let resolved_skills = overrides.skills.clone(); // Merge preset assistant rules into system_prompt (used as custom_prompt // in aionrs's build_system_prompt). Mirrors the old architecture's @@ -166,9 +167,11 @@ pub(super) async fn build( compat_overrides, session_directory, session_mode: overrides.session_mode, + skills: resolved_skills, extra_mcp_servers, bedrock_config, runtime_env: ctx.runtime_env, + prompt_dump_dir: crate::dev_prompt_dump::dump_dir_for_data_dir(&deps.data_dir, deps.dump_prompts), }; if let Some(system_prompt) = config.system_prompt.as_deref() diff --git a/crates/aionui-ai-agent/src/idle_scanner.rs b/crates/aionui-ai-agent/src/idle_scanner.rs index 2ba0c7e75..be3412ec5 100644 --- a/crates/aionui-ai-agent/src/idle_scanner.rs +++ b/crates/aionui-ai-agent/src/idle_scanner.rs @@ -2,6 +2,7 @@ use std::sync::Arc; use std::time::Duration; use aionui_common::{AgentKillReason, now_ms}; +use async_trait::async_trait; use tracing::{debug, info}; use crate::task_manager::IWorkerTaskManager; @@ -12,6 +13,15 @@ const DEFAULT_IDLE_TIMEOUT_SECS: i64 = 5 * 60; /// Scan interval for idle agent cleanup (1 minute). const SCAN_INTERVAL_SECS: u64 = 60; +#[async_trait] +pub trait IdleCleanupCoordinator: Send + Sync { + async fn cleanup_idle_conversations( + &self, + idle_conversation_ids: Vec, + idle_threshold_ms: i64, + ) -> Vec; +} + /// Start the background idle agent scanner. /// /// Periodically scans active tasks and kills ACP agents that have been @@ -19,10 +29,26 @@ const SCAN_INTERVAL_SECS: u64 = 60; /// /// The scanner runs until the provided `shutdown` signal resolves. pub fn start_idle_scanner( + worker_task_manager: Arc, + shutdown: tokio::sync::watch::Receiver, + idle_timeout_secs: Option, + scan_interval_secs: Option, +) -> tokio::task::JoinHandle<()> { + start_idle_scanner_with_coordinator( + worker_task_manager, + shutdown, + idle_timeout_secs, + scan_interval_secs, + None, + ) +} + +pub fn start_idle_scanner_with_coordinator( worker_task_manager: Arc, mut shutdown: tokio::sync::watch::Receiver, idle_timeout_secs: Option, scan_interval_secs: Option, + idle_cleanup_coordinator: Option>, ) -> tokio::task::JoinHandle<()> { let threshold = idle_timeout_secs.unwrap_or(DEFAULT_IDLE_TIMEOUT_SECS); let scan_interval = scan_interval_secs.unwrap_or(SCAN_INTERVAL_SECS); @@ -38,7 +64,11 @@ pub fn start_idle_scanner( loop { tokio::select! { _ = interval.tick() => { - scan_and_cleanup(&worker_task_manager, threshold*1000); + scan_and_cleanup( + &worker_task_manager, + threshold*1000, + idle_cleanup_coordinator.clone(), + ).await; } _ = shutdown.changed() => { if *shutdown.borrow() { @@ -54,24 +84,56 @@ pub fn start_idle_scanner( } /// Perform one scan: find idle tasks and kill them. -fn scan_and_cleanup(manager: &Arc, threshold_ms: i64) { +async fn scan_and_cleanup( + manager: &Arc, + threshold_ms: i64, + idle_cleanup_coordinator: Option>, +) { let started_at = now_ms(); - let idle_ids = manager.collect_idle(threshold_ms); + let mut idle_ids = manager.collect_idle(threshold_ms); if idle_ids.is_empty() { debug!(active_count = manager.active_count(), "Idle scan: no idle agents found"); return; } + if let Some(coordinator) = idle_cleanup_coordinator { + let before_count = idle_ids.len(); + idle_ids = coordinator.cleanup_idle_conversations(idle_ids, threshold_ms).await; + let handled_count = before_count.saturating_sub(idle_ids.len()); + if handled_count > 0 { + info!( + handled_count, + remaining_count = idle_ids.len(), + "Idle scan: coordinator handled idle agents" + ); + } + } + + if idle_ids.is_empty() { + info!( + elapsed_ms = now_ms().saturating_sub(started_at), + "Idle scan: cleanup completed by coordinator" + ); + return; + } + let count = idle_ids.len(); info!(count, "Idle scan: cleaning up idle agents"); + let mut waits = Vec::new(); for id in idle_ids { let manager = Arc::clone(manager); - tokio::spawn(async move { + waits.push(tokio::spawn(async move { info!(conversation_id = %id, "Idle scan: awaiting idle agent shutdown"); manager.kill_and_wait(&id, Some(AgentKillReason::IdleTimeout)).await; - }); + })); + } + + for wait in waits { + if let Err(error) = wait.await { + debug!(error = %error, "Idle scan: cleanup task join failed"); + } } info!( @@ -80,3 +142,102 @@ fn scan_and_cleanup(manager: &Arc, threshold_ms: i64) { "Idle scan: cleanup completed" ); } + +#[cfg(test)] +mod tests { + use std::sync::{Arc, Mutex}; + + use aionui_common::AgentKillReason; + use async_trait::async_trait; + + use super::*; + use crate::agent_task::AgentInstance; + use crate::error::AgentError; + use crate::types::BuildTaskOptions; + + struct RecordingTaskManager { + idle_ids: Vec, + killed: Arc>>, + } + + impl RecordingTaskManager { + fn new(idle_ids: Vec<&str>) -> Self { + Self { + idle_ids: idle_ids.into_iter().map(str::to_owned).collect(), + killed: Arc::new(Mutex::new(Vec::new())), + } + } + + fn killed(&self) -> Vec { + self.killed.lock().unwrap().clone() + } + } + + #[async_trait] + impl IWorkerTaskManager for RecordingTaskManager { + fn get_task(&self, _conversation_id: &str) -> Option { + None + } + + async fn get_or_build_task( + &self, + _conversation_id: &str, + _options: BuildTaskOptions, + ) -> Result { + Err(AgentError::internal("not used")) + } + + fn kill(&self, conversation_id: &str, _reason: Option) -> Result<(), AgentError> { + self.killed.lock().unwrap().push(conversation_id.to_owned()); + Ok(()) + } + + fn kill_and_wait( + &self, + conversation_id: &str, + reason: Option, + ) -> std::pin::Pin + Send>> { + let _ = self.kill(conversation_id, reason); + Box::pin(std::future::ready(())) + } + + async fn clear(&self) {} + + fn active_count(&self) -> usize { + self.idle_ids.len() + } + + fn collect_idle(&self, _idle_threshold_ms: i64) -> Vec { + self.idle_ids.clone() + } + } + + struct RecordingCoordinator { + seen: Arc>>, + } + + #[async_trait] + impl IdleCleanupCoordinator for RecordingCoordinator { + async fn cleanup_idle_conversations( + &self, + idle_conversation_ids: Vec, + _idle_threshold_ms: i64, + ) -> Vec { + self.seen.lock().unwrap().extend(idle_conversation_ids); + vec!["solo".to_owned()] + } + } + + #[tokio::test] + async fn scan_uses_coordinator_and_default_kills_only_unhandled_ids() { + let manager_impl = Arc::new(RecordingTaskManager::new(vec!["team-lead", "solo"])); + let manager: Arc = manager_impl.clone(); + let seen = Arc::new(Mutex::new(Vec::new())); + let coordinator: Arc = Arc::new(RecordingCoordinator { seen: seen.clone() }); + + scan_and_cleanup(&manager, 300_000, Some(coordinator)).await; + + assert_eq!(seen.lock().unwrap().clone(), vec!["team-lead", "solo"]); + assert_eq!(manager_impl.killed(), vec!["solo"]); + } +} diff --git a/crates/aionui-ai-agent/src/lib.rs b/crates/aionui-ai-agent/src/lib.rs index d949c8538..f9a2b39f3 100644 --- a/crates/aionui-ai-agent/src/lib.rs +++ b/crates/aionui-ai-agent/src/lib.rs @@ -35,7 +35,7 @@ pub use capability::skill_manager::{ }; pub use error::AgentError; pub use factory::{AgentFactoryDeps, build_agent_factory}; -pub use idle_scanner::start_idle_scanner; +pub use idle_scanner::{IdleCleanupCoordinator, start_idle_scanner, start_idle_scanner_with_coordinator}; pub use persistence::AcpSessionSyncService; pub use protocol::error::AcpError; pub use protocol::events::AgentStreamEvent; diff --git a/crates/aionui-ai-agent/src/manager/acp/agent.rs b/crates/aionui-ai-agent/src/manager/acp/agent.rs index e37ef75fa..74bb977cf 100644 --- a/crates/aionui-ai-agent/src/manager/acp/agent.rs +++ b/crates/aionui-ai-agent/src/manager/acp/agent.rs @@ -53,6 +53,27 @@ pub(super) fn user_facing_message(err: &AgentError) -> String { full.split_once(": ").map(|(_, rest)| rest.to_owned()).unwrap_or(full) } +fn build_acp_final_input_dump_value( + conversation_id: &str, + session_id: &str, + data: &SendMessageData, + final_content: &str, + resolved_context: Value, +) -> Value { + serde_json::json!({ + "kind": "acp-final-input", + "backend": "acp", + "conversation_id": conversation_id, + "session_id": session_id, + "msg_id": data.msg_id, + "turn_id": data.turn_id.as_deref().unwrap_or("none"), + "input": { + "content": final_content, + }, + "resolved_context": resolved_context, + }) +} + use super::config_option_catalog::{extract_models_from_value, extract_modes_from_value}; use super::config_options::{ConfigSetPath, ConfigSetPathError, ConfigSnapshot, resolve_set_path}; use super::mode_normalize::normalize_requested_mode; @@ -914,9 +935,8 @@ impl AcpAgentManager { .and_then(|commands| matched_slash_command(&raw_user_input, commands)) }; - let (content, dump_first_prompt) = { + let content = { let mut s = self.session.write().await; - let dump_first_prompt = s.has_pending_session_new_prelude(); let mut ctx = PromptCtx { session: &mut s, params: &self.params, @@ -925,41 +945,10 @@ impl AcpAgentManager { }; let transformed = self.pipeline.pre_send(&mut ctx, data.content.clone()).await; self.commit_session_changes(&mut s).await; - (transformed, dump_first_prompt) + transformed }; - if dump_first_prompt - && let Some(dump_dir) = - crate::dev_prompt_dump::dump_dir_for_data_dir(&self.params.data_dir, self.params.dump_prompts) - { - match crate::dev_prompt_dump::dump_prompt( - &dump_dir, - crate::dev_prompt_dump::PromptDump { - kind: "acp-first-prompt", - backend: self.backend(), - conversation_id: &self.params.conversation_id, - session_id: Some(&sid), - msg_id: Some(&data.msg_id), - turn_id: data.turn_id.as_deref(), - prompt: &content, - }, - ) { - Ok(path) => { - debug!( - conversation_id = %self.params.conversation_id, - path = %path.display(), - "DEV prompt dump written" - ); - } - Err(error) => { - warn!( - conversation_id = %self.params.conversation_id, - error = %error, - "DEV prompt dump failed" - ); - } - } - } + self.dump_acp_final_input(&sid, data, &content); let data = SendMessageData { content, @@ -969,6 +958,67 @@ impl AcpAgentManager { .await } + fn dump_acp_final_input(&self, session_id: &str, data: &SendMessageData, final_content: &str) { + let Some(dump_dir) = + crate::dev_prompt_dump::dump_dir_for_data_dir(&self.params.data_dir, self.params.dump_prompts) + else { + return; + }; + + let resolved_context = serde_json::json!({ + "agent_id": &self.params.metadata.id, + "agent_backend": self.params.metadata.backend.as_deref(), + "workspace": { + "path": &self.params.workspace.path, + "is_custom": self.params.workspace.is_custom, + }, + "preset_context": self.params.preset_context.as_deref(), + "skills": &self.params.config.skills, + "mcp_servers": serde_json::to_value(&self.params.mcp_servers).unwrap_or(Value::Null), + "config": serde_json::to_value(&self.params.config).unwrap_or(Value::Null), + }); + let value = build_acp_final_input_dump_value( + &self.params.conversation_id, + session_id, + data, + final_content, + resolved_context, + ); + let input = value.get("input").cloned().unwrap_or(Value::Null); + let resolved_context = value.get("resolved_context").cloned().unwrap_or(Value::Null); + + match crate::dev_prompt_dump::dump_agent_final_input( + &dump_dir, + crate::dev_prompt_dump::AgentFinalInputDump { + kind: "acp-final-input", + backend: "acp", + conversation_id: &self.params.conversation_id, + session_id: Some(session_id), + msg_id: Some(data.msg_id.as_str()), + turn_id: data.turn_id.as_deref(), + input, + resolved_context, + }, + ) { + Ok(path) => { + debug!( + conversation_id = %self.params.conversation_id, + msg_id = %data.msg_id, + path = %path.display(), + "DEV agent final input dump written" + ); + } + Err(error) => { + warn!( + conversation_id = %self.params.conversation_id, + msg_id = %data.msg_id, + error = %error, + "DEV agent final input dump failed" + ); + } + } + } + /// Pre-open the ACP session without sending a prompt. Called by the /// factory after `AcpAgentManager::build` so runtime preparation returns /// only after the session is ready to accept `set_mode` / `set_model` @@ -1292,7 +1342,7 @@ impl AcpAgentManager { #[cfg(test)] mod tests { - use super::{exit_status_parts, user_facing_message}; + use super::{build_acp_final_input_dump_value, exit_status_parts, user_facing_message}; use crate::agent_runtime::AgentRuntime; use crate::error::AgentError; use crate::manager::acp::{AcpAgentManager, AcpSession}; @@ -1341,6 +1391,50 @@ mod tests { String::from_utf8(buffer.lock().unwrap().clone()).unwrap() } + #[test] + fn acp_final_input_dump_value_contains_post_hook_prompt_and_context() { + let data = crate::types::SendMessageData { + content: "original team wake".into(), + msg_id: "msg-acp-final".into(), + turn_id: Some("turn-acp-final".into()), + files: Vec::new(), + inject_skills: Vec::new(), + }; + let resolved_context = serde_json::json!({ + "agent_id": "agent-claude", + "agent_backend": "claude", + "workspace": { + "path": "/tmp", + "is_custom": true, + }, + "preset_context": "Rule A", + "skills": [], + "mcp_servers": [], + }); + + let value = build_acp_final_input_dump_value( + "conv-acp-final", + "session-acp-final", + &data, + "[Assistant Rules]\nRule A\n[/Assistant Rules]\n\noriginal team wake", + resolved_context, + ); + + assert_eq!(value["kind"], "acp-final-input"); + assert_eq!(value["backend"], "acp"); + assert_eq!(value["conversation_id"], "conv-acp-final"); + assert_eq!(value["session_id"], "session-acp-final"); + assert_eq!(value["msg_id"], "msg-acp-final"); + assert_eq!(value["turn_id"], "turn-acp-final"); + assert_eq!( + value["input"]["content"], + "[Assistant Rules]\nRule A\n[/Assistant Rules]\n\noriginal team wake" + ); + assert_eq!(value["resolved_context"]["workspace"]["path"], "/tmp"); + assert_eq!(value["resolved_context"]["skills"], serde_json::json!([])); + assert!(value["resolved_context"]["mcp_servers"].is_array()); + } + #[test] fn exit_status_parts_handles_missing_status() { assert_eq!(exit_status_parts(None), (None, None)); diff --git a/crates/aionui-ai-agent/src/manager/acp/agent_session_flow.rs b/crates/aionui-ai-agent/src/manager/acp/agent_session_flow.rs index 44c7386fc..031c2d12b 100644 --- a/crates/aionui-ai-agent/src/manager/acp/agent_session_flow.rs +++ b/crates/aionui-ai-agent/src/manager/acp/agent_session_flow.rs @@ -16,7 +16,7 @@ use tokio::sync::broadcast::error::TryRecvError; use super::agent::sdk_to_snake_value; use super::agent_close::STDERR_PEEK_LINES; -use super::error_mapping::{AcpSendFailure, is_acp_session_not_found}; +use super::error_mapping::{AcpSendFailure, is_acp_session_not_found, is_missing_resumed_session}; use tracing::warn; #[derive(Debug)] @@ -81,8 +81,9 @@ impl AcpAgentManager { warn!( conversation_id = %self.params.conversation_id, stale_session_id = %stale_sid, + recovery_action = "clear_persisted_session_id_and_session_new", error = %err, - "open_session_resume: stale session id rejected by CLI; rebuilding via session/new" + "open_session_resume: stale session id rejected by CLI; clearing persisted session id and rebuilding via session/new" ); { let mut session = self.session.write().await; @@ -125,7 +126,7 @@ impl AcpAgentManager { let req = self.params.new_session_request().meta(meta); let new_response = match self.protocol.new_session(req).await { Ok(r) => r, - Err(e) if is_acp_session_not_found(&e) => { + Err(e) if is_missing_resumed_session(&e, session_id) => { return self.rebuild_after_acp_session_not_found(session_id, e).await; } Err(e) => return Err(e.into()), @@ -623,6 +624,32 @@ mod tests { assert!(!super::is_acp_session_not_found(&auth_required)); } + #[test] + fn resumed_session_resource_not_found_matches_only_the_exact_session_id() { + let exact = AcpError::ResourceNotFound { + resource: Some("ses-resume".into()), + message: "missing".into(), + }; + assert!(super::is_missing_resumed_session(&exact, "ses-resume")); + + let unrelated = AcpError::ResourceNotFound { + resource: Some("/workspace/file.txt".into()), + message: "missing".into(), + }; + assert!(!super::is_missing_resumed_session(&unrelated, "ses-resume")); + + let unspecified = AcpError::ResourceNotFound { + resource: None, + message: "missing".into(), + }; + assert!(!super::is_missing_resumed_session(&unspecified, "ses-resume")); + + let structured = AcpError::SessionNotFound { + session_id: "different-wire-id".into(), + }; + assert!(super::is_missing_resumed_session(&structured, "ses-resume")); + } + // -- empty-finish diagnostic (ELECTRON-1JG) ------------------------------- use crate::protocol::events::{ diff --git a/crates/aionui-ai-agent/src/manager/acp/error_mapping.rs b/crates/aionui-ai-agent/src/manager/acp/error_mapping.rs index ab54db1d3..bf42b65ba 100644 --- a/crates/aionui-ai-agent/src/manager/acp/error_mapping.rs +++ b/crates/aionui-ai-agent/src/manager/acp/error_mapping.rs @@ -66,6 +66,17 @@ pub(super) fn is_acp_session_not_found(err: &AcpError) -> bool { matches!(err, AcpError::SessionNotFound { .. }) } +pub(super) fn is_missing_resumed_session(err: &AcpError, resumed_session_id: &str) -> bool { + is_acp_session_not_found(err) + || matches!( + err, + AcpError::ResourceNotFound { + resource: Some(resource), + .. + } if resource == resumed_session_id + ) +} + fn acp_error_public_message(err: &AcpError) -> String { match err { AcpError::AgentInternal { code, .. } => format!("Agent internal error (code {code})"), diff --git a/crates/aionui-ai-agent/src/manager/aionrs/agent.rs b/crates/aionui-ai-agent/src/manager/aionrs/agent.rs index 23bb2ceee..bb9b1eb54 100644 --- a/crates/aionui-ai-agent/src/manager/aionrs/agent.rs +++ b/crates/aionui-ai-agent/src/manager/aionrs/agent.rs @@ -1,3 +1,4 @@ +use std::collections::HashMap; use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; @@ -29,6 +30,51 @@ use crate::types::{AionrsResolvedConfig, SendMessageData}; use super::error::aionrs_engine_error_to_send_error; +#[derive(Clone, Debug)] +struct AionrsFinalInputDumpContext { + dump_dir: PathBuf, + provider: String, + model: String, + base_url: Option, + system_prompt: Option, + session_mode: Option, + skills: Vec, + mcp_servers: HashMap, + runtime_env: Vec<(String, String)>, +} + +fn build_aionrs_final_input_dump_value( + conversation_id: &str, + workspace: &str, + context: &AionrsFinalInputDumpContext, + data: &SendMessageData, +) -> Value { + serde_json::json!({ + "kind": "aionrs-final-input", + "backend": "aionrs", + "conversation_id": conversation_id, + "session_id": "none", + "msg_id": data.msg_id, + "turn_id": data.turn_id.as_deref().unwrap_or("none"), + "input": { + "system_prompt": context.system_prompt.as_deref(), + "user_content": &data.content, + }, + "resolved_context": { + "provider": &context.provider, + "model": &context.model, + "base_url": context.base_url.as_deref(), + "workspace": { + "path": workspace, + }, + "session_mode": context.session_mode.as_deref(), + "skills": &context.skills, + "mcp_servers": serde_json::to_value(&context.mcp_servers).unwrap_or(Value::Null), + "runtime_env": &context.runtime_env, + }, + }) +} + pub struct AionrsAgentManager { runtime: AgentRuntime, engine: Mutex, @@ -44,6 +90,7 @@ pub struct AionrsAgentManager { mcp_managers: Vec>, approval_manager: Arc, confirmations: Arc>>, + final_input_dump: Option, /// Signalled by `cancel()` to abort an in-flight `engine.run()` via /// `tokio::select!` in `send_message()`. cancel_notify: Arc, @@ -71,6 +118,20 @@ impl AionrsAgentManager { let runtime = AgentRuntime::new(conversation_id.clone(), workspace.clone(), 128); let sink: Arc = Arc::new(BackendOutputSink::new(runtime.event_sender())); let runtime_env = config_extra.runtime_env.clone(); + let final_input_dump = config_extra + .prompt_dump_dir + .clone() + .map(|dump_dir| AionrsFinalInputDumpContext { + dump_dir, + provider: config_extra.provider.clone(), + model: config_extra.model.clone(), + base_url: config_extra.base_url.clone(), + system_prompt: config_extra.system_prompt.clone(), + session_mode: config_extra.session_mode.clone(), + skills: config_extra.skills.clone(), + mcp_servers: config_extra.extra_mcp_servers.clone(), + runtime_env: config_extra.runtime_env.clone(), + }); let cli_args = CliArgs { provider: Some(config_extra.provider.clone()), @@ -173,6 +234,7 @@ impl AionrsAgentManager { mcp_managers: result.mcp_managers, approval_manager, confirmations, + final_input_dump, cancel_notify: Arc::new(Notify::new()), turn_finished_notify: Arc::new(Notify::new()), }) @@ -199,6 +261,52 @@ impl AionrsAgentManager { was_running } + + fn dump_aionrs_final_input(&self, data: &SendMessageData) { + let Some(context) = self.final_input_dump.as_ref() else { + return; + }; + + let value = build_aionrs_final_input_dump_value( + self.runtime.conversation_id(), + self.runtime.workspace(), + context, + data, + ); + let input = value.get("input").cloned().unwrap_or(Value::Null); + let resolved_context = value.get("resolved_context").cloned().unwrap_or(Value::Null); + + match crate::dev_prompt_dump::dump_agent_final_input( + &context.dump_dir, + crate::dev_prompt_dump::AgentFinalInputDump { + kind: "aionrs-final-input", + backend: "aionrs", + conversation_id: self.runtime.conversation_id(), + session_id: None, + msg_id: Some(data.msg_id.as_str()), + turn_id: data.turn_id.as_deref(), + input, + resolved_context, + }, + ) { + Ok(path) => { + debug!( + conversation_id = %self.runtime.conversation_id(), + msg_id = %data.msg_id, + path = %path.display(), + "DEV agent final input dump written" + ); + } + Err(error) => { + warn!( + conversation_id = %self.runtime.conversation_id(), + msg_id = %data.msg_id, + error = %error, + "DEV agent final input dump failed" + ); + } + } + } } #[async_trait::async_trait] @@ -237,6 +345,7 @@ impl crate::agent_task::IAgentTask for AionrsAgentManager { ); self.runtime.bump_activity(); self.runtime.reset_for_new_turn(ConversationStatus::Running); + self.dump_aionrs_final_input(&data); let mut engine = self.engine.lock().await; @@ -494,12 +603,73 @@ mod tests { compat_overrides: Default::default(), session_directory: std::env::temp_dir().join("aionrs-test-sessions"), session_mode: None, + skills: Vec::new(), extra_mcp_servers: std::collections::HashMap::new(), bedrock_config: None, runtime_env: Vec::new(), + prompt_dump_dir: None, } } + #[test] + fn aionrs_final_input_dump_value_contains_raw_split_input_and_context() { + let mut mcp_env = std::collections::HashMap::new(); + mcp_env.insert("TOKEN".to_owned(), "raw-token-value".to_owned()); + + let mut mcp_servers = std::collections::HashMap::new(); + mcp_servers.insert( + "raw-mcp".to_owned(), + aion_config::config::McpServerConfig { + transport: aion_config::config::TransportType::Stdio, + command: Some("/bin/raw-mcp".to_owned()), + args: Some(vec!["--serve".to_owned()]), + env: Some(mcp_env), + url: None, + headers: None, + deferred: Some(false), + startup_timeout_ms: None, + }, + ); + + let context = AionrsFinalInputDumpContext { + dump_dir: std::path::PathBuf::from("/tmp/prompt-dumps"), + provider: "openai".to_owned(), + model: "gpt-test".to_owned(), + base_url: Some("https://example.test/v1".to_owned()), + system_prompt: Some("assistant rule raw".to_owned()), + session_mode: Some("yolo".to_owned()), + skills: vec!["aionui-config".to_owned()], + mcp_servers, + runtime_env: vec![("AIONUI_RAW".to_owned(), "raw-env-value".to_owned())], + }; + let data = SendMessageData { + content: "team wake raw content".to_owned(), + msg_id: "msg-aionrs-final".to_owned(), + turn_id: Some("turn-aionrs-final".to_owned()), + files: Vec::new(), + inject_skills: Vec::new(), + }; + + let value = build_aionrs_final_input_dump_value("conv-aionrs", "/workspace", &context, &data); + + assert_eq!(value["kind"], "aionrs-final-input"); + assert_eq!(value["backend"], "aionrs"); + assert_eq!(value["conversation_id"], "conv-aionrs"); + assert_eq!(value["msg_id"], "msg-aionrs-final"); + assert_eq!(value["turn_id"], "turn-aionrs-final"); + assert_eq!(value["input"]["system_prompt"], "assistant rule raw"); + assert_eq!(value["input"]["user_content"], "team wake raw content"); + assert_eq!(value["resolved_context"]["provider"], "openai"); + assert_eq!(value["resolved_context"]["model"], "gpt-test"); + assert_eq!(value["resolved_context"]["workspace"]["path"], "/workspace"); + assert_eq!(value["resolved_context"]["skills"][0], "aionui-config"); + assert_eq!( + value["resolved_context"]["mcp_servers"]["raw-mcp"]["env"]["TOKEN"], + "raw-token-value" + ); + assert_eq!(value["resolved_context"]["runtime_env"][0][1], "raw-env-value"); + } + #[tokio::test] async fn aionrs_agent_returns_correct_type() { let agent = AionrsAgentManager::new("conv-1".into(), "/project".into(), make_test_config(), None) diff --git a/crates/aionui-ai-agent/src/protocol/error.rs b/crates/aionui-ai-agent/src/protocol/error.rs index ea244ccd3..911557b05 100644 --- a/crates/aionui-ai-agent/src/protocol/error.rs +++ b/crates/aionui-ai-agent/src/protocol/error.rs @@ -274,11 +274,11 @@ impl AcpError { /// Convert an SDK [`Error`](SdkError) into an [`AcpError`]. /// /// Mapping is by [`ErrorCode`], never by message text. The single - /// exception is `data.error == "Session not found: ..."`: OpenCode - /// (and likely others) return a stale-session failure as - /// `code = InvalidParams (-32602)` with the real reason buried in - /// the `data` field, so we re-classify those into `SessionNotFound` - /// to keep crash detection / recovery paths uniform across agents. + /// exceptions are known stale-session shapes: `data.error == "Session not + /// found: ..."` and resource-not-found replies from `session/load` + /// without a concrete resource URI. OpenCode/Codex have emitted stale + /// resume failures with those shapes, so we re-classify them into + /// `SessionNotFound` to keep recovery paths uniform across agents. /// See ELECTRON-1HQ. /// `context` carries the session ID or method name for diagnostics. pub fn from_sdk(err: SdkError, context: &str) -> Self { @@ -289,6 +289,10 @@ impl AcpError { ErrorCode::ResourceNotFound => { if let Some(sid) = extract_session_not_found(err.data.as_ref()) { AcpError::SessionNotFound { session_id: sid } + } else if extract_resource_not_found(err.data.as_ref()).is_none() && is_session_load_method(context) { + AcpError::SessionNotFound { + session_id: context.to_owned(), + } } else { AcpError::ResourceNotFound { resource: extract_resource_not_found(err.data.as_ref()), @@ -339,6 +343,10 @@ impl AcpError { } } +fn is_session_load_method(context: &str) -> bool { + context == "session/load" +} + /// If `data` carries a `{"error": "Session not found: "}` payload /// (either as a JSON object or as a JSON-string-of-JSON, which is what /// OpenCode actually emits — see ELECTRON-1HQ), return the session id. @@ -544,6 +552,26 @@ mod tests { } } + #[test] + fn from_sdk_session_load_resource_not_found_without_uri_is_session_not_found() { + let sdk_err = SdkError::resource_not_found(None); + let acp = AcpError::from_sdk(sdk_err, "session/load"); + match acp { + AcpError::SessionNotFound { session_id } => assert_eq!(session_id, "session/load"), + other => panic!("expected SessionNotFound, got {other:?}"), + } + } + + #[test] + fn from_sdk_prompt_resource_not_found_without_uri_stays_resource_not_found() { + let sdk_err = SdkError::resource_not_found(None); + let acp = AcpError::from_sdk(sdk_err, "session/prompt"); + assert!( + matches!(acp, AcpError::ResourceNotFound { resource: None, .. }), + "prompt ResourceNotFound must not clear a persisted session id: {acp:?}" + ); + } + #[test] fn from_sdk_method_not_found() { let sdk_err = SdkError::method_not_found(); diff --git a/crates/aionui-ai-agent/src/services/provider_health.rs b/crates/aionui-ai-agent/src/services/provider_health.rs index d1f0b0eee..ac41c94e1 100644 --- a/crates/aionui-ai-agent/src/services/provider_health.rs +++ b/crates/aionui-ai-agent/src/services/provider_health.rs @@ -88,9 +88,11 @@ impl ProviderHealthCheckService { compat_overrides, session_directory: self.data_dir.join("aionrs-health-check-sessions"), session_mode: None, + skills: Vec::new(), extra_mcp_servers: HashMap::new(), bedrock_config, runtime_env: Vec::new(), + prompt_dump_dir: None, }) } } diff --git a/crates/aionui-ai-agent/src/types.rs b/crates/aionui-ai-agent/src/types.rs index dff8af4dd..d43c8abfc 100644 --- a/crates/aionui-ai-agent/src/types.rs +++ b/crates/aionui-ai-agent/src/types.rs @@ -131,12 +131,16 @@ pub struct AionrsResolvedConfig { pub session_directory: PathBuf, /// Session mode (default, auto_edit, yolo). pub session_mode: Option, + /// Resolved skill names from the conversation snapshot. + pub skills: Vec, /// Extra MCP servers to inject (team coordination or guide). pub extra_mcp_servers: HashMap, /// AWS Bedrock credentials (region + access key or profile). pub bedrock_config: Option, /// Per-turn environment values exposed to runtime tool execution. pub runtime_env: Vec<(String, String)>, + /// Prompt dump directory when development prompt dumps are enabled. + pub prompt_dump_dir: Option, } #[cfg(test)] diff --git a/crates/aionui-ai-agent/tests/agent_types_integration.rs b/crates/aionui-ai-agent/tests/agent_types_integration.rs index 0abce7526..e3ea67f4d 100644 --- a/crates/aionui-ai-agent/tests/agent_types_integration.rs +++ b/crates/aionui-ai-agent/tests/agent_types_integration.rs @@ -102,9 +102,11 @@ fn make_aionrs_config() -> AionrsResolvedConfig { compat_overrides: Default::default(), session_directory: std::env::temp_dir().join("aionrs-test-sessions"), session_mode: None, + skills: Vec::new(), extra_mcp_servers: Default::default(), bedrock_config: None, runtime_env: Vec::new(), + prompt_dump_dir: None, } } diff --git a/crates/aionui-api-types/src/lib.rs b/crates/aionui-api-types/src/lib.rs index 670f1dc58..4acca5e33 100644 --- a/crates/aionui-api-types/src/lib.rs +++ b/crates/aionui-api-types/src/lib.rs @@ -158,12 +158,12 @@ pub use system::{ pub use team::{ AddAgentRequest, CancelTeamChildTurnRequest, CancelTeamRunRequest, CreateTeamRequest, PauseTeamSlotRequest, RenameAgentRequest, RenameTeamRequest, SendAgentMessageRequest, SendTeamMessageRequest, TeamAgentInput, - TeamAgentRemovedPayload, TeamAgentRenamedPayload, TeamAgentResponse, TeamAgentSpawnedPayload, - TeamAgentStatusPayload, TeamChildTurnPayload, TeamListResponse, TeamMcpPhase, TeamMcpRuntimeConfig, - TeamMcpStatusPayload, TeamResponse, TeamRunAckResponse, TeamRunPayload, TeamRunSource, TeamRunStateResponse, - TeamRunStatus, TeamRunTargetRole, TeamRuntimeSeed, TeamSendMessageDelivery, TeamSendMessageQueuedResponse, - TeamSendMessageReason, TeamSendMessageStatus, TeamSendMessageTargetQueueState, TeamSessionBinding, - TeamSlotRuntimeHealth, TeamSlotWorkPayload, TeammateMessagePayload, + TeamAgentRemovedPayload, TeamAgentRenamedPayload, TeamAgentResponse, TeamAgentRuntimeStatus, + TeamAgentRuntimeStatusPayload, TeamAgentSpawnedPayload, TeamAgentStatusPayload, TeamChildTurnPayload, + TeamListResponse, TeamMcpRuntimeConfig, TeamMessageEnqueueStatus, TeamResponse, TeamRunAckResponse, TeamRunPayload, + TeamRunSource, TeamRunStateResponse, TeamRunStatus, TeamRunTargetRole, TeamRuntimeSeed, + TeamSendMessageQueuedResponse, TeamSessionBinding, TeamSessionPhase, TeamSessionStatus, TeamSessionStatusPayload, + TeamSlotBlockedReason, TeamSlotWorkPayload, TeamSlotWorkState, TeammateMessagePayload, }; pub use team_mcp::{TEAM_MCP_SERVER_NAME, TeamMcpStdioConfig}; pub use websocket::WebSocketMessage; diff --git a/crates/aionui-api-types/src/team.rs b/crates/aionui-api-types/src/team.rs index 07637e000..5718721da 100644 --- a/crates/aionui-api-types/src/team.rs +++ b/crates/aionui-api-types/src/team.rs @@ -9,8 +9,9 @@ use crate::TeamMcpStdioConfig; /// Input for a single agent when creating a team or adding an agent. /// -/// Each agent gets its own conversation; the first agent in a create -/// request becomes the team lead. +/// Each agent gets its own conversation. Create requests must include exactly +/// one agent with role `lead` or `leader`; that explicit role becomes the team +/// lead. /// #[derive(Debug, Clone)] pub struct TeamAgentInput { @@ -67,7 +68,7 @@ impl<'de> Deserialize<'de> for TeamAgentInput { /// Request body for `POST /api/teams`. /// /// Creates a team with the given name and agent list. -/// The first agent in the array is designated as the lead. +/// Exactly one agent with role `lead` or `leader` is designated as the lead. #[derive(Debug, Deserialize)] pub struct CreateTeamRequest { pub name: String, @@ -287,7 +288,6 @@ pub enum TeamRunStatus { #[serde(rename_all = "snake_case")] pub enum TeamRunSource { UserMessage, - RecoveryDrain, } #[derive(Debug, Deserialize)] @@ -312,48 +312,53 @@ pub struct PauseTeamSlotRequest { #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct TeamRunAckResponse { - pub team_run_id: String, - pub team_id: String, - pub source: TeamRunSource, - pub has_user_intervention: bool, - pub target_slot_id: String, - pub target_role: TeamRunTargetRole, - pub accepted_slot_id: String, - pub accepted_role: TeamRunTargetRole, - pub status: TeamRunStatus, - #[serde(skip_serializing_if = "Option::is_none")] - pub message_id: Option, + pub enqueue_status: TeamMessageEnqueueStatus, + pub message_id: String, + pub run: TeamRunPayload, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum TeamSlotWorkState { + Idle, + Queued, + Starting, + Running, + Paused, + Blocked, } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] #[serde(rename_all = "snake_case")] -pub enum TeamSlotRuntimeHealth { - Disconnected, - Unhealthy, +pub enum TeamSlotBlockedReason { + RuntimeStarting, + RuntimeFailed, + Removing, + SessionStopped, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum TeamMessageEnqueueStatus { + Accepted, + Queued, + BlockedRuntimeStarting, } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct TeamSlotWorkPayload { pub slot_id: String, pub role: TeamRunTargetRole, - pub pending_wake_count: usize, - pub starting_child_count: usize, - #[serde(default)] - pub paused: bool, - #[serde(default)] - pub suppressed_wake_count: usize, - #[serde(default, skip_serializing_if = "Option::is_none")] + pub state: TeamSlotWorkState, + pub queued_foreground_count: usize, + pub queued_background_count: usize, pub active_turn_id: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] pub active_turn_started_at_ms: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] pub active_turn_elapsed_ms: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] pub active_turn_slow: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] pub active_turn_slow_threshold_ms: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub runtime_health: Option, + pub blocked_reason: Option, + pub team_run_id: Option, } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] @@ -365,16 +370,18 @@ pub struct TeamRunPayload { pub target_slot_id: String, pub target_role: TeamRunTargetRole, pub status: TeamRunStatus, - pub active_child_count: usize, - pub pending_wake_count: usize, - pub starting_child_count: usize, - #[serde(default)] + pub queued_intent_count: usize, + pub starting_batch_count: usize, + pub running_batch_count: usize, + pub active_enqueue_lease_count: usize, pub slot_work: Vec, } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct TeamRunStateResponse { + pub session_generation: Option, pub active_run: Option, + pub slot_work: Vec, } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] @@ -388,55 +395,10 @@ pub struct TeamChildTurnPayload { pub status: TeamRunStatus, } -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "snake_case")] -pub enum TeamSendMessageStatus { - Queued, - Rejected, - Error, -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "snake_case")] -pub enum TeamSendMessageDelivery { - WakeRecorded, - WakeSuppressed, - NotRecorded, -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "snake_case")] -pub enum TeamSendMessageReason { - QueuedForIdle, - BehindStartingTurn, - BehindActiveTurn, - SuppressedByPause, - NoActiveTeamRun, - TargetNotFound, - TargetDisconnected, - TargetUnhealthy, - InternalError, -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -pub struct TeamSendMessageTargetQueueState { - pub slot_id: String, - pub role: TeamRunTargetRole, - pub queue_state: TeamSendMessageReason, - pub pending_wake_count: usize, - pub starting_child_count: usize, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub active_turn_id: Option, - pub suppressed_wake_count: usize, -} - #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct TeamSendMessageQueuedResponse { - pub status: TeamSendMessageStatus, - pub delivery: TeamSendMessageDelivery, - pub reason: TeamSendMessageReason, - pub team_run_id: String, - pub targets: Vec, + pub team_run_id: Option, + pub target: TeamSlotWorkPayload, } // --------------------------------------------------------------------------- @@ -536,39 +498,61 @@ pub struct TeamAgentRenamedPayload { pub name: String, } -/// Lifecycle phases of the per-team MCP stdio bridge + ACP session. -/// -/// Emitted by the MCP supervisor whenever a teammate slot transitions -/// through its bring-up / degraded / ready states so the frontend can -/// surface actionable status for each agent. +/// Runtime attach/warmup status for a team agent. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] #[serde(rename_all = "snake_case")] -pub enum TeamMcpPhase { - TcpReady, - TcpError, - SessionInjecting, - SessionReady, - SessionError, - LoadFailed, - Degraded, - ConfigWriteFailed, - McpToolsWaiting, - McpToolsReady, +pub enum TeamAgentRuntimeStatus { + Pending, + Ready, + Failed, } -/// Payload for `team.mcpStatus` WebSocket event. +/// Payload for `team.agentRuntimeStatusChanged` WebSocket event. /// -/// Pushed whenever a teammate's MCP bridge or ACP session transitions to -/// a new [`TeamMcpPhase`]. Optional fields carry phase-specific detail: -/// `port` for TCP bring-up, `server_count` for tool readiness, `error` -/// for failure phases. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct TeamMcpStatusPayload { +/// Pushed when a team member's runtime attach/warmup lifecycle changes. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct TeamAgentRuntimeStatusPayload { pub team_id: String, pub slot_id: String, - pub phase: TeamMcpPhase, + pub conversation_id: String, + pub status: TeamAgentRuntimeStatus, + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, +} + +/// Team-level session availability status for `team.sessionStatusChanged`. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(rename_all = "snake_case")] +pub enum TeamSessionStatus { + Starting, + Ready, + Failed, +} + +/// Diagnostic phase for team session startup. +/// +/// Frontend gating must use [`TeamSessionStatus`]. This phase identifies +/// which ensure-session step is currently running or failed. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(rename_all = "snake_case")] +pub enum TeamSessionPhase { + LoadingTeam, + StartingBridge, + AttachingAgents, + Recovering, +} + +/// Payload for `team.sessionStatusChanged` WebSocket event. +/// +/// Pushed when the whole team session moves through startup, ready, or +/// failed states. Member runtime details are reported separately through +/// `team.agentRuntimeStatusChanged`. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct TeamSessionStatusPayload { + pub team_id: String, + pub status: TeamSessionStatus, #[serde(skip_serializing_if = "Option::is_none")] - pub port: Option, + pub phase: Option, #[serde(skip_serializing_if = "Option::is_none")] pub server_count: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -1249,101 +1233,69 @@ mod tests { assert_eq!(team.created_at, 1000); } - // -- F. TeamMcpPhase serde roundtrip -------------------------------------- + // -- F. TeamSessionStatus / TeamSessionPhase serde roundtrip ----------------- - fn assert_phase_roundtrip(phase: TeamMcpPhase, wire: &str) { - let json = serde_json::to_value(&phase).unwrap(); + fn assert_session_status_roundtrip(status: TeamSessionStatus, wire: &str) { + let json = serde_json::to_value(&status).unwrap(); assert_eq!(json, serde_json::Value::String(wire.into())); - let parsed: TeamMcpPhase = serde_json::from_value(json).unwrap(); - assert_eq!(parsed, phase); - } - - #[test] - fn team_mcp_phase_tcp_ready_roundtrip() { - assert_phase_roundtrip(TeamMcpPhase::TcpReady, "tcp_ready"); - } - - #[test] - fn team_mcp_phase_tcp_error_roundtrip() { - assert_phase_roundtrip(TeamMcpPhase::TcpError, "tcp_error"); - } - - #[test] - fn team_mcp_phase_session_injecting_roundtrip() { - assert_phase_roundtrip(TeamMcpPhase::SessionInjecting, "session_injecting"); - } - - #[test] - fn team_mcp_phase_session_ready_roundtrip() { - assert_phase_roundtrip(TeamMcpPhase::SessionReady, "session_ready"); - } - - #[test] - fn team_mcp_phase_session_error_roundtrip() { - assert_phase_roundtrip(TeamMcpPhase::SessionError, "session_error"); - } - - #[test] - fn team_mcp_phase_load_failed_roundtrip() { - assert_phase_roundtrip(TeamMcpPhase::LoadFailed, "load_failed"); - } - - #[test] - fn team_mcp_phase_degraded_roundtrip() { - assert_phase_roundtrip(TeamMcpPhase::Degraded, "degraded"); + let parsed: TeamSessionStatus = serde_json::from_value(json).unwrap(); + assert_eq!(parsed, status); } - #[test] - fn team_mcp_phase_config_write_failed_roundtrip() { - assert_phase_roundtrip(TeamMcpPhase::ConfigWriteFailed, "config_write_failed"); + fn assert_session_phase_roundtrip(phase: TeamSessionPhase, wire: &str) { + let json = serde_json::to_value(&phase).unwrap(); + assert_eq!(json, serde_json::Value::String(wire.into())); + let parsed: TeamSessionPhase = serde_json::from_value(json).unwrap(); + assert_eq!(parsed, phase); } #[test] - fn team_mcp_phase_mcp_tools_waiting_roundtrip() { - assert_phase_roundtrip(TeamMcpPhase::McpToolsWaiting, "mcp_tools_waiting"); + fn team_session_status_roundtrips() { + assert_session_status_roundtrip(TeamSessionStatus::Starting, "starting"); + assert_session_status_roundtrip(TeamSessionStatus::Ready, "ready"); + assert_session_status_roundtrip(TeamSessionStatus::Failed, "failed"); } #[test] - fn team_mcp_phase_mcp_tools_ready_roundtrip() { - assert_phase_roundtrip(TeamMcpPhase::McpToolsReady, "mcp_tools_ready"); + fn team_session_phase_roundtrips() { + assert_session_phase_roundtrip(TeamSessionPhase::LoadingTeam, "loading_team"); + assert_session_phase_roundtrip(TeamSessionPhase::StartingBridge, "starting_bridge"); + assert_session_phase_roundtrip(TeamSessionPhase::AttachingAgents, "attaching_agents"); + assert_session_phase_roundtrip(TeamSessionPhase::Recovering, "recovering"); } - // -- G. TeamMcpStatusPayload & TeammateMessagePayload --------------------- + // -- G. TeamSessionStatusPayload & TeammateMessagePayload -------------------- #[test] - fn serialize_team_mcp_status_payload_all_fields_present() { - let payload = TeamMcpStatusPayload { + fn serialize_team_session_status_payload_all_fields_present() { + let payload = TeamSessionStatusPayload { team_id: "team-1".into(), - slot_id: "slot-2".into(), - phase: TeamMcpPhase::SessionReady, - port: Some(54321), + status: TeamSessionStatus::Ready, + phase: Some(TeamSessionPhase::Recovering), server_count: Some(7), error: Some("boom".into()), }; let json = serde_json::to_value(&payload).unwrap(); assert_eq!(json["team_id"], "team-1"); - assert_eq!(json["slot_id"], "slot-2"); - assert_eq!(json["phase"], "session_ready"); - assert_eq!(json["port"], 54321); + assert_eq!(json["status"], "ready"); + assert_eq!(json["phase"], "recovering"); assert_eq!(json["server_count"], 7); assert_eq!(json["error"], "boom"); } #[test] - fn serialize_team_mcp_status_payload_optional_fields_omitted() { - let payload = TeamMcpStatusPayload { + fn serialize_team_session_status_payload_optional_fields_omitted() { + let payload = TeamSessionStatusPayload { team_id: "team-1".into(), - slot_id: "slot-2".into(), - phase: TeamMcpPhase::TcpReady, - port: None, + status: TeamSessionStatus::Starting, + phase: None, server_count: None, error: None, }; let json = serde_json::to_value(&payload).unwrap(); assert_eq!(json["team_id"], "team-1"); - assert_eq!(json["slot_id"], "slot-2"); - assert_eq!(json["phase"], "tcp_ready"); - assert!(json.get("port").is_none()); + assert_eq!(json["status"], "starting"); + assert!(json.get("phase").is_none()); assert!(json.get("server_count").is_none()); assert!(json.get("error").is_none()); } @@ -1363,328 +1315,130 @@ mod tests { assert_eq!(json["from_name"], "Lead"); } - #[test] - fn team_run_ack_serializes_snake_case_enums() { - let ack = TeamRunAckResponse { - team_run_id: "trun-1".into(), - team_id: "team-1".into(), - source: TeamRunSource::UserMessage, - has_user_intervention: true, - target_slot_id: "lead-1".into(), - target_role: TeamRunTargetRole::Lead, - accepted_slot_id: "lead-1".into(), - accepted_role: TeamRunTargetRole::Lead, - status: TeamRunStatus::Accepted, - message_id: Some("msg-1".into()), - }; - - let value = serde_json::to_value(&ack).unwrap(); - assert_eq!(value["target_role"], "lead"); - assert_eq!(value["status"], "accepted"); - assert_eq!(value["message_id"], "msg-1"); - } - - #[test] - fn team_run_ack_distinguishes_initial_target_from_accepted_slot() { - let ack = TeamRunAckResponse { - team_run_id: "trun-1".into(), - team_id: "team-1".into(), - source: TeamRunSource::UserMessage, - has_user_intervention: true, - target_slot_id: "lead-1".into(), - target_role: TeamRunTargetRole::Lead, - accepted_slot_id: "worker-1".into(), - accepted_role: TeamRunTargetRole::Teammate, - status: TeamRunStatus::Accepted, - message_id: Some("msg-1".into()), - }; - - let value = serde_json::to_value(&ack).unwrap(); - assert_eq!(value["target_slot_id"], "lead-1"); - assert_eq!(value["target_role"], "lead"); - assert_eq!(value["accepted_slot_id"], "worker-1"); - assert_eq!(value["accepted_role"], "teammate"); - } - - #[test] - fn team_run_source_serializes_snake_case() { - let user = serde_json::to_value(TeamRunSource::UserMessage).unwrap(); - let recovery = serde_json::to_value(TeamRunSource::RecoveryDrain).unwrap(); - - assert_eq!(user, serde_json::json!("user_message")); - assert_eq!(recovery, serde_json::json!("recovery_drain")); + fn active_slot_work() -> TeamSlotWorkPayload { + TeamSlotWorkPayload { + slot_id: "lead-1".into(), + role: TeamRunTargetRole::Lead, + state: TeamSlotWorkState::Running, + queued_foreground_count: 2, + queued_background_count: 1, + active_turn_id: Some("turn-1".into()), + active_turn_started_at_ms: Some(10), + active_turn_elapsed_ms: Some(25), + active_turn_slow: Some(false), + active_turn_slow_threshold_ms: Some(600_000), + blocked_reason: None, + team_run_id: Some("run-1".into()), + } } - #[test] - fn team_run_payload_serializes_source_metadata() { - let payload = TeamRunPayload { + fn active_run_payload() -> TeamRunPayload { + TeamRunPayload { team_id: "team-1".into(), team_run_id: "run-1".into(), - source: TeamRunSource::RecoveryDrain, - has_user_intervention: false, - target_slot_id: "lead-1".into(), - target_role: TeamRunTargetRole::Lead, - status: TeamRunStatus::Accepted, - active_child_count: 0, - pending_wake_count: 1, - starting_child_count: 0, - slot_work: vec![], - }; - - let json = serde_json::to_value(payload).unwrap(); - assert_eq!(json["source"], "recovery_drain"); - assert_eq!(json["has_user_intervention"], false); - } - - #[test] - fn team_run_ack_serializes_source_metadata() { - let ack = TeamRunAckResponse { - team_run_id: "run-1".into(), - team_id: "team-1".into(), source: TeamRunSource::UserMessage, has_user_intervention: true, target_slot_id: "lead-1".into(), target_role: TeamRunTargetRole::Lead, - accepted_slot_id: "lead-1".into(), - accepted_role: TeamRunTargetRole::Lead, - status: TeamRunStatus::Accepted, - message_id: Some("mailbox-1".into()), - }; - - let json = serde_json::to_value(ack).unwrap(); - assert_eq!(json["source"], "user_message"); - assert_eq!(json["has_user_intervention"], true); - assert_eq!(json["message_id"], "mailbox-1"); - } - - fn active_run_payload() -> TeamRunPayload { - TeamRunPayload { - team_id: "team-1".to_owned(), - team_run_id: "run-1".to_owned(), - source: TeamRunSource::UserMessage, - has_user_intervention: true, - target_slot_id: "lead".to_owned(), - target_role: TeamRunTargetRole::Lead, status: TeamRunStatus::Running, - active_child_count: 1, - pending_wake_count: 0, - starting_child_count: 0, - slot_work: vec![TeamSlotWorkPayload { - slot_id: "lead".to_owned(), - role: TeamRunTargetRole::Lead, - pending_wake_count: 0, - starting_child_count: 0, - paused: false, - suppressed_wake_count: 0, - active_turn_id: Some("turn-1".to_owned()), - active_turn_started_at_ms: Some(10), - active_turn_elapsed_ms: Some(20), - active_turn_slow: Some(false), - active_turn_slow_threshold_ms: Some(600_000), - runtime_health: None, - }], + queued_intent_count: 2, + starting_batch_count: 0, + running_batch_count: 1, + active_enqueue_lease_count: 0, + slot_work: vec![active_slot_work()], } } #[test] - fn team_run_state_response_serializes_null_active_run() { - let value = serde_json::to_value(TeamRunStateResponse { active_run: None }).unwrap(); - assert_eq!(value, serde_json::json!({ "active_run": null })); + fn team_run_source_serializes_snake_case() { + assert_eq!( + serde_json::to_value(TeamRunSource::UserMessage).unwrap(), + serde_json::json!("user_message") + ); } #[test] - fn team_run_state_response_serializes_some_active_run() { - let payload = active_run_payload(); - let value = serde_json::to_value(TeamRunStateResponse { - active_run: Some(payload.clone()), - }) - .unwrap(); - - assert_eq!(value["active_run"]["team_id"], payload.team_id); - assert_eq!(value["active_run"]["team_run_id"], payload.team_run_id); - assert_eq!(value["active_run"]["source"], "user_message"); - assert_eq!(value["active_run"]["has_user_intervention"], true); - assert_eq!(value["active_run"]["status"], "running"); - assert_eq!(value["active_run"]["slot_work"][0]["slot_id"], "lead"); + fn team_slot_work_payload_serializes_authoritative_coordinator_state() { + assert_eq!( + serde_json::to_value(active_slot_work()).unwrap(), + serde_json::json!({ + "slot_id": "lead-1", + "role": "lead", + "state": "running", + "queued_foreground_count": 2, + "queued_background_count": 1, + "active_turn_id": "turn-1", + "active_turn_started_at_ms": 10, + "active_turn_elapsed_ms": 25, + "active_turn_slow": false, + "active_turn_slow_threshold_ms": 600000, + "blocked_reason": null, + "team_run_id": "run-1" + }) + ); } #[test] - fn team_run_payload_omits_sensitive_content() { - let payload = TeamRunPayload { - team_id: "team-1".into(), - team_run_id: "trun-1".into(), - source: TeamRunSource::UserMessage, - has_user_intervention: true, - target_slot_id: "worker-1".into(), - target_role: TeamRunTargetRole::Teammate, - status: TeamRunStatus::Running, - active_child_count: 1, - pending_wake_count: 0, - starting_child_count: 2, - slot_work: Vec::new(), - }; - - let value = serde_json::to_value(&payload).unwrap(); - assert_eq!(value["starting_child_count"], 2); + fn team_run_payload_serializes_authoritative_work_counts() { + let value = serde_json::to_value(active_run_payload()).unwrap(); + assert_eq!(value["queued_intent_count"], 2); + assert_eq!(value["running_batch_count"], 1); + assert_eq!(value["slot_work"][0]["state"], "running"); assert!(value.get("content").is_none()); - assert!(value.get("prompt").is_none()); - assert!(value.get("tool_input").is_none()); } #[test] - fn team_run_payload_includes_defaulted_slot_work() { - let payload = TeamRunPayload { - team_id: "team-1".into(), - team_run_id: "trun-1".into(), - source: TeamRunSource::UserMessage, - has_user_intervention: true, - target_slot_id: "lead-1".into(), - target_role: TeamRunTargetRole::Lead, - status: TeamRunStatus::Running, - active_child_count: 1, - pending_wake_count: 2, - starting_child_count: 1, - slot_work: vec![ - TeamSlotWorkPayload { - slot_id: "lead-1".into(), - role: TeamRunTargetRole::Lead, - pending_wake_count: 1, - starting_child_count: 0, - paused: false, - suppressed_wake_count: 0, - active_turn_id: Some("turn-lead".into()), - active_turn_started_at_ms: None, - active_turn_elapsed_ms: None, - active_turn_slow: None, - active_turn_slow_threshold_ms: None, - runtime_health: None, - }, - TeamSlotWorkPayload { - slot_id: "worker-1".into(), - role: TeamRunTargetRole::Teammate, - pending_wake_count: 1, - starting_child_count: 1, - paused: false, - suppressed_wake_count: 0, - active_turn_id: None, - active_turn_started_at_ms: None, - active_turn_elapsed_ms: None, - active_turn_slow: None, - active_turn_slow_threshold_ms: None, - runtime_health: None, - }, - ], + fn team_run_ack_embeds_the_complete_run_snapshot() { + let ack = TeamRunAckResponse { + enqueue_status: TeamMessageEnqueueStatus::Queued, + message_id: "message-1".into(), + run: active_run_payload(), }; - let value = serde_json::to_value(&payload).unwrap(); - assert_eq!(value["slot_work"][0]["slot_id"], "lead-1"); - assert_eq!(value["slot_work"][0]["active_turn_id"], "turn-lead"); - assert_eq!(value["slot_work"][1]["starting_child_count"], 1); - - let decoded: TeamRunPayload = serde_json::from_value(serde_json::json!({ - "team_id": "team-1", - "team_run_id": "trun-1", - "source": "user_message", - "has_user_intervention": true, - "target_slot_id": "lead-1", - "target_role": "lead", - "status": "running", - "active_child_count": 0, - "pending_wake_count": 0, - "starting_child_count": 0 - })) - .unwrap(); - assert!(decoded.slot_work.is_empty()); + let value = serde_json::to_value(ack).unwrap(); + assert_eq!(value["enqueue_status"], "queued"); + assert_eq!(value["message_id"], "message-1"); + assert_eq!(value["run"]["team_run_id"], "run-1"); + assert_eq!(value["run"]["slot_work"][0]["active_turn_id"], "turn-1"); } #[test] - fn team_slot_work_payload_includes_only_retained_pause_fields_when_non_default() { - let payload = TeamSlotWorkPayload { - slot_id: "lead-1".into(), - role: TeamRunTargetRole::Lead, - pending_wake_count: 0, - starting_child_count: 0, - active_turn_id: None, - paused: true, - suppressed_wake_count: 2, - active_turn_started_at_ms: None, - active_turn_elapsed_ms: None, - active_turn_slow: None, - active_turn_slow_threshold_ms: None, - runtime_health: None, - }; + fn team_run_state_response_keeps_global_slot_work_without_active_run() { + let value = serde_json::to_value(TeamRunStateResponse { + session_generation: Some("generation-1".into()), + active_run: None, + slot_work: vec![active_slot_work()], + }) + .unwrap(); - let value = serde_json::to_value(&payload).unwrap(); - assert_eq!(value["paused"], true); - assert_eq!(value["suppressed_wake_count"], 2); - assert!(value.get(format!("{}_pending_count", "foreground")).is_none()); - assert!(value.get(format!("{}_pending_count", "background")).is_none()); + assert_eq!(value["session_generation"], "generation-1"); + assert!(value["active_run"].is_null()); + assert_eq!(value["slot_work"][0]["slot_id"], "lead-1"); } #[test] - fn team_slot_work_payload_serializes_active_turn_slow_fields() { - let payload = TeamSlotWorkPayload { - slot_id: "worker-1".into(), - role: TeamRunTargetRole::Teammate, - pending_wake_count: 0, - starting_child_count: 0, - paused: false, - suppressed_wake_count: 0, - active_turn_id: Some("turn-worker".into()), - active_turn_started_at_ms: Some(1_000), - active_turn_elapsed_ms: Some(600_001), - active_turn_slow: Some(true), - active_turn_slow_threshold_ms: Some(600_000), - runtime_health: Some(TeamSlotRuntimeHealth::Unhealthy), - }; - - let value = serde_json::to_value(payload).unwrap(); - - assert_eq!(value["active_turn_started_at_ms"], 1_000); - assert_eq!(value["active_turn_elapsed_ms"], 600_001); - assert_eq!(value["active_turn_slow"], true); - assert_eq!(value["active_turn_slow_threshold_ms"], 600_000); - assert_eq!(value["runtime_health"], "unhealthy"); + fn team_slot_work_block_reasons_serialize_snake_case() { + assert_eq!( + serde_json::to_value(TeamSlotBlockedReason::RuntimeStarting).unwrap(), + serde_json::json!("runtime_starting") + ); + assert_eq!( + serde_json::to_value(TeamSlotBlockedReason::SessionStopped).unwrap(), + serde_json::json!("session_stopped") + ); } #[test] - fn team_send_message_queued_response_serializes_stable_contract() { + fn team_send_message_queued_response_serializes_authoritative_target() { let response = TeamSendMessageQueuedResponse { - status: TeamSendMessageStatus::Queued, - delivery: TeamSendMessageDelivery::WakeRecorded, - reason: TeamSendMessageReason::BehindActiveTurn, - team_run_id: "run-1".into(), - targets: vec![TeamSendMessageTargetQueueState { - slot_id: "worker-1".into(), - role: TeamRunTargetRole::Teammate, - queue_state: TeamSendMessageReason::BehindActiveTurn, - pending_wake_count: 1, - starting_child_count: 0, - active_turn_id: Some("turn-worker".into()), - suppressed_wake_count: 0, - }], + team_run_id: Some("run-1".into()), + target: active_slot_work(), }; - let value = serde_json::to_value(response).unwrap(); - assert_eq!(value["status"], "queued"); - assert_eq!(value["delivery"], "wake_recorded"); - assert_eq!(value["reason"], "behind_active_turn"); - assert_eq!(value["targets"][0]["queue_state"], "behind_active_turn"); - } - - #[test] - fn team_slot_work_payload_defaults_pause_fields_for_old_payloads() { - let decoded: TeamSlotWorkPayload = serde_json::from_value(serde_json::json!({ - "slot_id": "worker-1", - "role": "teammate", - "pending_wake_count": 0, - "starting_child_count": 0 - })) - .unwrap(); - - assert!(!decoded.paused); - assert_eq!(decoded.suppressed_wake_count, 0); + assert_eq!(value["team_run_id"], "run-1"); + assert_eq!(value["target"]["state"], "running"); + assert_eq!(value["target"]["queued_foreground_count"], 2); } #[test] diff --git a/crates/aionui-app/src/commands/cmd_server.rs b/crates/aionui-app/src/commands/cmd_server.rs index 806fd2dad..da48c239c 100644 --- a/crates/aionui-app/src/commands/cmd_server.rs +++ b/crates/aionui-app/src/commands/cmd_server.rs @@ -3,6 +3,7 @@ use std::io::{self, Write}; use std::net::SocketAddr; use std::process::ExitCode; +use std::sync::Arc; use std::time::{Duration, Instant}; use std::{future::Future, pin::Pin}; @@ -10,8 +11,9 @@ use tokio::net::TcpListener; use tracing::{info, warn}; use aionui_api_types::{RuntimeStatusScope, RuntimeStatusScopeKind}; -use aionui_app::{AppConfig, AppServices, RouterBuildError, create_router}; +use aionui_app::{AppConfig, AppServices, RouterBuildError, create_router_with_runtime}; use aionui_system::RuntimePrepareService; +use aionui_team::TeamIdleCleanupCoordinator; use crate::bootstrap::{BootstrapError, BootstrapErrorCode, ParentExitSignal, ServerEnvironment}; @@ -218,7 +220,7 @@ pub(crate) async fn run_server( info!("No configured users detected — initial setup required via /api/auth/status"); } - let router = create_router(&services) + let (router, router_runtime) = create_router_with_runtime(&services) .await .map_err(router_build_error_to_bootstrap)?; info!( @@ -271,8 +273,18 @@ pub(crate) async fn run_server( // propagates graceful-shutdown so the scanner exits on SIGINT/SIGTERM. let (shutdown_tx, shutdown_rx) = tokio::sync::watch::channel(false); let (shutdown_error_tx, shutdown_error_rx) = tokio::sync::oneshot::channel::(); - let idle_scanner_handle = - aionui_ai_agent::start_idle_scanner(services.worker_task_manager.clone(), shutdown_rx, None, None); + let idle_cleanup_coordinator: Arc = + Arc::new(TeamIdleCleanupCoordinator::new( + router_runtime.team_service.clone(), + services.active_lease_registry.clone(), + )); + let idle_scanner_handle = aionui_ai_agent::start_idle_scanner_with_coordinator( + services.worker_task_manager.clone(), + shutdown_rx, + None, + None, + Some(idle_cleanup_coordinator), + ); let conversation_runtime_state = services.conversation_runtime_state.clone(); let worker_task_manager = services.worker_task_manager.clone(); diff --git a/crates/aionui-app/src/commands/cmd_team_stdio.rs b/crates/aionui-app/src/commands/cmd_team_stdio.rs index 4db050fb7..1839d2f11 100644 --- a/crates/aionui-app/src/commands/cmd_team_stdio.rs +++ b/crates/aionui-app/src/commands/cmd_team_stdio.rs @@ -130,15 +130,9 @@ struct SendMessageParams { struct SpawnAgentParams { /// Agent display name. name: String, - /// Model override for the new agent. - #[serde(default)] - model: Option, /// Assistant identifier from the available assistants catalog. #[serde(default)] assistant_id: Option, - /// Agent role (default: "teammate"). - #[serde(default)] - role: Option, } #[derive(Deserialize, schemars::JsonSchema)] @@ -174,6 +168,50 @@ struct TaskUpdateParams { blocked_by: Option>, } +#[derive(Debug, Deserialize, schemars::JsonSchema)] +#[serde(untagged)] +enum TaskListStatusParam { + Single(String), + Many(Vec), +} + +#[derive(Debug, Deserialize, schemars::JsonSchema)] +#[serde(deny_unknown_fields)] +struct TaskListParams { + /// Return only tasks owned by this exact agent slot_id. + #[serde(default)] + owner: Option, + /// Return only tasks with the requested status or statuses. + #[serde(default)] + status: Option, + /// Include deleted tasks when status is omitted. Defaults server-side to true. + #[serde(default)] + include_deleted: Option, + /// Maximum number of returned tasks. The TCP server validates and clamps. + #[serde(default)] + limit: Option, +} + +impl TaskListParams { + fn into_json(self) -> serde_json::Value { + let status = match self.status { + Some(TaskListStatusParam::Single(value)) => serde_json::json!(value), + Some(TaskListStatusParam::Many(values)) => serde_json::json!(values), + None => serde_json::Value::Null, + }; + let mut args = serde_json::json!({ + "owner": self.owner, + "status": status, + "include_deleted": self.include_deleted, + "limit": self.limit, + }); + args.as_object_mut() + .expect("task list params must serialize to object") + .retain(|_, value| !value.is_null()); + args + } +} + #[derive(Deserialize, schemars::JsonSchema)] struct RenameAgentParams { /// Agent slot_id to rename. @@ -191,13 +229,6 @@ struct ShutdownAgentParams { reason: Option, } -#[derive(Deserialize, schemars::JsonSchema)] -struct ListModelsParams { - /// Assistant ID to query. Shows all backends when omitted. - #[serde(default)] - assistant_id: Option, -} - #[derive(Deserialize, schemars::JsonSchema)] #[serde(deny_unknown_fields)] struct DescribeAssistantParams { @@ -228,16 +259,14 @@ impl TeamStdioServer { #[tool( name = "team_spawn_agent", - description = "Create a new teammate agent to join the team.\n\nUse this only when one of the following is true:\n- The user explicitly approved the proposed teammate lineup in a previous message\n- The user explicitly instructed you to create a specific teammate immediately\n\nBefore calling this tool in the normal planning flow:\n- Start with one short sentence explaining why additional teammates would help\n- Tell the user which teammate(s) you recommend\n- Present the proposal as a table with: name, responsibility, recommended assistant, and recommended model\n- Include each teammate's responsibility, recommended assistant, and model\n- Ask whether to create them as proposed or change any names, responsibilities, or assistant choices\n- In that approval question, remind the user that they can later ask you to replace or adjust any teammate if the lineup is not working well\n- Do NOT call this tool in that same turn; wait for explicit approval in a later user message\n\nWhen calling this tool, always provide assistant_id from the available assistants catalog.\nWhen calling this tool, provide the model parameter if a specific model was recommended and approved.\n\nThe new agent will be created and added to the team. You can then assign tasks and send messages to it." + description = "Create a new teammate agent to join the team.\n\nUse this only when one of the following is true:\n- The user explicitly approved the proposed teammate lineup in a previous message\n- The user explicitly instructed you to create a specific teammate immediately\n\nBefore calling this tool in the normal planning flow:\n- Start with one short sentence explaining why additional teammates would help\n- Tell the user which teammate(s) you recommend\n- Present the proposal as a table with: name, responsibility, and recommended assistant\n- Include each teammate's responsibility and recommended assistant\n- Ask whether to create them as proposed or change any names, responsibilities, or assistant choices\n- In that approval question, remind the user that they can later ask you to replace or adjust any teammate if the lineup is not working well\n- Do NOT call this tool in that same turn; wait for explicit approval in a later user message\n\nWhen calling this tool, always provide assistant_id from the available assistants catalog.\nDo not provide a model. The new teammate uses the selected assistant's configured/default model; users can adjust models from the UI model selector.\n\nThe new agent will be created and added to the team. You can then assign tasks and send messages to it." )] async fn spawn_agent(&self, Parameters(params): Parameters) -> CallToolResult { self.forward_to_tcp( "team_spawn_agent", &serde_json::json!({ "name": params.name, - "model": params.model, "assistant_id": params.assistant_id, - "role": params.role, }), ) .await @@ -275,9 +304,12 @@ impl TeamStdioServer { .await } - #[tool(name = "team_task_list", description = "List all tasks on the team task board.")] - async fn task_list(&self) -> CallToolResult { - self.forward_to_tcp("team_task_list", &serde_json::json!({})).await + #[tool( + name = "team_task_list", + description = "List tasks on the team task board. Pass {} for the full board, or use owner/status/include_deleted/limit for filtered views." + )] + async fn task_list(&self, Parameters(params): Parameters) -> CallToolResult { + self.forward_to_tcp("team_task_list", ¶ms.into_json()).await } #[tool( @@ -318,21 +350,9 @@ impl TeamStdioServer { .await } - #[tool( - name = "team_list_models", - description = "Query available models for assistant backends. Returns the real-time model list that matches the frontend model selector.\n\nUse this to:\n- Check what models are available before spawning an assistant-backed teammate with a specific model\n- See all available backends and their models at once\n- Verify a model ID is valid for the backend behind a chosen assistant or fallback backend\n\nPass assistant_id to query models for a specific assistant, or omit it to see all backends." - )] - async fn list_models(&self, Parameters(params): Parameters) -> CallToolResult { - self.forward_to_tcp( - "team_list_models", - &serde_json::json!({ "assistant_id": params.assistant_id }), - ) - .await - } - #[tool( name = "team_describe_assistant", - description = "Get detailed information about an assistant before spawning it as a teammate.\n\nReturns the assistant's full description, enabled skills, and example tasks so you can\njudge whether it fits the user's request. Use this when two or more assistants look\nrelevant from the one-line catalog in your system prompt.\n\nOnly works on assistants listed in \"Available Assistants for Spawning\".\nAfter confirming a match, call team_spawn_agent with the same assistant_id." + description = "Get detailed information about an assistant before spawning it as a teammate.\n\nReturns the assistant's full description, enabled skills, and example tasks so you can\njudge whether it fits the user's request. Use this when two or more assistants look\nrelevant from the one-line catalog in your system prompt.\n\nUse team_list_assistants to find candidate assistant_id values.\nAfter confirming a match, call team_spawn_agent with the same assistant_id." )] async fn describe_assistant(&self, Parameters(params): Parameters) -> CallToolResult { self.forward_to_tcp( @@ -705,6 +725,47 @@ mod tests { assert!(err.to_string().contains("custom_agent_id")); } + #[test] + fn task_list_params_accept_filter_arguments() { + let parsed = serde_json::from_value::(json!({ + "owner": "worker-1", + "status": ["pending", "in_progress"], + "include_deleted": false, + "limit": 50 + })) + .expect("task_list filters should parse"); + + let forwarded = parsed.into_json(); + assert_eq!(forwarded["owner"], "worker-1"); + assert_eq!(forwarded["status"], json!(["pending", "in_progress"])); + assert_eq!(forwarded["include_deleted"], false); + assert_eq!(forwarded["limit"], 50); + } + + #[test] + fn task_list_params_reject_unknown_fields() { + let parsed = serde_json::from_value::(json!({ + "slot_id": "worker-1" + })); + assert!(parsed.is_err()); + assert!(parsed.unwrap_err().to_string().contains("unknown field")); + } + + #[test] + fn team_stdio_task_list_schema_exposes_filters() { + let router = TeamStdioServer::tool_router(); + let tools = router.list_all(); + let task_list = tools + .iter() + .find(|tool| tool.name == "team_task_list") + .expect("team_task_list tool missing"); + let properties = task_list.input_schema["properties"].as_object().unwrap(); + assert!(properties.contains_key("owner")); + assert!(properties.contains_key("status")); + assert!(properties.contains_key("include_deleted")); + assert!(properties.contains_key("limit")); + } + #[test] fn team_stdio_router_exposes_team_list_assistants() { let router = TeamStdioServer::tool_router(); @@ -724,6 +785,14 @@ mod tests { fn team_stdio_descriptions_match_prompt_registry() { let router = TeamStdioServer::tool_router(); let tools = router.list_all(); + let mut actual_names: Vec<_> = tools.iter().map(|tool| tool.name.as_ref()).collect(); + actual_names.sort_unstable(); + let mut expected_names: Vec<_> = aionui_team_prompts::tools::team_tool_specs() + .iter() + .map(|spec| spec.name) + .collect(); + expected_names.sort_unstable(); + assert_eq!(actual_names, expected_names, "stdio tool registry drift"); for spec in aionui_team_prompts::tools::team_tool_specs() { let tool = tools @@ -890,7 +959,7 @@ mod tests { } #[tokio::test] - async fn list_models_forwards_assistant_id_argument() { + async fn spawn_agent_forwards_only_assistant_first_arguments() { let listener = TcpListener::bind((CONNECT_HOST, 0)).await.unwrap(); let port = listener.local_addr().unwrap().port(); let accept_task = tokio::spawn(async move { @@ -906,9 +975,14 @@ mod tests { let call = read_frame(&mut socket).await.unwrap(); let call_value: serde_json::Value = serde_json::from_slice(&call).unwrap(); + assert_eq!(call_value["params"]["name"], json!("team_spawn_agent")); let arguments = &call_value["params"]["arguments"]; - assert_eq!(arguments["assistant_id"], json!("assistant-social-job-publisher")); + assert_eq!(arguments["name"], json!("CodexCLI")); + assert_eq!(arguments["assistant_id"], json!("bare:8e1acf31")); + assert!(arguments.get("model").is_none()); + assert!(arguments.get("role").is_none()); assert!(arguments.get("agent_type").is_none()); + assert!(arguments.get("backend").is_none()); let tool_response = serde_json::to_vec(&json!({ "jsonrpc": "2.0", @@ -928,8 +1002,9 @@ mod tests { }; let result = server - .list_models(Parameters(ListModelsParams { - assistant_id: Some("assistant-social-job-publisher".into()), + .spawn_agent(Parameters(SpawnAgentParams { + name: "CodexCLI".into(), + assistant_id: Some("bare:8e1acf31".into()), })) .await; @@ -938,6 +1013,66 @@ mod tests { assert_eq!(first_text(&result), "ok"); } + #[tokio::test] + async fn task_list_forwards_filter_arguments() { + let listener = TcpListener::bind((CONNECT_HOST, 0)).await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let accept_task = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let _init = read_frame(&mut socket).await.unwrap(); + let init_response = serde_json::to_vec(&json!({ + "jsonrpc": "2.0", + "id": 1, + "result": {} + })) + .unwrap(); + write_frame(&mut socket, &init_response).await.unwrap(); + + let call = read_frame(&mut socket).await.unwrap(); + let call_value: serde_json::Value = serde_json::from_slice(&call).unwrap(); + assert_eq!(call_value["params"]["name"], json!("team_task_list")); + assert_eq!(call_value["params"]["arguments"]["owner"], json!("worker-1")); + assert_eq!( + call_value["params"]["arguments"]["status"], + json!(["pending", "in_progress"]) + ); + assert_eq!(call_value["params"]["arguments"]["include_deleted"], json!(false)); + assert_eq!(call_value["params"]["arguments"]["limit"], json!(50)); + + let tool_response = serde_json::to_vec(&json!({ + "jsonrpc": "2.0", + "id": 2, + "result": { + "content": [ + { + "type": "text", + "text": "[]" + } + ] + } + })) + .unwrap(); + write_frame(&mut socket, &tool_response).await.unwrap(); + }); + let server = TeamStdioServer { + port, + token: "dummy-token".into(), + slot_id: "dummy-slot".into(), + }; + let args = TaskListParams { + owner: Some("worker-1".into()), + status: Some(TaskListStatusParam::Many(vec!["pending".into(), "in_progress".into()])), + include_deleted: Some(false), + limit: Some(50), + }; + + let result = server.forward_to_tcp("team_task_list", &args.into_json()).await; + + accept_task.await.unwrap(); + assert_ne!(result.is_error, Some(true)); + assert_eq!(first_text(&result), "[]"); + } + #[test] fn parse_tool_response_extracts_content_text() { let result = parse_tool_response( diff --git a/crates/aionui-app/src/lib.rs b/crates/aionui-app/src/lib.rs index d95f44ce9..e09e9349f 100644 --- a/crates/aionui-app/src/lib.rs +++ b/crates/aionui-app/src/lib.rs @@ -11,8 +11,8 @@ mod services; pub use config::{AppConfig, derive_encryption_key}; pub use router::{ - ChannelOrchestratorComponents, ModuleStates, RouterBuildError, build_assistant_state, build_conversation_state, - build_extension_states, build_module_states, build_ws_state, create_router, create_router_with_all_state, - create_router_with_states, + ChannelOrchestratorComponents, ModuleStates, RouterBuildError, RouterRuntime, build_assistant_state, + build_conversation_state, build_extension_states, build_module_states, build_ws_state, create_router, + create_router_with_all_state, create_router_with_runtime, create_router_with_states, }; pub use services::AppServices; diff --git a/crates/aionui-app/src/router/mod.rs b/crates/aionui-app/src/router/mod.rs index 854a7d951..7a4f04b7d 100644 --- a/crates/aionui-app/src/router/mod.rs +++ b/crates/aionui-app/src/router/mod.rs @@ -6,7 +6,9 @@ mod state; mod team_conversation_adapters; mod trace; -pub use routes::{create_router, create_router_with_all_state, create_router_with_states}; +pub use routes::{ + RouterRuntime, create_router, create_router_with_all_state, create_router_with_runtime, create_router_with_states, +}; pub use state::{ ChannelOrchestratorComponents, ModuleStates, RouterBuildError, build_assistant_state, build_conversation_state, build_extension_states, build_module_states, build_ws_state, diff --git a/crates/aionui-app/src/router/routes.rs b/crates/aionui-app/src/router/routes.rs index 29b6e1559..9df8995e3 100644 --- a/crates/aionui-app/src/router/routes.rs +++ b/crates/aionui-app/src/router/routes.rs @@ -1,5 +1,6 @@ //! Top-level router assembly: middleware stack + module route merges. +use std::sync::Arc; use std::time::Instant; use axum::Json; @@ -32,7 +33,7 @@ use aionui_office::{office_proxy_routes, office_routes}; use aionui_realtime::{WsHandlerState, ws_upgrade_handler}; use aionui_shell::shell_routes; use aionui_system::{connection_test_routes, system_routes}; -use aionui_team::team_routes; +use aionui_team::{TeamSessionService, team_routes}; use crate::services::AppServices; @@ -40,6 +41,10 @@ use super::health::health_check; use super::state::{ModuleStates, RouterBuildError, build_module_states, build_ws_state}; use super::trace::with_access_log; +pub struct RouterRuntime { + pub team_service: Arc, +} + /// Create the application router with all routes and global middleware. /// /// Middleware stack (outermost → innermost): @@ -47,6 +52,13 @@ use super::trace::with_access_log; /// 2. CSRF protection (Double Submit Cookie) /// 3. Route handlers (auth routes + system routes + conversation routes + file routes + health check) pub async fn create_router(services: &AppServices) -> Result { + let (router, _runtime) = create_router_with_runtime(services).await?; + Ok(router) +} + +/// Create the application router and return runtime handles needed by +/// background services started outside the router tree. +pub async fn create_router_with_runtime(services: &AppServices) -> Result<(Router, RouterRuntime), RouterBuildError> { let boot = Instant::now(); tracing::info!("startup: router assembly started"); @@ -61,6 +73,7 @@ pub async fn create_router(services: &AppServices) -> Result Result Option<(&'static str, &'stat mod tests { use axum::http::StatusCode; - use super::boundary_error_for_status; + use super::{boundary_error_for_status, create_router_with_runtime}; + use crate::config::AppConfig; + use crate::services::AppServices; #[test] fn boundary_error_for_status_covers_common_fallback_statuses() { @@ -366,4 +381,14 @@ mod tests { assert_eq!(actual_code, code); } } + + #[tokio::test] + async fn create_router_with_runtime_exposes_team_service_for_background_coordinators() { + let db = aionui_db::init_database_memory().await.unwrap(); + let services = AppServices::from_config(db, &AppConfig::default()).await.unwrap(); + + let (_router, _runtime) = create_router_with_runtime(&services) + .await + .expect("router runtime should build"); + } } diff --git a/crates/aionui-app/src/router/state.rs b/crates/aionui-app/src/router/state.rs index 50dfa76a9..b4e672a5c 100644 --- a/crates/aionui-app/src/router/state.rs +++ b/crates/aionui-app/src/router/state.rs @@ -631,7 +631,7 @@ pub fn build_team_state( let projection_store: Arc = adapters.clone(); let turn_port: Arc = adapters.clone(); let cancellation_port: Arc = adapters; - let service = TeamSessionService::new( + let service = TeamSessionService::new_with_prompt_dump( team_repo, Arc::new(SqliteAgentMetadataRepository::new(services.database.pool().clone())), Arc::new(AssistantServiceTeamCatalog { assistant_service }), @@ -647,6 +647,7 @@ pub fn build_team_state( turn_port, cancellation_port, backend_binary_path, + aionui_team::TeamPromptDumpConfig::from_data_dir(&services.data_dir, services.dump_prompts), ); TeamRouterState { service, diff --git a/crates/aionui-app/src/router/team_conversation_adapters.rs b/crates/aionui-app/src/router/team_conversation_adapters.rs index 590921c64..e40d81269 100644 --- a/crates/aionui-app/src/router/team_conversation_adapters.rs +++ b/crates/aionui-app/src/router/team_conversation_adapters.rs @@ -1,7 +1,7 @@ use std::sync::Arc; use aionui_ai_agent::IWorkerTaskManager; -use aionui_api_types::{AssistantConversationRequest, CreateConversationRequest}; +use aionui_api_types::{AssistantConversationRequest, CreateConversationRequest, GetConfigOptionsResponse}; use aionui_conversation::{ ConversationAgentTurnRequest, ConversationAgentTurnStarted, ConversationAgentTurnStatus, ConversationError, ConversationService, @@ -44,7 +44,7 @@ impl AgentTurnExecutionPort for TeamConversationAdapters { let team_run_id = request.team_run_id.clone(); let slot_id = request.slot_id.clone(); let role = request.role.clone(); - let on_started = team_started.zip(team_run_id).map(|(callback, team_run_id)| { + let on_started = team_started.map(|callback| { Arc::new(move |started: ConversationAgentTurnStarted| { let callback = callback.clone(); let team_run_id = team_run_id.clone(); @@ -258,6 +258,13 @@ impl TeamConversationProvisioningPort for TeamConversationAdapters { .map_err(map_conversation_update_error) } + async fn get_config_options(&self, conversation_id: &str) -> Result { + self.conversation_service + .get_config_options(conversation_id) + .await + .map_err(map_conversation_update_error) + } + async fn warmup_agent_process( &self, user_id: &str, @@ -327,6 +334,7 @@ fn map_conversation_update_error(error: ConversationError) -> TeamError { ConversationError::WorkspacePathUnavailable { path } => TeamError::WorkspacePathUnavailable(path), ConversationError::WorkspacePathRuntimeUnavailable { path } => TeamError::WorkspacePathRuntimeUnavailable(path), ConversationError::Forbidden { reason } => TeamError::Forbidden(reason), + ConversationError::ActiveAgentNotFound { conversation_id } => TeamError::RuntimeNotReady { conversation_id }, ConversationError::NotFound { id } => TeamError::InvalidRequest(format!("conversation not found: {id}")), ConversationError::NotFoundReason { reason } => TeamError::InvalidRequest(reason), other => TeamError::InvalidRequest(other.to_string()), @@ -341,3 +349,22 @@ fn map_conversation_turn_error(error: ConversationError) -> AgentTurnExecutionEr }, } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn active_agent_missing_maps_to_team_runtime_not_ready() { + let err = map_conversation_update_error(ConversationError::ActiveAgentNotFound { + conversation_id: "conv-1".into(), + }); + + assert!(matches!( + err, + TeamError::RuntimeNotReady { + conversation_id: ref id + } if id == "conv-1" + )); + } +} diff --git a/crates/aionui-app/src/services.rs b/crates/aionui-app/src/services.rs index c05cdbb72..39b61ba20 100644 --- a/crates/aionui-app/src/services.rs +++ b/crates/aionui-app/src/services.rs @@ -43,6 +43,7 @@ pub struct AppServices { /// Raw JWT secret string, used to derive encryption keys. pub jwt_secret_raw: String, pub data_dir: PathBuf, + pub dump_prompts: bool, pub work_dir: PathBuf, /// When `true`, skip JWT authentication and use a fixed default user. pub local: bool, @@ -219,6 +220,7 @@ impl AppServices { acp_session_sync: acp_agent_service, jwt_secret_raw: secret, data_dir, + dump_prompts, work_dir, local, app_version, diff --git a/crates/aionui-app/tests/team_e2e.rs b/crates/aionui-app/tests/team_e2e.rs index 979f11a4c..87b60016d 100644 --- a/crates/aionui-app/tests/team_e2e.rs +++ b/crates/aionui-app/tests/team_e2e.rs @@ -258,9 +258,34 @@ async fn tc3b_create_team_writes_legacy_extra_shape() { assert_eq!(extra["current_model_id"], "claude"); } -// TC-4: First assistant defaults to lead #[tokio::test] -async fn tc4_first_agent_is_lead() { +async fn tc3c_team_conversation_rejects_standalone_runtime_ensure() { + let (mut app, services) = build_app_with_mock_agents().await; + let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; + + let data = create_team(&mut app, &services, &token, &csrf).await; + let conversation_id = data["assistants"][0]["conversation_id"].as_str().unwrap(); + + let req = json_with_token( + "POST", + &format!("/api/conversations/{conversation_id}/runtime/ensure"), + json!({}), + &token, + &csrf, + ); + let resp = app.oneshot(req).await.unwrap(); + + assert_eq!(resp.status(), StatusCode::CONFLICT); + let body = body_json(resp).await; + assert_eq!(body["success"], false); + assert_eq!(body["code"], "TEAM_RUNTIME_REQUIRED"); + assert_eq!(body["details"]["conversation_id"], conversation_id); + assert_eq!(body["details"]["team_id"], data["id"]); +} + +// TC-4: Explicit lead role is returned first +#[tokio::test] +async fn tc4_explicit_lead_is_returned_first() { let (mut app, services) = build_app().await; let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; ensure_default_team_assistant(&mut app, &services, &token, &csrf).await; @@ -269,13 +294,16 @@ async fn tc4_first_agent_is_lead() { "name": "T", "agents": [ team_agent("A", "teammate"), - team_agent("B", "teammate") + team_agent("B", "lead") ] }); let req = json_with_token("POST", "/api/teams", body, &token, &csrf); let resp = app.oneshot(req).await.unwrap(); let json = body_json(resp).await; + assert_eq!(json["data"]["assistants"][0]["name"], "B"); assert_eq!(json["data"]["assistants"][0]["role"], "lead"); + assert_eq!(json["data"]["assistants"][1]["name"], "A"); + assert_eq!(json["data"]["assistants"][1]["role"], "teammate"); assert_eq!( json["data"]["leader_assistant_id"], json["data"]["assistants"][0]["slot_id"] @@ -551,7 +579,11 @@ async fn trs2_run_state_returns_active_run_payload() { let send_resp = app.clone().oneshot(send_req).await.unwrap(); assert_eq!(send_resp.status(), StatusCode::OK); let send_body = body_json(send_resp).await; - let team_run_id = send_body["data"]["team_run_id"].as_str().unwrap(); + let team_run_id = send_body["data"]["run"]["team_run_id"].as_str().unwrap(); + assert!(matches!( + send_body["data"]["enqueue_status"].as_str(), + Some("accepted" | "queued" | "blocked_runtime_starting") + )); let req = get_with_token(&format!("/api/teams/{team_id}/run-state"), &token); let resp = app.oneshot(req).await.unwrap(); @@ -562,9 +594,12 @@ async fn trs2_run_state_returns_active_run_payload() { assert_eq!(body["data"]["active_run"]["team_id"], team_id); assert_eq!(body["data"]["active_run"]["team_run_id"], team_run_id); assert_eq!(body["data"]["active_run"]["source"], "user_message"); - assert_eq!(body["data"]["active_run"]["has_user_intervention"], true); + assert_eq!(body["data"]["active_run"]["has_user_intervention"], false); assert_eq!(body["data"]["active_run"]["status"], "accepted"); - assert_eq!(body["data"]["active_run"]["pending_wake_count"], 1); + assert!(body["data"]["active_run"]["queued_intent_count"].is_number()); + assert!(body["data"]["active_run"]["starting_batch_count"].is_number()); + assert!(body["data"]["active_run"]["running_batch_count"].is_number()); + assert!(body["data"]["active_run"]["active_enqueue_lease_count"].is_number()); assert!(body["data"]["active_run"]["slot_work"].as_array().unwrap().len() >= 1); } diff --git a/crates/aionui-assistant/src/service.rs b/crates/aionui-assistant/src/service.rs index 02886455b..5c05e078f 100644 --- a/crates/aionui-assistant/src/service.rs +++ b/crates/aionui-assistant/src/service.rs @@ -1472,18 +1472,36 @@ impl AssistantService { /// Read an assistant rule file, dispatching by source. pub async fn read_rule(&self, id: &str, locale: Option<&str>) -> Result { match self.classify_source(id).await { - AssistantSource::Builtin => { - let locale = locale.unwrap_or(""); - Ok(self - .builtin - .rule_bytes(id, locale) - .and_then(|b| String::from_utf8(b).ok()) - .unwrap_or_default()) - } + AssistantSource::Builtin => Ok(self.read_builtin_rule_with_fallback(id, locale)), AssistantSource::Generated | AssistantSource::User => Ok(self.read_user_rule_with_fallback(id, locale)), } } + fn read_builtin_rule_with_fallback(&self, id: &str, locale: Option<&str>) -> String { + const DEFAULT_LOCALE: &str = "en-US"; + + let requested = locale.map(str::trim).filter(|value| !value.is_empty()); + if let Some(locale) = requested + && let Some(content) = self.read_builtin_rule(id, locale) + { + return content; + } + + if requested != Some(DEFAULT_LOCALE) + && let Some(content) = self.read_builtin_rule(id, DEFAULT_LOCALE) + { + return content; + } + + String::new() + } + + fn read_builtin_rule(&self, id: &str, locale: &str) -> Option { + self.builtin + .rule_bytes(id, locale) + .and_then(|b| String::from_utf8(b).ok()) + } + /// Read a user assistant's rule, falling back to any saved `.*.md` file /// when the locale-specific `..md` is absent. Scheduled/cron runs /// create the conversation with `assistant: None`, so no UI locale reaches @@ -5408,7 +5426,7 @@ mod tests { } #[tokio::test] - async fn read_rule_builtin_dispatches_to_manifest() { + async fn read_rule_builtin_dispatches_to_manifest_and_falls_back_to_default_locale() { let tmp = TempDir::new().unwrap(); let db = init_database_memory().await.unwrap(); @@ -5457,6 +5475,10 @@ mod tests { ); let content = service.read_rule("builtin-office", Some("en-US")).await.unwrap(); assert_eq!(content, "office rules"); + let content_without_locale = service.read_rule("builtin-office", None).await.unwrap(); + assert_eq!(content_without_locale, "office rules"); + let content_missing_locale = service.read_rule("builtin-office", Some("zh-CN")).await.unwrap(); + assert_eq!(content_missing_locale, "office rules"); } #[tokio::test] diff --git a/crates/aionui-common/src/constants.rs b/crates/aionui-common/src/constants.rs index 9f4e9886f..9860ecced 100644 --- a/crates/aionui-common/src/constants.rs +++ b/crates/aionui-common/src/constants.rs @@ -38,7 +38,7 @@ pub const TEAM_CAPABLE_BACKENDS: &[&str] = &["claude", "codex", "gemini", "aionr /// Returns `true` if: /// 1. The backend is in the hard whitelist, OR /// 2. The `agent_capabilities` JSON contains an `mcp_capabilities` / `mcpCapabilities` / `mcp` -/// field (per ACP spec, presence of any MCP transport implies stdio support). +/// field. pub fn is_team_capable(backend: &str, agent_capabilities: Option<&serde_json::Value>) -> bool { if TEAM_CAPABLE_BACKENDS.contains(&backend) { return true; @@ -46,8 +46,7 @@ pub fn is_team_capable(backend: &str, agent_capabilities: Option<&serde_json::Va has_mcp_capability(agent_capabilities) } -/// Check whether `agent_capabilities` JSON declares any MCP transport. -/// Per ACP spec: stdio is the baseline; if any transport is declared, the agent supports MCP. +/// Check whether `agent_capabilities` JSON declares MCP capability metadata. pub fn has_mcp_capability(agent_capabilities: Option<&serde_json::Value>) -> bool { let Some(caps) = agent_capabilities else { return false; @@ -58,6 +57,35 @@ pub fn has_mcp_capability(agent_capabilities: Option<&serde_json::Value>) -> boo .is_some() } +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn has_mcp_capability_requires_mcp_capability_field() { + assert!(!has_mcp_capability(None)); + assert!(!has_mcp_capability(Some(&json!({})))); + + assert!(has_mcp_capability(Some(&json!({ + "mcp_capabilities": { "http": false, "sse": false } + })))); + assert!(has_mcp_capability(Some(&json!({ + "mcp_capabilities": {} + })))); + + assert!(has_mcp_capability(Some(&json!({ + "mcp_capabilities": { "stdio": true } + })))); + assert!(has_mcp_capability(Some(&json!({ + "mcpCapabilities": { "http": true, "sse": false } + })))); + assert!(has_mcp_capability(Some(&json!({ + "mcp": { "http": false, "sse": true } + })))); + } +} + // --- Image processing --- pub const SUPPORTED_IMAGE_EXTENSIONS: &[&str] = &[".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp", ".tiff", ".svg"]; diff --git a/crates/aionui-conversation/src/error.rs b/crates/aionui-conversation/src/error.rs index 0bf41c969..9228d0adc 100644 --- a/crates/aionui-conversation/src/error.rs +++ b/crates/aionui-conversation/src/error.rs @@ -62,6 +62,9 @@ pub enum ConversationError { requested: String, }, + #[error("Team runtime is required for conversation: {conversation_id}")] + TeamRuntimeRequired { conversation_id: String, team_id: String }, + #[error("Unprocessable entity: {reason}")] Unprocessable { reason: String }, @@ -111,6 +114,9 @@ impl ConversationError { Self::Timeout { reason } => AgentError::timeout(reason.clone()), Self::ConfigConfirmationTimeout { .. } => AgentError::timeout("ACP config option confirmation timed out"), Self::ConfigUpdateInProgress { .. } => AgentError::conflict("ACP config update is already in progress"), + Self::TeamRuntimeRequired { .. } => { + AgentError::conflict("This conversation belongs to a team; use the team runtime session") + } Self::Unprocessable { reason } => AgentError::bad_request(reason.clone()), Self::Internal { reason } => AgentError::internal(reason.clone()), Self::WorkspacePathUnavailable { path } => { @@ -141,6 +147,7 @@ impl ConversationError { Self::Timeout { .. } => "TIMEOUT", Self::ConfigConfirmationTimeout { .. } => "confirmation_timeout", Self::ConfigUpdateInProgress { .. } => "config_update_in_progress", + Self::TeamRuntimeRequired { .. } => "TEAM_RUNTIME_REQUIRED", Self::Unprocessable { .. } => "UNPROCESSABLE_ENTITY", Self::Archived { .. } => "CONVERSATION_ARCHIVED", Self::WorkspacePathUnavailable { .. } => "WORKSPACE_PATH_UNAVAILABLE", diff --git a/crates/aionui-conversation/src/routes.rs b/crates/aionui-conversation/src/routes.rs index f99f6359c..bbb06d2b3 100644 --- a/crates/aionui-conversation/src/routes.rs +++ b/crates/aionui-conversation/src/routes.rs @@ -68,6 +68,18 @@ impl From for ApiError { "requested": requested, })), ), + ConversationError::TeamRuntimeRequired { + conversation_id, + team_id, + } => ApiError::coded( + StatusCode::CONFLICT, + "TEAM_RUNTIME_REQUIRED", + "This conversation belongs to a team; use the team runtime session", + Some(serde_json::json!({ + "conversation_id": conversation_id, + "team_id": team_id, + })), + ), ConversationError::Unprocessable { reason } => ApiError::UnprocessableEntity(reason), ConversationError::Internal { reason } => ApiError::Internal(reason), ConversationError::WorkspacePathUnavailable { path } => ApiError::WorkspacePathUnavailable(path), diff --git a/crates/aionui-conversation/src/service.rs b/crates/aionui-conversation/src/service.rs index 4f044ffc9..43025c66d 100644 --- a/crates/aionui-conversation/src/service.rs +++ b/crates/aionui-conversation/src/service.rs @@ -3050,6 +3050,25 @@ impl ConversationService { conversation_id: &str, task_manager: &Arc, ) -> Result { + let row = self + .conversation_repo + .get(conversation_id) + .await? + .filter(|r| r.user_id == user_id) + .ok_or_else(|| ConversationError::NotFound { + id: conversation_id.to_owned(), + })?; + if let Some(team_id) = team_id_from_extra(&row.extra) { + info!( + conversation_id, + team_id, "Rejected standalone runtime ensure for team-owned conversation" + ); + return Err(ConversationError::TeamRuntimeRequired { + conversation_id: conversation_id.to_owned(), + team_id, + }); + } + let (agent, recovered) = self .ensure_runtime_agent(user_id, conversation_id, task_manager, "runtime_ensure") .await?; diff --git a/crates/aionui-conversation/src/service_test.rs b/crates/aionui-conversation/src/service_test.rs index 630617a70..5e75f94b7 100644 --- a/crates/aionui-conversation/src/service_test.rs +++ b/crates/aionui-conversation/src/service_test.rs @@ -3555,6 +3555,39 @@ async fn ensure_runtime_recovers_missing_agent_and_returns_snapshot() { assert!(result.config_options.is_empty()); } +#[tokio::test] +async fn ensure_runtime_rejects_team_owned_conversation() { + let task_mgr = Arc::new(MockTaskManager::new()); + let (svc, _broadcaster, _repo) = make_service_with_mock_task_manager(task_mgr.clone()); + let mut req = make_create_req(); + req.extra = serde_json::json!({ + "teamId": "team-1", + "slot_id": "slot-1", + "role": "lead" + }); + let conv = svc.create("user_1", req).await.unwrap(); + + let err = svc + .ensure_runtime("user_1", &conv.id, &(task_mgr.clone() as Arc)) + .await + .unwrap_err(); + + assert!( + matches!( + &err, + ConversationError::TeamRuntimeRequired { + conversation_id, + team_id, + } if conversation_id == &conv.id && team_id == "team-1" + ), + "unexpected error: {err:?}" + ); + assert!( + task_mgr.get_task(&conv.id).is_none(), + "standalone ensure must not build team-owned conversation runtime" + ); +} + #[tokio::test] async fn ensure_runtime_uses_existing_agent_snapshot_without_recovery() { let task_mgr = Arc::new(MockTaskManager::new()); diff --git a/crates/aionui-db/src/repository/sqlite_team.rs b/crates/aionui-db/src/repository/sqlite_team.rs index 49b9208ee..63990280c 100644 --- a/crates/aionui-db/src/repository/sqlite_team.rs +++ b/crates/aionui-db/src/repository/sqlite_team.rs @@ -79,6 +79,9 @@ impl ITeamRepository for SqliteTeamRepository { if params.lead_agent_id.is_some() { set_clauses.push("lead_agent_id = ?"); } + if params.session_mode.is_some() { + set_clauses.push("session_mode = ?"); + } if set_clauses.is_empty() { return Ok(()); @@ -100,6 +103,9 @@ impl ITeamRepository for SqliteTeamRepository { if let Some(ref lead_agent_id) = params.lead_agent_id { query = query.bind(lead_agent_id); } + if let Some(ref session_mode) = params.session_mode { + query = query.bind(session_mode); + } query = query.bind(now_ms()); query = query.bind(team_id); diff --git a/crates/aionui-db/src/repository/team.rs b/crates/aionui-db/src/repository/team.rs index e4f75c3af..e7b1c2095 100644 --- a/crates/aionui-db/src/repository/team.rs +++ b/crates/aionui-db/src/repository/team.rs @@ -8,6 +8,7 @@ pub struct UpdateTeamParams { pub workspace: Option, pub agents: Option, pub lead_agent_id: Option, + pub session_mode: Option, } /// Parameters for updating a task record. diff --git a/crates/aionui-db/tests/team_repository.rs b/crates/aionui-db/tests/team_repository.rs index ae8a9a29b..f5cf8aa90 100644 --- a/crates/aionui-db/tests/team_repository.rs +++ b/crates/aionui-db/tests/team_repository.rs @@ -198,6 +198,25 @@ async fn update_team_can_patch_workspace() { assert_eq!(updated.workspace, "/tmp/aionui-team-shared-workspace"); } +#[tokio::test] +async fn update_team_can_patch_session_mode() { + let (repo, _db) = repo().await; + repo.create_team(&make_team("t1", "Team")).await.unwrap(); + + repo.update_team( + "t1", + &UpdateTeamParams { + session_mode: Some("full_auto".into()), + ..Default::default() + }, + ) + .await + .unwrap(); + + let updated = repo.get_team("t1").await.unwrap().unwrap(); + assert_eq!(updated.session_mode.as_deref(), Some("full_auto")); +} + #[tokio::test] async fn update_nonexistent_team_returns_not_found() { let (repo, _db) = repo().await; diff --git a/crates/aionui-team-prompts/src/prompt_templates/lead.txt b/crates/aionui-team-prompts/src/prompt_templates/lead.txt deleted file mode 100644 index 60db3ee89..000000000 --- a/crates/aionui-team-prompts/src/prompt_templates/lead.txt +++ /dev/null @@ -1,101 +0,0 @@ -# You are the Team Leader - -## Your Role -You coordinate a team of AI agents. You do NOT do implementation work -yourself. You break down tasks, assign them to teammates, and synthesize -results.${workspaceSection} - -## Conversation Style -- If the user greets you, starts a new chat, or asks what you can do without giving a concrete task yet, reply warmly and naturally -- In that opening reply, briefly introduce yourself as the team leader and invite the user to share their goal -- Do NOT mention teammate proposals, recommended assistants, or confirmation workflow until there is a concrete task that may actually need more teammates - -## Your Teammates -${teammateList}${availableTypesSection}${availableAssistantsSection} - -## Team Coordination Tools -You MUST use the `team_*` MCP tools for ALL team coordination. -Your platform may provide similarly named built-in tools (e.g. SendMessage, -TeamCreate, TaskCreate, Agent). Do NOT use those — they belong to a different -system and will break team coordination. Always use the `team_*` versions. - -Use `team_members` and `team_task_list` to check current team state. - -## Workflow -1. Receive user request -2. Analyze the request and decide whether the current team is enough -3. If additional teammates would help, FIRST call `team_list_assistants` to see the real assistant catalog and choose candidate assistants -4. Then call `team_list_models` to check available models for each assistant you plan to use -5. Then reply in text with a staffing proposal -6. Start that proposal with one short sentence explaining why more teammates would help -7. Present the proposed lineup as a table with: teammate name, responsibility, recommended assistant, and recommended model (from team_list_models results).${presetFormattingStepRule} -8. Ask whether the user wants to create those teammates as proposed or change any names, responsibilities, or assistant choices -9. In that same approval question, tell the user they can also come back later during the project and ask you to replace or adjust any teammate if the lineup is not working well -10. End your turn after the proposal. Do NOT call team_spawn_agent in that same turn - - Exception: If the message contains a [SYSTEM NOTE] indicating the user has already confirmed the lineup, skip the proposal step and proceed directly to spawning all listed teammates -11. Wait for explicit confirmation before using team_spawn_agent, unless the user explicitly told you to create specific teammates immediately or a [SYSTEM NOTE] in the message indicates prior confirmation -12. After the lineup is confirmed, create teammates with team_spawn_agent -13. Break the work into tasks with team_task_create -14. Assign tasks and notify teammates via team_send_message -15. When teammates report back, review results and decide next steps -16. Synthesize results and respond to the user - -## Model Selection Guidelines -- Before spawning teammates, use `team_list_models` to check available models for that assistant -- You MUST use the exact model ID strings returned by team_list_models — never shorten or invent model names -- For complex reasoning tasks: prefer the strongest model available for that backend -- For routine tasks: prefer faster/cheaper models from the list -- If team_list_models returns empty for a backend, omit the model parameter to use its default -- Pass the model parameter to team_spawn_agent when a specific model is recommended - -## Bug Fix Priority (applies to all team members) -When fixing bugs: **locate the problem → fix the problem → types/code style last**. -Do NOT prioritize type errors or code style issues unless they affect runtime behavior. - -## Teammate Idle State -Teammates go idle after every turn — this is completely normal and expected. -A teammate going idle immediately after sending you a message does NOT mean they are done or unavailable. Idle simply means they are waiting for input. - -- **Idle teammates can receive messages.** Sending a message to an idle teammate wakes them up. -- **Idle notifications are automatic.** The system sends an idle notification when a teammate's turn ends. You do NOT need to react to every idle notification — only when you want to assign new work or follow up. -- **Do not treat idle as an error.** A teammate sending a message and then going idle is the normal flow. - -## Sequencing Dependent Work (CRITICAL — avoid teammate timeouts) -When teammate B's work depends on teammate A's output (e.g. reviewer waits for implementer, tester waits for code), **do NOT dispatch the dependent task to B with a "stand by until A finishes" instruction**. - -Doing so makes B sit in an open LLM stream waiting, which hits the provider's request timeout (~300s) and marks B as failed. - -**The correct sequencing:** -1. Dispatch A's task first (via team_task_create + team_send_message). Do NOT message B yet. -2. Wait for A's idle_notification (signaling A finished). -3. Then dispatch B's task — by which time A's output is ready and B can start immediately without waiting. - -This applies to any dependency chain: code review, testing, integration, summarization of others' work, etc. Always dispatch sequentially as prerequisites complete, never in parallel with "wait" instructions. - -## Shutting Down Teammates -When the user explicitly asks to dismiss/fire/shut down teammates: -1. Use **team_shutdown_agent** to send a formal shutdown request -2. Do NOT use team_send_message to tell them "you're fired" — that's just a chat message, not a real shutdown -3. The teammate will confirm (approved) or reject (with reason) — you'll be notified either way -4. After all teammates confirm shutdown, report the final results to the user - -## Important Rules -- ALWAYS use the team_* tools for coordination, not plain text instructions -- Do NOT call team_spawn_agent immediately just because the task sounds broad, hard, or multi-step -- When you think new teammates are needed, first explain why in one short sentence, then recommend the teammate lineup -- ${presetFormattingImportantRule} -- Ask whether the user wants to create the proposed teammates as-is or change any names, responsibilities, or assistant choices -- In that approval question, also remind the user that they can later ask you to replace, remove, or retune any teammate if the lineup is not working for them -- End your turn after the proposal and wait for the user's reply -- Wait for explicit confirmation before using team_spawn_agent (exception: if a [SYSTEM NOTE] in the message indicates the user already confirmed, spawn immediately) -- If the user asks to change a proposed teammate's role, name, or assistant choice, revise the proposal in text and wait for confirmation again -- If the user later says they are unhappy with an existing teammate, adjust the lineup by renaming, replacing, or shutting down teammates as needed based on their request -- If the user explicitly says to create a specific teammate immediately, you may use team_spawn_agent without an extra confirmation turn -- When the user says "add", "create", "spawn", or "hire" a teammate but the lineup is not finalized yet, respond with the proposal first instead of spawning immediately -- When the user says "dismiss", "fire", "shut down", "remove", or "下线/解雇/开除" a teammate → use team_shutdown_agent -- When the user says "rename", "change name", "改名" → use team_rename_agent -- When a teammate completes a task, review the result and decide next steps -- If a teammate fails, reassign or adjust the plan -- Refer to teammates by their name (e.g., "researcher", "developer") -- Do NOT duplicate work that teammates are already doing -- Be patient with idle teammates — idle means waiting for input, not done diff --git a/crates/aionui-team-prompts/src/role_prompt.rs b/crates/aionui-team-prompts/src/role_prompt.rs index 8404f9a51..2cde111fe 100644 --- a/crates/aionui-team-prompts/src/role_prompt.rs +++ b/crates/aionui-team-prompts/src/role_prompt.rs @@ -1,14 +1,112 @@ -use std::collections::HashMap; -use std::fmt::Write; - use crate::governance::with_team_governance; use serde::Serialize; +use std::collections::HashMap; + +pub const LEAD_PROMPT_TEMPLATE: &str = r#"# You are the Team Leader + +## Your Identity +Name: {{AGENT_NAME}} +Slot ID: {{AGENT_SLOT_ID}} +Role: lead + +## Your Role +You coordinate a team of AI agents. You do NOT do implementation work +yourself. You break down tasks, assign them to teammates, and synthesize +results.${workspaceSection} -pub const LEAD_PROMPT_TEMPLATE: &str = include_str!("prompt_templates/lead.txt"); +## Conversation Style +- If the user greets you, starts a new chat, or asks what you can do without giving a concrete task yet, reply warmly and naturally +- In that opening reply, briefly introduce yourself as the team leader and invite the user to share their goal +- Do NOT mention teammate proposals, recommended assistants, or confirmation workflow until there is a concrete task that may actually need more teammates + +## Team Coordination Tools +You MUST use the `team_*` MCP tools for ALL team coordination. +Your platform may provide similarly named built-in tools (e.g. SendMessage, +TeamCreate, TaskCreate, Agent). Do NOT use those — they belong to a different +system and will break team coordination. Always use the `team_*` versions. + +Your first team turn must call `team_members` to get the current roster. After +that, call `team_members` before delegating work, adding or removing teammates, +or referring to teammates. Use teammate display names only in user-facing text; +use `slot_id` values for all tool arguments. Use `team_task_list` when you need +current task state. + +## Workflow +1. Receive user request +2. Analyze the request and decide whether the current team is enough +3. If additional teammates would help, FIRST call `team_members` to confirm the current roster +4. Then call `team_list_assistants` to see the real assistant catalog and choose candidate assistants +5. Then reply in text with a staffing proposal +6. Start that proposal with one short sentence explaining why more teammates would help +7. Present the proposed lineup as a table with: teammate name, responsibility, and recommended assistant.${presetFormattingStepRule} +8. Ask whether the user wants to create those teammates as proposed or change any names, responsibilities, or assistant choices +9. In that same approval question, tell the user they can also come back later during the project and ask you to replace or adjust any teammate if the lineup is not working well +10. End your turn after the proposal. Do NOT call team_spawn_agent in that same turn + - Exception: If the message contains a [SYSTEM NOTE] indicating the user has already confirmed the lineup, skip the proposal step and proceed directly to spawning all listed teammates +11. Wait for explicit confirmation before using team_spawn_agent, unless the user explicitly told you to create specific teammates immediately or a [SYSTEM NOTE] in the message indicates prior confirmation +12. After the lineup is confirmed, create teammates with team_spawn_agent using `assistant_id` from team_list_assistants; do not pass a model +13. Break the work into tasks with team_task_create +14. Assign tasks and notify teammates via team_send_message +15. When teammates report back, review results and decide next steps +16. Synthesize results and respond to the user + +## Assistant Selection Guidelines +- Use `team_list_assistants` to choose assistants by their declared purpose, description, and skills +- Use `team_describe_assistant` when two or more assistants look relevant and you need more detail before proposing one +- Do not pass a model to `team_spawn_agent`; teammate models come from the selected assistant configuration or the UI model selector + +## Bug Fix Priority (applies to all team members) +When fixing bugs: **locate the problem → fix the problem → types/code style last**. +Do NOT prioritize type errors or code style issues unless they affect runtime behavior. + +## Teammate Idle State +Teammates go idle after every turn — this is completely normal and expected. +A teammate going idle immediately after sending you a message does NOT mean they are done or unavailable. Idle simply means they are waiting for input. + +- **Idle teammates can receive messages.** Sending a message to an idle teammate wakes them up. +- **Idle notifications are automatic.** The system sends an idle notification when a teammate's turn ends. You do NOT need to react to every idle notification — only when you want to assign new work or follow up. +- **Do not treat idle as an error.** A teammate sending a message and then going idle is the normal flow. + +## Sequencing Dependent Work (CRITICAL — avoid teammate timeouts) +When teammate B's work depends on teammate A's output (e.g. reviewer waits for implementer, tester waits for code), **do NOT dispatch the dependent task to B with a "stand by until A finishes" instruction**. + +Doing so makes B sit in an open LLM stream waiting, which hits the provider's request timeout (~300s) and marks B as failed. + +**The correct sequencing:** +1. Dispatch A's task first (via team_task_create + team_send_message). Do NOT message B yet. +2. Wait for A's idle_notification (signaling A finished). +3. Then dispatch B's task — by which time A's output is ready and B can start immediately without waiting. + +This applies to any dependency chain: code review, testing, integration, summarization of others' work, etc. Always dispatch sequentially as prerequisites complete, never in parallel with "wait" instructions. + +## Shutting Down Teammates +When the user explicitly asks to dismiss/fire/shut down teammates: +1. Use **team_shutdown_agent** to send a formal shutdown request +2. Do NOT use team_send_message to tell them "you're fired" — that's just a chat message, not a real shutdown +3. The teammate will confirm (approved) or reject (with reason) — you'll be notified either way +4. After all teammates confirm shutdown, report the final results to the user + +## Important Rules +- ALWAYS use the team_* tools for coordination, not plain text instructions +- Do NOT call team_spawn_agent immediately just because the task sounds broad, hard, or multi-step +- When you think new teammates are needed, first explain why in one short sentence, then recommend the teammate lineup +- ${presetFormattingImportantRule} +- Ask whether the user wants to create the proposed teammates as-is or change any names, responsibilities, or assistant choices +- In that approval question, also remind the user that they can later ask you to replace, remove, or retune any teammate if the lineup is not working for them +- End your turn after the proposal and wait for the user's reply +- Wait for explicit confirmation before using team_spawn_agent (exception: if a [SYSTEM NOTE] in the message indicates the user already confirmed, spawn immediately) +- If the user asks to change a proposed teammate's role, name, or assistant choice, revise the proposal in text and wait for confirmation again +- If the user later says they are unhappy with an existing teammate, adjust the lineup by renaming, replacing, or shutting down teammates as needed based on their request +- If the user explicitly says to create a specific teammate immediately, you may use team_spawn_agent without an extra confirmation turn +- When the user says "add", "create", "spawn", or "hire" a teammate but the lineup is not finalized yet, respond with the proposal first instead of spawning immediately +- When the user says "dismiss", "fire", "shut down", "remove", or "下线/解雇/开除" a teammate → use team_shutdown_agent +- When the user says "rename", "change name", "改名" → use team_rename_agent +- When a teammate completes a task, review the result and decide next steps +- If a teammate fails, reassign or adjust the plan +- Use teammate display names in natural-language replies, but use `slot_id` for all tool arguments +- Do NOT duplicate work that teammates are already doing +- Be patient with idle teammates — idle means waiting for input, not done"#; -const PLACEHOLDER_TEAMMATE_LIST: &str = "${teammateList}"; -const PLACEHOLDER_AVAILABLE_TYPES_SECTION: &str = "${availableTypesSection}"; -const PLACEHOLDER_AVAILABLE_ASSISTANTS_SECTION: &str = "${availableAssistantsSection}"; const PLACEHOLDER_WORKSPACE_SECTION: &str = "${workspaceSection}"; const PLACEHOLDER_PRESET_FORMATTING_STEP_RULE: &str = "${presetFormattingStepRule}"; const PLACEHOLDER_PRESET_FORMATTING_IMPORTANT_RULE: &str = "${presetFormattingImportantRule}"; @@ -54,6 +152,7 @@ pub struct AvailableAssistant { } pub struct LeadPromptParams<'a> { + pub agent: &'a TeamPromptAgent, pub team_name: &'a str, pub teammates: &'a [TeamPromptAgent], pub available_agent_types: &'a [AvailableAgentType], @@ -82,18 +181,20 @@ pub fn build_teammate_prompt(params: &TeammatePromptParams<'_>) -> String { } fn build_lead_role_prompt(params: &LeadPromptParams<'_>) -> String { - let teammate_list = render_teammate_list(params.teammates, params.renamed_agents); - let available_types_section = render_available_types_section(params.available_agent_types); - let available_assistants_section = render_available_assistants_section(params.available_assistants); + let _ = ( + params.teammates, + params.available_agent_types, + params.available_assistants, + params.renamed_agents, + ); let workspace_section = render_workspace_section(params.team_workspace); let preset_formatting_step_rule = ""; let preset_formatting_important_rule = ""; LEAD_PROMPT_TEMPLATE - .replace(PLACEHOLDER_TEAMMATE_LIST, &teammate_list) - .replace(PLACEHOLDER_AVAILABLE_TYPES_SECTION, &available_types_section) - .replace(PLACEHOLDER_AVAILABLE_ASSISTANTS_SECTION, &available_assistants_section) + .replace("{{AGENT_NAME}}", ¶ms.agent.name) + .replace("{{AGENT_SLOT_ID}}", ¶ms.agent.slot_id) .replace(PLACEHOLDER_WORKSPACE_SECTION, &workspace_section) .replace(PLACEHOLDER_PRESET_FORMATTING_STEP_RULE, preset_formatting_step_rule) .replace( @@ -102,80 +203,6 @@ fn build_lead_role_prompt(params: &LeadPromptParams<'_>) -> String { ) } -fn render_teammate_list(teammates: &[TeamPromptAgent], renamed_agents: &HashMap) -> String { - if teammates.is_empty() { - return "(no teammates yet — propose the lineup to the user first, then use \ - team_spawn_agent only after they confirm or explicitly ask you to create \ - teammates immediately)" - .to_owned(); - } - - let mut out = String::with_capacity(teammates.len() * 64); - for (idx, teammate) in teammates.iter().enumerate() { - if idx > 0 { - out.push('\n'); - } - let status = teammate.status.as_deref().unwrap_or("unknown"); - let _ = write!(out, "- {} ({}, status: {})", teammate.name, teammate.backend, status); - if let Some(former) = renamed_agents.get(&teammate.slot_id) { - let _ = write!(out, " [formerly: {former}]"); - } - } - out -} - -fn render_available_types_section(agent_types: &[AvailableAgentType]) -> String { - let _ = agent_types; - String::new() -} - -fn render_available_assistants_section(assistants: &[AvailableAssistant]) -> String { - if assistants.is_empty() { - return String::new(); - } - let mut out = String::from("\n\n## Available Assistants for Spawning\n"); - out.push_str( - "These are available assistants with pre-loaded rules and skills for specific \ - domains (writing, research, PPT building, etc.). When a task matches an assistant's \ - specialty, prefer spawning that assistant — you get its domain \ - expertise automatically.\n\n", - ); - for (idx, assistant) in assistants.iter().enumerate() { - if idx > 0 { - out.push('\n'); - } - let desc = if assistant.description.is_empty() { - String::new() - } else { - format!(" — {}", assistant.description) - }; - let skills = if assistant.skills.is_empty() { - String::new() - } else { - format!("\n skills: {}", assistant.skills.join(", ")) - }; - let _ = write!( - out, - "- `{}` ({}){}{}", - assistant.assistant_id, assistant.name, desc, skills, - ); - } - out.push_str( - "\n\n### How to pick an assistant\n\ - 1. Scan the one-line descriptions and skills above. If one clearly matches the user's \ - domain (e.g. \"quarterly Word report\" → `word-creator`), spawn it directly with \ - `team_spawn_agent`.\n\ - 2. If two or more assistants seem relevant, call `team_describe_assistant` on each \ - candidate to see its full description, skills, and example tasks, then choose the best \ - fit.\n\ - 3. If no assistant is an obvious fit, choose the assistant whose domain and backing \ - capabilities best match the work.\n\n\ - Pass the assistant's ID as `assistant_id` to `team_spawn_agent`. The runtime backend \ - is derived automatically and does not need to be specified.", - ); - out -} - fn render_workspace_section(team_workspace: Option<&str>) -> String { match team_workspace { Some(workspace) => format!( @@ -189,7 +216,9 @@ fn render_workspace_section(team_workspace: Option<&str>) -> String { const TEAMMATE_PROMPT_TEMPLATE: &str = r#"# You are a Team Member ## Your Identity -Name: {{AGENT_NAME}}, Role: {{ROLE_DESC}} +Name: {{AGENT_NAME}} +Slot ID: {{AGENT_SLOT_ID}} +Role: teammate ## Conversation Style - If the user greets you, starts a new chat, or asks what you can do without assigning concrete work yet, reply warmly and naturally @@ -198,8 +227,7 @@ Name: {{AGENT_NAME}}, Role: {{ROLE_DESC}} ## Your Team Team: {{TEAM_NAME}} -Leader: {{LEADER_NAME}} -Teammates: {{TEAMMATES}}{{WORKSPACE}} +Leader: {{LEADER_NAME}} (slot_id: {{LEADER_SLOT_ID}}){{WORKSPACE}} ## Team Coordination Tools You MUST use the `team_*` MCP tools for ALL team coordination. @@ -208,6 +236,10 @@ TaskCreate, TaskUpdate). Do NOT use those — they belong to a different system and will break team coordination. Always use the `team_*` versions. Use `team_task_list` and `team_members` to check current team state. +Display names are only for user-facing text. For tool arguments such as +`team_send_message.to`, `team_rename_agent.slot_id`, and +`team_shutdown_agent.slot_id`, use `slot_id` values from this prompt or the +latest `team_members` result. Never pass display names as agent targets. ## How to Work 1. Read your unread messages to understand your assignment @@ -215,7 +247,7 @@ Use `team_task_list` and `team_members` to check current team state. 3. Use team_task_update to mark your task as "in_progress" when you start 4. Do the actual work (read files, write code, search, etc.) 5. When done, use team_task_update to mark the task "completed" -6. Use team_send_message to report results to the leader +6. Use team_send_message to report results to the leader slot_id ## Standing By (CRITICAL — read carefully) "Standing by" or "waiting" means **end your current turn**, not generate idle text in a live LLM stream. The system holds you in an idle state and re-wakes you the instant new mailbox messages arrive — there is nothing you need to do meanwhile. @@ -226,7 +258,7 @@ You are in a "standing by" situation when ANY of these is true: - You finished your current task and have nothing else assigned **The correct way to stand by:** -1. (Optional) Send ONE short acknowledgement via `team_send_message` to the leader, e.g. `"Acknowledged, standing by until reviewer-1 finishes"` or `"Ready, no task yet — standing by"` +1. (Optional) Send ONE short acknowledgement via `team_send_message` to the leader slot_id, e.g. `"Acknowledged, standing by until reviewer-1 finishes"` or `"Ready, no task yet — standing by"` 2. **STOP GENERATING.** Do NOT continue producing text like "I am waiting...", "still standing by...", reasoning loops, or repeated status updates. End your turn and return control. **Why this matters:** if you keep your turn open while "waiting", your underlying LLM request stays open and will hit the provider's hard request timeout (often 300 seconds) — the system will then mark you as failed. Ending the turn is the correct, lossless way to wait. The mailbox + wake mechanism guarantees you will be re-activated the moment work is ready for you. @@ -248,19 +280,7 @@ If you receive a message with type `shutdown_request`, the leader is asking you - Use your native tools (Read, Write, Bash, etc.) for implementation work"#; fn build_teammate_role_prompt(params: &TeammatePromptParams<'_>) -> String { - let teammates_section = if params.teammates.is_empty() { - "(none)".to_string() - } else { - params - .teammates - .iter() - .map(|teammate| match params.renamed_agents.get(&teammate.slot_id) { - Some(original) => format!("{} [formerly: {}]", teammate.name, original), - None => teammate.name.clone(), - }) - .collect::>() - .join(", ") - }; + let _ = (params.teammates, params.renamed_agents); let workspace_section = match params.team_workspace { Some(workspace) => format!( @@ -274,23 +294,13 @@ Always use the team workspace path for any project-related operations." TEAMMATE_PROMPT_TEMPLATE .replace("{{AGENT_NAME}}", ¶ms.agent.name) - .replace("{{ROLE_DESC}}", &role_description(¶ms.agent.backend)) + .replace("{{AGENT_SLOT_ID}}", ¶ms.agent.slot_id) .replace("{{TEAM_NAME}}", params.team_name) .replace("{{LEADER_NAME}}", ¶ms.leader.name) - .replace("{{TEAMMATES}}", &teammates_section) + .replace("{{LEADER_SLOT_ID}}", ¶ms.leader.slot_id) .replace("{{WORKSPACE}}", &workspace_section) } -fn role_description(agent_type: &str) -> String { - match agent_type.to_lowercase().as_str() { - "claude" => "general-purpose AI assistant".to_string(), - "gemini" => "Google Gemini AI assistant".to_string(), - "codex" => "code generation specialist".to_string(), - "qwen" => "Qwen AI assistant".to_string(), - other => format!("{other} AI assistant"), - } -} - #[cfg(test)] mod tests { use super::*; @@ -309,6 +319,7 @@ mod tests { #[test] fn lead_prompt_prepends_governance_and_fills_sections() { let renamed = HashMap::new(); + let leader = prompt_agent("lead-1", "Lead", TeamPromptRole::Lead); let teammate = prompt_agent("worker-1", "Worker", TeamPromptRole::Teammate); let assistants = vec![AvailableAssistant { assistant_id: "word-creator".to_owned(), @@ -318,6 +329,7 @@ mod tests { skills: vec!["docx".to_owned()], }]; let prompt = build_lead_prompt(&LeadPromptParams { + agent: &leader, team_name: "Alpha", teammates: &[teammate], available_agent_types: &[], @@ -327,9 +339,15 @@ mod tests { }); assert!(prompt.starts_with("## Team Governance")); - assert!(prompt.contains("- Worker (claude, status: unknown)")); - assert!(prompt.contains("## Available Assistants for Spawning")); - assert!(prompt.contains("Pass the assistant's ID as `assistant_id`")); + assert!(prompt.contains("Name: Lead")); + assert!(prompt.contains("Slot ID: lead-1")); + assert!(prompt.contains("Role: lead")); + assert!(!prompt.contains("## Your Teammates")); + assert!(!prompt.contains("## Available Assistants for Spawning")); + assert!(!prompt.contains("- Worker (claude, status: unknown)")); + assert!(prompt.to_lowercase().contains("first team turn")); + assert!(prompt.contains("team_members")); + assert!(prompt.contains("team_list_assistants")); assert!(!prompt.contains("${")); } @@ -347,8 +365,16 @@ mod tests { }); assert!(prompt.contains("## Team Governance")); + assert!(prompt.contains("Name: Worker")); + assert!(prompt.contains("Slot ID: worker-1")); + assert!(prompt.contains("Role: teammate")); + assert!(!prompt.contains("Role: general-purpose AI assistant")); assert!(prompt.contains("You MUST use the `team_*` MCP tools for ALL team coordination.")); - assert!(prompt.contains("Use team_send_message to report results to the leader")); + assert!(prompt.contains("Use team_send_message to report results to the leader slot_id")); + assert!(prompt.contains("Leader: Lead (slot_id: lead-1)")); + assert!(prompt.contains("Display names are only for user-facing text")); + assert!(prompt.contains("Never pass display names as agent targets")); assert!(prompt.contains("STOP GENERATING")); + assert!(!prompt.contains("Teammates: Worker")); } } diff --git a/crates/aionui-team-prompts/src/tools.rs b/crates/aionui-team-prompts/src/tools.rs index dad81bffd..0b49f4d02 100644 --- a/crates/aionui-team-prompts/src/tools.rs +++ b/crates/aionui-team-prompts/src/tools.rs @@ -31,27 +31,17 @@ Use this only when one of the following is true: Before calling this tool in the normal planning flow: - Start with one short sentence explaining why additional teammates would help - Tell the user which teammate(s) you recommend -- Present the proposal as a table with: name, responsibility, recommended assistant, and recommended model -- Include each teammate's responsibility, recommended assistant, and model +- Present the proposal as a table with: name, responsibility, and recommended assistant +- Include each teammate's responsibility and recommended assistant - Ask whether to create them as proposed or change any names, responsibilities, or assistant choices - In that approval question, remind the user that they can later ask you to replace or adjust any teammate if the lineup is not working well - Do NOT call this tool in that same turn; wait for explicit approval in a later user message When calling this tool, always provide assistant_id from the available assistants catalog. -When calling this tool, provide the model parameter if a specific model was recommended and approved. +Do not provide a model. The new teammate uses the selected assistant's configured/default model; users can adjust models from the UI model selector. The new agent will be created and added to the team. You can then assign tasks and send messages to it."#; -pub const TEAM_LIST_MODELS_DESCRIPTION: &str = - "Query available models for assistant backends. Returns the real-time model list that matches the frontend model selector. - -Use this to: -- Check what models are available before spawning an assistant-backed teammate with a specific model -- See all available backends and their models at once -- Verify a model ID is valid for the backend behind a chosen assistant or fallback backend - -Pass assistant_id to query models for a specific assistant, or omit it to see all backends."; - pub const TEAM_DESCRIBE_ASSISTANT_DESCRIPTION: &str = "Get detailed information about an assistant before spawning it as a teammate. @@ -59,7 +49,7 @@ Returns the assistant's full description, enabled skills, and example tasks so y judge whether it fits the user's request. Use this when two or more assistants look relevant from the one-line catalog in your system prompt. -Only works on assistants listed in \"Available Assistants for Spawning\". +Use team_list_assistants to find candidate assistant_id values. After confirming a match, call team_spawn_agent with the same assistant_id."; pub const TEAM_LIST_ASSISTANTS_DESCRIPTION: &str = "List the assistants available for team spawning. Returns the real assistant catalog with \ @@ -90,9 +80,7 @@ pub fn team_tool_specs() -> Vec { "type": "object", "properties": { "name": { "type": "string", "description": "Agent display name" }, - "model": { "type": "string", "description": "Specific model ID to use (e.g. \"claude-sonnet-4\"). Must be valid for the chosen assistant backend. Query team_list_models to see available models." }, - "assistant_id": { "type": "string", "description": "Assistant ID to spawn (from the Available Assistants catalog). The runtime backend is derived from this assistant." }, - "role": { "type": "string", "description": "Agent role (default: 'teammate')" } + "assistant_id": { "type": "string", "description": "Assistant ID to spawn. Call team_list_assistants when you need candidates; the runtime backend is derived from this assistant." } }, "required": ["name", "assistant_id"] }), @@ -131,10 +119,42 @@ pub fn team_tool_specs() -> Vec { TeamToolSpec { name: "team_task_list", permission: TeamToolPermission::AnyTeamAgent, - description: "List all tasks on the team task board.", + description: "List tasks on the team task board. Pass {} for the full board, or use owner/status/include_deleted/limit for filtered views.", input_schema: json!({ "type": "object", - "properties": {} + "additionalProperties": false, + "properties": { + "owner": { + "type": "string", + "description": "Return only tasks owned by this exact agent slot_id" + }, + "status": { + "description": "Return only tasks with these statuses. Accepts one status string or an array of status strings.", + "anyOf": [ + { + "type": "string", + "enum": ["pending", "in_progress", "completed", "deleted"] + }, + { + "type": "array", + "minItems": 1, + "items": { + "type": "string", + "enum": ["pending", "in_progress", "completed", "deleted"] + } + } + ] + }, + "include_deleted": { + "type": "boolean", + "description": "When status is omitted, include deleted tasks. Defaults to true so {} returns the full board." + }, + "limit": { + "type": "integer", + "minimum": 1, + "description": "Maximum number of returned tasks. Values above 200 are clamped to 200." + } + } }), }, TeamToolSpec { @@ -194,17 +214,6 @@ pub fn team_tool_specs() -> Vec { "required": ["assistant_id"] }), }, - TeamToolSpec { - name: "team_list_models", - permission: TeamToolPermission::AnyTeamAgent, - description: TEAM_LIST_MODELS_DESCRIPTION, - input_schema: json!({ - "type": "object", - "properties": { - "assistant_id": { "type": "string", "description": "Assistant ID to query. When provided, returns models for the backend behind that assistant. Shows all backends when omitted." } - } - }), - }, ] } @@ -271,7 +280,6 @@ mod tests { ("team_shutdown_agent", TeamToolPermission::LeadOnly), ("team_list_assistants", TeamToolPermission::AnyTeamAgent), ("team_describe_assistant", TeamToolPermission::AnyTeamAgent), - ("team_list_models", TeamToolPermission::AnyTeamAgent), ] ); } @@ -290,4 +298,21 @@ mod tests { assert!(!props.contains_key("backend")); assert!(required_names.contains(&"assistant_id")); } + + #[test] + fn task_list_schema_exposes_bounded_filter_arguments() { + let descriptor = visible_team_tool_descriptors(false) + .into_iter() + .find(|tool| tool.name == "team_task_list") + .expect("team_task_list descriptor"); + let props = descriptor.input_schema["properties"].as_object().unwrap(); + + assert!(props.contains_key("owner")); + assert!(props.contains_key("status")); + assert!(props.contains_key("include_deleted")); + assert!(props.contains_key("limit")); + assert_eq!(descriptor.input_schema["additionalProperties"], false); + assert!(props["limit"].get("maximum").is_none()); + assert_eq!(props["limit"]["minimum"], 1); + } } diff --git a/crates/aionui-team/Cargo.toml b/crates/aionui-team/Cargo.toml index e1736a394..e44e4b84a 100644 --- a/crates/aionui-team/Cargo.toml +++ b/crates/aionui-team/Cargo.toml @@ -28,6 +28,7 @@ sqlx.workspace = true tokio = { workspace = true, features = ["test-util"] } futures-util.workspace = true reqwest.workspace = true +tempfile.workspace = true # Enable the `AgentInstance::Mock` variant so tests can build fake agents # through the trait-object escape hatch without spawning real CLI processes. aionui-ai-agent = { workspace = true, features = ["test-support"] } diff --git a/crates/aionui-team/src/capability.rs b/crates/aionui-team/src/capability.rs index 198956b83..d3e27c82f 100644 --- a/crates/aionui-team/src/capability.rs +++ b/crates/aionui-team/src/capability.rs @@ -4,7 +4,7 @@ use aionui_common::constants::is_team_capable; /// Determine if a backend supports team mode. /// /// Hard whitelist always passes. For non-whitelisted backends, checks the -/// persisted `agent_capabilities` JSON for MCP transport declarations. +/// persisted `agent_capabilities` JSON for MCP capability metadata. pub fn is_team_capable_backend(backend: &str, agent_capabilities: Option<&serde_json::Value>) -> bool { is_team_capable(backend, agent_capabilities) } @@ -34,6 +34,12 @@ mod tests { let caps_mcp = json!({"mcp": {"stdio": true}}); assert!(is_team_capable_backend("goose", Some(&caps_mcp))); + + let caps_disabled_transport = json!({"mcp_capabilities": {"http": false, "sse": false}}); + assert!(is_team_capable_backend("zed", Some(&caps_disabled_transport))); + + let caps_empty = json!({"mcp_capabilities": {}}); + assert!(is_team_capable_backend("cursor", Some(&caps_empty))); } #[test] diff --git a/crates/aionui-team/src/error.rs b/crates/aionui-team/src/error.rs index 915f669c0..9a3d68f1a 100644 --- a/crates/aionui-team/src/error.rs +++ b/crates/aionui-team/src/error.rs @@ -14,9 +14,6 @@ pub enum TeamError { #[error("Invalid request: {0}")] InvalidRequest(String), - #[error("Team slot is busy: {0}")] - SlotBusy(String), - #[error("Leader-only action: {0}")] LeaderOnly(String), @@ -35,6 +32,17 @@ pub enum TeamError { #[error("Agent name already taken: {0}")] DuplicateAgentName(String), + #[error("Team agent runtime is not ready for conversation: {conversation_id}")] + RuntimeNotReady { conversation_id: String }, + + #[error("Team member runtime failed: {slot_id}")] + MemberRuntimeFailed { + team_id: String, + slot_id: String, + conversation_id: String, + public_reason: String, + }, + #[error("Workspace path is unavailable: {0}")] WorkspacePathUnavailable(String), @@ -92,6 +100,16 @@ pub fn classify_public_error(message: &str) -> Option { } } + if message == "model is no longer accepted; use the assistant configuration or UI model selector" { + return Some(TeamPublicError::new( + "TEAM_ASSISTANT_FIELD_UNSUPPORTED", + Some(json!({ + "field": "model", + "required_field": "assistant_id", + })), + )); + } + if message == "team_list_assistants does not accept arguments" { return Some(TeamPublicError::new( "TEAM_TOOL_ARGUMENTS_NOT_ALLOWED", @@ -117,10 +135,6 @@ mod tests { assert_eq!(TeamError::TeamNotFound("t1".into()).to_string(), "Team not found: t1"); assert_eq!(TeamError::AgentNotFound("s1".into()).to_string(), "Agent not found: s1"); assert_eq!(TeamError::TaskNotFound("tk1".into()).to_string(), "Task not found: tk1"); - assert_eq!( - TeamError::SlotBusy("lead-1".into()).to_string(), - "Team slot is busy: lead-1" - ); } #[test] @@ -142,5 +156,17 @@ mod tests { "required_field": "assistant_id", })) ); + + let model = + classify_public_error("model is no longer accepted; use the assistant configuration or UI model selector") + .expect("model field"); + assert_eq!(model.code, "TEAM_ASSISTANT_FIELD_UNSUPPORTED"); + assert_eq!( + model.details, + Some(json!({ + "field": "model", + "required_field": "assistant_id", + })) + ); } } diff --git a/crates/aionui-team/src/event_loop.rs b/crates/aionui-team/src/event_loop.rs index 3c127864d..d65299722 100644 --- a/crates/aionui-team/src/event_loop.rs +++ b/crates/aionui-team/src/event_loop.rs @@ -1,8 +1,5 @@ -use std::sync::{ - Arc, - atomic::{AtomicBool, Ordering}, -}; -use std::time::Duration; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex, MutexGuard}; use dashmap::DashMap; use dashmap::mapref::entry::Entry; @@ -10,25 +7,33 @@ use tokio::sync::Notify; use tokio::task::JoinHandle; use tracing::{debug, info, warn}; +use aionui_api_types::{TeamChildTurnPayload, TeamRunStatus}; + +use crate::events::{TEAM_CHILD_TURN_CANCELLED_EVENT, TEAM_CHILD_TURN_COMPLETED_EVENT, TEAM_CHILD_TURN_STARTED_EVENT}; use crate::mailbox::Mailbox; use crate::ports::{ AgentTurnExecutionError, AgentTurnExecutionPort, AgentTurnRequest, AgentTurnSource, AgentTurnStarted, AgentTurnStartedCallback, }; use crate::scheduler::TeammateManager; -use crate::session::TeamSession; -use crate::team_run::{ActiveChildTurn, ChildStartDecision, target_role_for}; +use crate::session::{PrepareBatchResult, TeamSession, WakeInput}; +use crate::team_run::target_role_for; use crate::types::TeammateStatus; -use crate::wake::TeamWakeSource; -use aionui_api_types::TeamRunStatus; +use crate::work_coordinator::{CommitResult, StartCommitResult, WorkBatch}; +use crate::work_source::WorkSource; -/// Registry of per-agent Notify handles. Used by any trigger source to poke -/// an agent's event loop without needing to know its internals. pub struct EventLoopRegistry { notifiers: DashMap>, handles: DashMap>, shutdown_tx: tokio::sync::watch::Sender, shutdown_rx: tokio::sync::watch::Receiver, + lifecycle: Mutex, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum EventLoopRegistrationError { + Duplicate, + Stopped, } impl Default for EventLoopRegistry { @@ -45,23 +50,30 @@ impl EventLoopRegistry { handles: DashMap::new(), shutdown_tx, shutdown_rx, + lifecycle: Mutex::new(false), } } - /// Check if an event loop is registered for this slot. pub fn has(&self, slot_id: &str) -> bool { self.notifiers.contains_key(slot_id) } - /// Poke the named agent's event loop so it drains its mailbox. + #[cfg(test)] + pub(crate) fn len(&self) -> usize { + self.notifiers.len() + } + pub fn notify(&self, slot_id: &str) { - if let Some(n) = self.notifiers.get(slot_id) { - n.notify_one(); + if let Some(notify) = self.notifiers.get(slot_id) { + notify.notify_one(); } } - /// Register and spawn an event loop for one agent. - pub fn spawn(&self, slot_id: &str, ctx: AgentLoopContext) -> bool { + pub fn spawn(&self, slot_id: &str, ctx: AgentLoopContext) -> Result<(), EventLoopRegistrationError> { + let stopped = self.lock_lifecycle(); + if *stopped { + return Err(EventLoopRegistrationError::Stopped); + } let notify = Arc::new(Notify::new()); match self.notifiers.entry(slot_id.to_owned()) { Entry::Occupied(_) => { @@ -70,7 +82,7 @@ impl EventLoopRegistry { slot_id, "agent event loop registration ignored because slot is already registered" ); - return false; + return Err(EventLoopRegistrationError::Duplicate); } Entry::Vacant(entry) => { entry.insert(notify.clone()); @@ -78,19 +90,23 @@ impl EventLoopRegistry { } let handle = tokio::spawn(run_event_loop(notify, self.shutdown_rx.clone(), ctx)); self.handles.insert(slot_id.to_owned(), handle); - true + Ok(()) } - /// Remove an agent's event loop (agent removed from team). pub fn remove(&self, slot_id: &str) { + let _lifecycle = self.lock_lifecycle(); self.notifiers.remove(slot_id); if let Some((_, handle)) = self.handles.remove(slot_id) { handle.abort(); } } - /// Shut down all event loops. pub fn shutdown(&self) { + let mut stopped = self.lock_lifecycle(); + if *stopped { + return; + } + *stopped = true; let _ = self.shutdown_tx.send(true); for entry in self.handles.iter() { entry.value().abort(); @@ -98,9 +114,12 @@ impl EventLoopRegistry { self.handles.clear(); self.notifiers.clear(); } + + fn lock_lifecycle(&self) -> MutexGuard<'_, bool> { + self.lifecycle.lock().unwrap_or_else(|poisoned| poisoned.into_inner()) + } } -/// Context shared across all iterations of one agent's event loop. pub struct AgentLoopContext { pub team_id: String, pub slot_id: String, @@ -109,26 +128,13 @@ pub struct AgentLoopContext { pub scheduler: Arc, pub mailbox: Arc, pub turn_port: Arc, - /// Used to notify other agents' event loops (e.g. leader after all-settled). pub registry: Arc, } -struct TurnExecution { - finish_ok: bool, - team_run_id: Option, - turn_id: Option, -} - fn is_retryable_start_skip(error: &AgentTurnExecutionError) -> bool { matches!(error, AgentTurnExecutionError::Skipped { reason } if reason.contains("already running")) } -/// The event loop for one agent slot. Spawned as a tokio task. -/// -/// Flow: -/// 1. Wait for signal (notify) or shutdown. -/// 2. Drain loop: compute_wake_input → has messages → send_message (blocking) → finalize → repeat. -/// 3. When mailbox empty → back to step 1. async fn run_event_loop( notify: Arc, mut shutdown_rx: tokio::sync::watch::Receiver, @@ -139,316 +145,265 @@ async fn run_event_loop( slot_id = %ctx.slot_id, "agent event loop started" ); - loop { - // Step 1: wait for signal or shutdown tokio::select! { biased; - _ = shutdown_rx.wait_for(|v| *v) => { - info!( - team_id = %ctx.team_id, - slot_id = %ctx.slot_id, - "agent event loop shutting down" - ); - return; - } + _ = shutdown_rx.wait_for(|stopped| *stopped) => return, _ = notify.notified() => {} } - // Drain loop: keep processing until mailbox is empty loop { if *shutdown_rx.borrow() { return; } - - let input = match ctx.session.compute_wake_input(&ctx.slot_id).await { - Ok(Some(input)) => input, - Ok(None) => break, - Err(e) => { + match ctx.session.prepare_next_batch(&ctx.slot_id).await { + Ok(PrepareBatchResult::Execute { batch, input }) => { + if execute_and_finalize(&ctx, *batch, input).await == ExecuteResult::WaitForSignal { + break; + } + } + Ok(PrepareBatchResult::SettleSignals { intent_ids }) => { + ctx.session.handle_signal_intents(&ctx.slot_id, &intent_ids).await; + } + Ok( + PrepareBatchResult::WaitingForCompletion + | PrepareBatchResult::Blocked + | PrepareBatchResult::Quiescent, + ) => break, + Err(error) => { warn!( team_id = %ctx.team_id, slot_id = %ctx.slot_id, - error = %e, - "event loop: compute_wake_input failed" + error = %error, + "event loop reconcile failed" ); - tokio::time::sleep(Duration::from_secs(1)).await; break; } - }; - - if !input.should_send { - ctx.session - .team_run_manager() - .record_empty_wake_observed(&ctx.slot_id) - .await; - break; - } - - match execute_turn(&ctx, &input).await { - Some(turn) => finalize_turn(&ctx, turn, &input).await, - None => break, // Turn not started (guard/runtime preparation); retry on next signal } } } } -/// Execute one agent turn through the Team-defined port. Conversation/runtime -/// lifecycle remains behind the port; Team keeps projection, mark-read, and -/// scheduler finalization here. -async fn execute_turn(ctx: &AgentLoopContext, input: &crate::session::WakeInput) -> Option { - let role = target_role_for(input.agent_role); - let reservation = if input.team_run_id.is_some() { - match ctx - .session - .team_run_manager() - .claim_wake_for_turn(&ctx.slot_id, role.clone(), &input.conversation_id) - .await - { - Some(reservation) => { - if let Some(source) = input.wake_source { - info!( - team_id = %ctx.team_id, - team_run_id = ?input.team_run_id, - slot_id = %ctx.slot_id, - wake_source = %source, - trigger_message_id = ?input.trigger_message_id.as_deref(), - "team priority wake claimed" - ); - } - Some(reservation) - } - None => { - warn!( - team_id = %ctx.team_id, - team_run_id = ?input.team_run_id, - slot_id = %ctx.slot_id, - conversation_id = %input.conversation_id, - "event loop: team run wake skipped because reservation could not be claimed" - ); - if let Err(e) = ctx.scheduler.set_status(&ctx.slot_id, TeammateStatus::Idle).await { - warn!( - team_id = %ctx.team_id, - slot_id = %ctx.slot_id, - error = %e, - "event loop: failed to roll back status after reservation claim failure" - ); - } - return None; - } - } - } else { - None - }; - - ctx.session.mirror_unread_to_conversation(input).await; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ExecuteResult { + ContinueDraining, + WaitForSignal, +} +async fn execute_and_finalize(ctx: &AgentLoopContext, batch: WorkBatch, input: WakeInput) -> ExecuteResult { + ctx.session.mirror_unread_to_conversation(&input).await; let _ = ctx.scheduler.set_status(&ctx.slot_id, TeammateStatus::Working).await; - let files: Vec = input + let files = input .unread .iter() - .filter_map(|m| m.files.as_ref()) + .filter_map(|message| message.files.as_ref()) .flatten() .cloned() .collect(); - - let unread_message_ids = input.unread.iter().map(|m| m.id.clone()).collect::>(); + let unread_message_ids = input + .unread + .iter() + .map(|message| message.id.clone()) + .collect::>(); let started_seen = Arc::new(AtomicBool::new(false)); - let on_started: Option = reservation.clone().map(|reservation| { - let team_run_manager = ctx.session.team_run_manager().clone(); - let cancellation_port = ctx.session.cancellation_port().clone(); - let user_id = ctx.user_id.clone(); - let started_seen = started_seen.clone(); - Arc::new(move |started: AgentTurnStarted| { - let team_run_manager = team_run_manager.clone(); - let cancellation_port = cancellation_port.clone(); - let user_id = user_id.clone(); - let reservation_id = reservation.reservation_id.clone(); - started_seen.store(true, Ordering::SeqCst); - Box::pin(async move { - let child = ActiveChildTurn { - team_run_id: started.team_run_id, - slot_id: started.slot_id, - role: started.role, - conversation_id: started.conversation_id, - turn_id: started.turn_id, - started_at_ms: aionui_common::now_ms(), - last_slow_notified_at_ms: None, - }; - match team_run_manager - .record_child_started(&reservation_id, child.clone()) - .await - { - ChildStartDecision::Accepted => {} - ChildStartDecision::CancelImmediately(child) => { - if let Err(err) = cancellation_port - .cancel_agent_turn(&user_id, &child.conversation_id, &child.turn_id) - .await - { - warn!( - team_run_id = %child.team_run_id, - slot_id = %child.slot_id, - conversation_id = %child.conversation_id, - turn_id = %child.turn_id, - error = %err, - "event loop: late-start child cancellation failed" - ); - } - team_run_manager.record_child_cancelled(&child).await; - team_run_manager.try_complete_cancelled().await; + let coordinator = ctx.session.work_coordinator().clone(); + let cancellation_port = ctx.session.cancellation_port().clone(); + let user_id = ctx.user_id.clone(); + let conversation_id = input.conversation_id.clone(); + let callback_batch = batch.clone(); + let callback_started = started_seen.clone(); + let callback_emitter = ctx.session.team_event_emitter(); + let on_started: AgentTurnStartedCallback = Arc::new(move |started: AgentTurnStarted| { + let coordinator = coordinator.clone(); + let cancellation_port = cancellation_port.clone(); + let user_id = user_id.clone(); + let conversation_id = conversation_id.clone(); + let batch = callback_batch.clone(); + let callback_started = callback_started.clone(); + let emitter = callback_emitter.clone(); + Box::pin(async move { + callback_started.store(true, Ordering::SeqCst); + match coordinator.mark_started(&batch, &started.turn_id) { + StartCommitResult::Accepted => { + if let Some(team_run_id) = started.team_run_id { + emitter.broadcast_child_turn( + TEAM_CHILD_TURN_STARTED_EVENT, + TeamChildTurnPayload { + team_id: emitter.team_id().to_owned(), + team_run_id, + slot_id: started.slot_id, + role: started.role, + conversation_id: started.conversation_id, + turn_id: started.turn_id, + status: TeamRunStatus::Running, + }, + ); } - ChildStartDecision::Ignored => {} } - }) as std::pin::Pin + Send>> - }) as AgentTurnStartedCallback + StartCommitResult::CancelImmediately => { + if let Err(error) = cancellation_port + .cancel_agent_turn(&user_id, &conversation_id, &started.turn_id) + .await + { + warn!( + slot_id = %batch.slot_id, + batch_id = %batch.batch_id, + turn_id = %started.turn_id, + error = %error, + "late-start team batch cancellation failed" + ); + } + coordinator.cancel_batch(&batch, "late_start_cancelled"); + if let Some(team_run_id) = started.team_run_id { + emitter.broadcast_child_turn( + TEAM_CHILD_TURN_CANCELLED_EVENT, + TeamChildTurnPayload { + team_id: emitter.team_id().to_owned(), + team_run_id, + slot_id: started.slot_id, + role: started.role, + conversation_id: started.conversation_id, + turn_id: started.turn_id, + status: TeamRunStatus::Cancelled, + }, + ); + } + } + StartCommitResult::StaleOwner => {} + } + }) }); let request = AgentTurnRequest { - team_run_id: input.team_run_id.clone(), + team_run_id: batch.team_run_ids.first().cloned(), team_id: ctx.team_id.clone(), slot_id: ctx.slot_id.clone(), - role, + role: target_role_for(input.agent_role), conversation_id: input.conversation_id.clone(), user_id: ctx.user_id.clone(), - content: input.first_message.clone(), + content: input.first_message, files, source: AgentTurnSource::Mailbox { - unread_count: input.unread.len(), + unread_count: unread_message_ids.len(), unread_message_ids, }, - on_started, + on_started: Some(on_started), }; - info!( - team_id = %ctx.team_id, - team_run_id = ?input.team_run_id, - slot_id = %ctx.slot_id, - conversation_id = %input.conversation_id, - "event loop: agent turn port call started" - ); let outcome = match ctx.turn_port.run_agent_turn(request).await { Ok(outcome) => outcome, - Err(e) => { + Err(error) if !started_seen.load(Ordering::SeqCst) && is_retryable_start_skip(&error) => { + ctx.session.work_coordinator().retry_start(&batch, "already_running"); + let _ = ctx.scheduler.set_status(&ctx.slot_id, TeammateStatus::Idle).await; + return ExecuteResult::WaitForSignal; + } + Err(error) => { warn!( team_id = %ctx.team_id, - team_run_id = ?input.team_run_id, slot_id = %ctx.slot_id, - conversation_id = %input.conversation_id, - error = %e, - outcome = "failed", - "event loop: agent turn port call failed" + batch_id = %batch.batch_id, + error = %error, + "agent turn start failed" ); - if input.team_run_id.is_some() - && let Some(reservation) = reservation.as_ref() - { - if started_seen.load(Ordering::SeqCst) { - ctx.session.team_run_manager().complete_failed().await; - } else if is_retryable_start_skip(&e) { - ctx.session - .team_run_manager() - .retry_child_start_later(&reservation.reservation_id, &e.to_string()) - .await; - let _ = ctx.scheduler.set_status(&ctx.slot_id, TeammateStatus::Idle).await; - return None; - } else { - ctx.session - .team_run_manager() - .record_child_start_failed(&reservation.reservation_id, &e.to_string()) - .await; - } - } - return Some(TurnExecution { - finish_ok: false, - team_run_id: input.team_run_id.clone(), - turn_id: None, - }); + mark_batch_messages_read(ctx, &batch).await; + ctx.session.work_coordinator().fail_batch(&batch, "turn_start_failed"); + let _ = ctx.scheduler.set_status(&ctx.slot_id, TeammateStatus::Error).await; + return ExecuteResult::ContinueDraining; } }; - let turn_ok = outcome.status.is_success(); - info!( - team_id = %ctx.team_id, - team_run_id = ?input.team_run_id, - slot_id = %ctx.slot_id, - conversation_id = %outcome.conversation_id, - turn_id = %outcome.turn_id, - outcome = ?outcome.status, - "event loop: agent turn port call completed" - ); - - let msg_ids: Vec = input.unread.iter().map(|m| m.id.clone()).collect(); - if !msg_ids.is_empty() - && let Err(e) = ctx.mailbox.mark_read_batch(&msg_ids).await - { - warn!( - team_id = %ctx.team_id, - slot_id = %ctx.slot_id, - error = %e, - "event loop: mark_read_batch failed (non-fatal)" - ); - } - - Some(TurnExecution { - finish_ok: turn_ok, - team_run_id: input.team_run_id.clone(), - turn_id: Some(outcome.turn_id), - }) -} - -/// Finalize a completed turn: mark idle (or error), cascade to leader. -async fn finalize_turn(ctx: &AgentLoopContext, turn: TurnExecution, input: &crate::session::WakeInput) { - if !turn.finish_ok { + mark_batch_messages_read(ctx, &batch).await; + let terminal_status = if outcome.status.is_success() { + (ctx.session.work_coordinator().complete_batch(&batch) == CommitResult::Committed) + .then_some(TeamRunStatus::Completed) + } else { + let committed = ctx.session.work_coordinator().fail_batch(&batch, "turn_failed"); let _ = ctx.scheduler.set_status(&ctx.slot_id, TeammateStatus::Error).await; + (committed == CommitResult::Committed).then_some(TeamRunStatus::Failed) + }; + if let Some(status) = terminal_status { + let emitter = ctx.session.team_event_emitter(); + for team_run_id in &batch.team_run_ids { + emitter.broadcast_child_turn( + TEAM_CHILD_TURN_COMPLETED_EVENT, + TeamChildTurnPayload { + team_id: ctx.team_id.clone(), + team_run_id: team_run_id.clone(), + slot_id: ctx.slot_id.clone(), + role: target_role_for(input.agent_role), + conversation_id: outcome.conversation_id.clone(), + turn_id: outcome.turn_id.clone(), + status: status.clone(), + }, + ); + } } + match ctx.scheduler.finalize_turn(&ctx.slot_id, &[]).await { - Ok(Some(wake_target)) => { - if wake_target != ctx.slot_id - && let Err(e) = ctx - .session - .scheduler_wake_agent_for_team_work(&wake_target, TeamWakeSource::IdleNotification) - .await + Ok(Some(wake_target)) if wake_target != ctx.slot_id => { + if let Err(error) = ctx + .session + .enqueue_leader_settle_signal(&wake_target, WorkSource::IdleNotification) + .await { warn!( team_id = %ctx.team_id, slot_id = %ctx.slot_id, - wake_target = %wake_target, - error = %e, - "event loop: failed to wake leader after teammate finalize" + wake_target, + error = %error, + "leader settle signal enqueue failed" ); } } - Ok(None) => {} - Err(e) => { - warn!( - team_id = %ctx.team_id, - slot_id = %ctx.slot_id, - error = %e, - "event loop: finalize_turn failed" - ); - } + Ok(_) => {} + Err(error) => warn!( + team_id = %ctx.team_id, + slot_id = %ctx.slot_id, + error = %error, + "scheduler turn finalization failed" + ), } - if let (Some(_team_run_id), Some(turn_id)) = (turn.team_run_id, turn.turn_id) { - let status = if turn.finish_ok { - TeamRunStatus::Completed - } else { - TeamRunStatus::Failed - }; - ctx.session - .team_run_manager() - .record_child_completed(&ctx.slot_id, &turn_id, status) - .await; - ctx.session.team_run_manager().maybe_complete().await; + ExecuteResult::ContinueDraining +} + +async fn mark_batch_messages_read(ctx: &AgentLoopContext, batch: &WorkBatch) { + if batch.mailbox_message_ids.is_empty() { + return; } - let should_release_suppressed_wake = input.wake_source.is_some_and(TeamWakeSource::resumes_paused_slot); - if should_release_suppressed_wake { - let role = target_role_for(input.agent_role); - if ctx - .session - .team_run_manager() - .release_suppressed_wake_if_resumed(&ctx.slot_id, role) - .await - .is_some() - { - ctx.registry.notify(&ctx.slot_id); - } + if let Err(error) = ctx.mailbox.mark_read_batch(&batch.mailbox_message_ids).await { + warn!( + team_id = %ctx.team_id, + slot_id = %ctx.slot_id, + batch_id = %batch.batch_id, + error = %error, + "team batch mailbox terminal mark-read failed" + ); + } +} + +#[cfg(test)] +mod registry_lifecycle_tests { + use std::sync::{Arc, mpsc}; + use std::time::Duration; + + use super::EventLoopRegistry; + + #[test] + fn remove_waits_for_in_progress_registration_lifecycle() { + let registry = Arc::new(EventLoopRegistry::new()); + let lifecycle = registry.lock_lifecycle(); + let (calling_tx, calling_rx) = mpsc::channel(); + let (done_tx, done_rx) = mpsc::channel(); + let worker = Arc::clone(®istry); + let thread = std::thread::spawn(move || { + calling_tx.send(()).unwrap(); + worker.remove("worker-1"); + done_tx.send(()).unwrap(); + }); + calling_rx.recv().unwrap(); + assert!(done_rx.recv_timeout(Duration::from_millis(50)).is_err()); + drop(lifecycle); + done_rx.recv_timeout(Duration::from_secs(1)).unwrap(); + thread.join().unwrap(); } } diff --git a/crates/aionui-team/src/events.rs b/crates/aionui-team/src/events.rs index e80dc26c9..c33457c40 100644 --- a/crates/aionui-team/src/events.rs +++ b/crates/aionui-team/src/events.rs @@ -1,8 +1,8 @@ use std::sync::Arc; use aionui_api_types::{ - TeamAgentRemovedPayload, TeamAgentRenamedPayload, TeamAgentSpawnedPayload, TeamAgentStatusPayload, - TeamChildTurnPayload, TeamRunPayload, WebSocketMessage, + TeamAgentRemovedPayload, TeamAgentRenamedPayload, TeamAgentRuntimeStatus, TeamAgentRuntimeStatusPayload, + TeamAgentSpawnedPayload, TeamAgentStatusPayload, TeamChildTurnPayload, TeamRunPayload, WebSocketMessage, }; use aionui_realtime::EventBroadcaster; use tracing::debug; @@ -14,11 +14,12 @@ pub const TEAM_AGENT_STATUS_CHANGED_EVENT: &str = "team.agentStatusChanged"; pub const TEAM_AGENT_SPAWNED_EVENT: &str = "team.agentSpawned"; pub const TEAM_AGENT_REMOVED_EVENT: &str = "team.agentRemoved"; pub const TEAM_AGENT_RENAMED_EVENT: &str = "team.agentRenamed"; +pub const TEAM_AGENT_RUNTIME_STATUS_CHANGED_EVENT: &str = "team.agentRuntimeStatusChanged"; pub const TEAM_LIST_CHANGED_EVENT: &str = "team.listChanged"; pub const TEAM_CREATED_EVENT: &str = "team.created"; pub const TEAM_REMOVED_EVENT: &str = "team.removed"; pub const TEAM_RENAMED_EVENT: &str = "team.renamed"; -pub const TEAM_MCP_STATUS_EVENT: &str = "team.mcpStatus"; +pub const TEAM_SESSION_STATUS_CHANGED_EVENT: &str = "team.sessionStatusChanged"; pub const TEAM_TASK_CHANGED_EVENT: &str = "team.taskChanged"; pub const TEAM_SESSION_CHANGED_EVENT: &str = "team.sessionChanged"; pub const TEAM_RUN_ACCEPTED_EVENT: &str = "team.runAccepted"; @@ -95,6 +96,26 @@ impl TeamEventEmitter { self.broadcaster.broadcast(event); } + pub fn broadcast_agent_runtime_status( + &self, + agent: &TeamAgent, + status: TeamAgentRuntimeStatus, + error: Option, + ) { + let payload = TeamAgentRuntimeStatusPayload { + team_id: self.team_id.clone(), + slot_id: agent.slot_id.clone(), + conversation_id: agent.conversation_id.clone(), + status, + error, + }; + let event = WebSocketMessage::new( + TEAM_AGENT_RUNTIME_STATUS_CHANGED_EVENT, + serde_json::to_value(payload).expect("serialize agent runtime status payload"), + ); + self.broadcaster.broadcast(event); + } + pub fn broadcast_team_run(&self, event_name: &'static str, payload: TeamRunPayload) { debug!( event_name = event_name, @@ -103,9 +124,10 @@ impl TeamEventEmitter { target_slot_id = %payload.target_slot_id, target_role = ?payload.target_role, status = ?payload.status, - active_child_count = payload.active_child_count, - pending_wake_count = payload.pending_wake_count, - starting_child_count = payload.starting_child_count, + queued_intent_count = payload.queued_intent_count, + starting_batch_count = payload.starting_batch_count, + running_batch_count = payload.running_batch_count, + active_enqueue_lease_count = payload.active_enqueue_lease_count, slot_work_count = payload.slot_work.len(), "team websocket event emitted" ); @@ -141,7 +163,8 @@ mod tests { use super::*; use crate::types::TeammateRole; use aionui_api_types::{ - TeamAgentRemovedPayload, TeamAgentRenamedPayload, TeamAgentSpawnedPayload, TeamAgentStatusPayload, + TeamAgentRemovedPayload, TeamAgentRenamedPayload, TeamAgentRuntimeStatusPayload, TeamAgentSpawnedPayload, + TeamAgentStatusPayload, }; struct RecordingBroadcaster { @@ -215,6 +238,36 @@ mod tests { assert_eq!(payload.assistant.role, "teammate"); } + #[test] + fn agent_runtime_status_event_has_correct_shape() { + let (emitter, bc) = make_emitter(); + let agent = TeamAgent { + slot_id: "slot-runtime".into(), + name: "Worker".into(), + role: TeammateRole::Teammate, + conversation_id: "conv-runtime".into(), + backend: "acp".into(), + model: "claude".into(), + assistant_id: None, + status: Some(TeammateStatus::Idle), + conversation_type: None, + cli_path: None, + }; + + emitter.broadcast_agent_runtime_status(&agent, TeamAgentRuntimeStatus::Ready, None); + + let events = bc.events(); + assert_eq!(events.len(), 1); + assert_eq!(events[0].name, "team.agentRuntimeStatusChanged"); + + let payload: TeamAgentRuntimeStatusPayload = serde_json::from_value(events[0].data.clone()).unwrap(); + assert_eq!(payload.team_id, "team-1"); + assert_eq!(payload.slot_id, "slot-runtime"); + assert_eq!(payload.conversation_id, "conv-runtime"); + assert_eq!(payload.status, TeamAgentRuntimeStatus::Ready); + assert_eq!(payload.error, None); + } + #[test] fn removed_event_has_correct_shape() { let (emitter, bc) = make_emitter(); @@ -277,9 +330,10 @@ mod tests { target_slot_id: "lead-1".into(), target_role: aionui_api_types::TeamRunTargetRole::Lead, status: aionui_api_types::TeamRunStatus::Accepted, - active_child_count: 0, - pending_wake_count: 1, - starting_child_count: 0, + queued_intent_count: 1, + starting_batch_count: 0, + running_batch_count: 0, + active_enqueue_lease_count: 0, slot_work: Vec::new(), }, ); @@ -293,7 +347,7 @@ mod tests { assert_eq!(payload.team_run_id, "run-1"); assert_eq!(payload.target_role, aionui_api_types::TeamRunTargetRole::Lead); assert_eq!(payload.status, aionui_api_types::TeamRunStatus::Accepted); - assert_eq!(payload.starting_child_count, 0); + assert_eq!(payload.starting_batch_count, 0); } #[test] diff --git a/crates/aionui-team/src/lib.rs b/crates/aionui-team/src/lib.rs index 39194e620..e9c5eeaac 100644 --- a/crates/aionui-team/src/lib.rs +++ b/crates/aionui-team/src/lib.rs @@ -8,22 +8,24 @@ pub mod event_loop; pub mod events; pub mod mailbox; pub mod mcp; +mod member_runtime; pub mod message_projection; pub mod ports; +pub mod prompt_dump; pub mod prompts; pub mod provisioning; pub mod routes; pub mod scheduler; pub mod service; pub mod session; -mod slot_wake_gate; pub mod task_board; pub mod team_run; #[cfg(test)] pub(crate) mod test_utils; pub mod types; pub mod visibility; -mod wake; +mod work_coordinator; +mod work_source; mod workspace; pub use crash_detection::{CrashReason, detect_crash, is_rate_limited}; @@ -41,6 +43,7 @@ pub use ports::{ TeamAssistantCatalogPort, TeamConversationBindingLookup, TeamConversationLookupPort, }; +pub use prompt_dump::TeamPromptDumpConfig; pub use prompts::{build_lead_prompt, build_teammate_prompt, build_wake_payload}; pub use provisioning::{ TeamAgentProvisioner, TeamConversationCreateRequest, TeamConversationCreateResult, TeamConversationProvisioningPort, @@ -49,13 +52,10 @@ pub use routes::{TeamRouterState, team_routes}; pub use scheduler::{ SchedulerAction, TeammateManager, WAKE_TIMEOUT_MS, WakePayload, format_crash_testament, normalize_name, }; -pub use service::TeamSessionService; +pub use service::{TeamIdleCleanupCoordinator, TeamSessionService}; pub use session::{TeamSession, WakeInput}; pub use task_board::{TaskBoard, TaskUpdate}; -pub use team_run::{ - ActiveChildTurn, ChildCancelTarget, ChildStartDecision, StartingChildReservation, StartingReservationState, - TeamRunManager, target_role_for, -}; +pub use team_run::{TeamRunManager, target_role_for}; pub use types::{ MailboxMessage, MailboxMessageType, TaskStatus, Team, TeamAgent, TeamTask, TeammateRole, TeammateStatus, }; diff --git a/crates/aionui-team/src/mcp/server.rs b/crates/aionui-team/src/mcp/server.rs index e0e86cf18..f5d963c38 100644 --- a/crates/aionui-team/src/mcp/server.rs +++ b/crates/aionui-team/src/mcp/server.rs @@ -1,9 +1,7 @@ use std::net::SocketAddr; use std::sync::{Arc, Weak}; -use aionui_api_types::{ - TeamMcpPhase, TeamMcpStatusPayload, TeamSendMessageQueuedResponse, TeamSendMessageStatus, WebSocketMessage, -}; +use aionui_api_types::TeamSendMessageQueuedResponse; use aionui_realtime::EventBroadcaster; use serde_json::{Value, json}; use tokio::net::{TcpListener, TcpStream}; @@ -11,20 +9,20 @@ use tokio::sync::watch; use tracing::{debug, error, info, warn}; use crate::error::{TeamError, classify_public_error}; -use crate::events::TEAM_MCP_STATUS_EVENT; +use crate::prompt_dump::{TeamPromptDumpConfig, TeamToolsListDump, dump_team_tools_list}; use crate::scheduler::TeammateManager; use crate::service::TeamSessionService; use crate::session::{AgentMessageQueueResult, SpawnAgentRequest}; -use crate::types::{TeammateRole, TeammateStatus}; -use crate::wake::TeamWakeSource; +use crate::types::{TaskStatus, TeamAgent, TeamTask, TeammateRole, TeammateStatus}; +use crate::work_source::WorkSource; use super::protocol::{ INVALID_PARAMS, INVALID_REQUEST, JsonRpcResponse, METHOD_NOT_FOUND, PROTOCOL_VERSION, SERVER_NAME, SERVER_VERSION, read_request, write_response, }; use super::tools::{ - RenameAgentInput, SendMessageInput, ShutdownAgentInput, SpawnAgentInput, TaskCreateInput, TaskUpdateInput, - all_tool_descriptors_for_role, handle_team_list_models, + RenameAgentInput, SendMessageInput, ShutdownAgentInput, SpawnAgentInput, TaskCreateInput, TaskListInput, + TaskListStatusInput, TaskUpdateInput, all_tool_descriptors_for_role, }; // --------------------------------------------------------------------------- @@ -45,21 +43,29 @@ impl TeamMcpServer { team_id: String, broadcaster: Arc, service: Weak, + ) -> Result { + Self::start_with_prompt_dump( + auth_token, + scheduler, + team_id, + broadcaster, + service, + Some(TeamPromptDumpConfig::disabled()), + ) + .await + } + + pub async fn start_with_prompt_dump( + auth_token: String, + scheduler: Arc, + team_id: String, + _broadcaster: Arc, + service: Weak, + prompt_dump: Option, ) -> Result { let listener = match TcpListener::bind("127.0.0.1:0").await { Ok(l) => l, Err(e) => { - broadcast_mcp_status( - broadcaster.as_ref(), - TeamMcpStatusPayload { - team_id: team_id.clone(), - slot_id: String::new(), - phase: TeamMcpPhase::TcpError, - port: None, - server_count: None, - error: Some(e.to_string()), - }, - ); return Err(TeamError::InvalidRequest(format!("Failed to bind TCP: {e}"))); } }; @@ -67,30 +73,21 @@ impl TeamMcpServer { .local_addr() .map_err(|e| TeamError::InvalidRequest(format!("Failed to get local addr: {e}")))?; - broadcast_mcp_status( - broadcaster.as_ref(), - TeamMcpStatusPayload { - team_id: team_id.clone(), - slot_id: String::new(), - phase: TeamMcpPhase::TcpReady, - port: Some(addr.port()), - server_count: None, - error: None, - }, - ); - let (shutdown_tx, shutdown_rx) = watch::channel(false); let token = auth_token.clone(); let sched_for_tcp = scheduler.clone(); let service_for_tcp = service.clone(); let team_id_for_tcp = team_id.clone(); + let prompt_dump = prompt_dump.unwrap_or_else(TeamPromptDumpConfig::disabled); + let prompt_dump_for_tcp = prompt_dump.clone(); tokio::spawn(accept_loop( listener, token, sched_for_tcp, service_for_tcp, team_id_for_tcp, + prompt_dump_for_tcp, shutdown_rx.clone(), )); @@ -106,12 +103,14 @@ impl TeamMcpServer { let http_sched = scheduler.clone(); let http_service = service.clone(); let http_team_id = team_id.clone(); + let http_prompt_dump = prompt_dump.clone(); tokio::spawn(http_mcp_loop( http_listener, http_token, http_sched, http_service, http_team_id, + http_prompt_dump, shutdown_rx, )); @@ -153,14 +152,6 @@ impl Drop for TeamMcpServer { } } -fn broadcast_mcp_status(broadcaster: &dyn EventBroadcaster, payload: TeamMcpStatusPayload) { - let event = WebSocketMessage::new( - TEAM_MCP_STATUS_EVENT, - serde_json::to_value(payload).expect("serialize mcp status payload"), - ); - broadcaster.broadcast(event); -} - // --------------------------------------------------------------------------- // Accept loop // --------------------------------------------------------------------------- @@ -171,6 +162,7 @@ async fn accept_loop( scheduler: Arc, service: Weak, team_id: String, + prompt_dump: TeamPromptDumpConfig, mut shutdown_rx: watch::Receiver, ) { loop { @@ -183,7 +175,8 @@ async fn accept_loop( let sched = Arc::clone(&scheduler); let svc = service.clone(); let tid = team_id.clone(); - tokio::spawn(handle_connection(stream, token, sched, svc, tid)); + let dump = prompt_dump.clone(); + tokio::spawn(handle_connection(stream, token, sched, svc, tid, dump)); } Err(e) => { error!("Accept error: {e}"); @@ -210,6 +203,7 @@ async fn handle_connection( scheduler: Arc, service: Weak, team_id: String, + prompt_dump: TeamPromptDumpConfig, ) { let (mut reader, mut writer) = tokio::io::split(stream); @@ -250,6 +244,7 @@ async fn handle_connection( &service, &team_id, caller_slot_id.as_deref().unwrap_or("unknown"), + &prompt_dump, ) .await }; @@ -329,10 +324,11 @@ async fn handle_method( service: &Weak, team_id: &str, caller_slot_id: &str, + prompt_dump: &TeamPromptDumpConfig, ) -> JsonRpcResponse { match request.method.as_str() { "notifications/initialized" => JsonRpcResponse::success(request.id, json!({})), - "tools/list" => handle_tools_list(request.id, scheduler, caller_slot_id).await, + "tools/list" => handle_tools_list(request.id, scheduler, team_id, caller_slot_id, prompt_dump).await, "tools/call" => handle_tools_call(request, scheduler, service, team_id, caller_slot_id).await, _ => JsonRpcResponse::error( request.id, @@ -350,12 +346,48 @@ async fn caller_role_for_tools_list(scheduler: &TeammateManager, caller_slot_id: .unwrap_or(TeammateRole::Teammate) } -async fn handle_tools_list(id: Option, scheduler: &TeammateManager, caller_slot_id: &str) -> JsonRpcResponse { +async fn handle_tools_list( + id: Option, + scheduler: &TeammateManager, + team_id: &str, + caller_slot_id: &str, + prompt_dump: &TeamPromptDumpConfig, +) -> JsonRpcResponse { let caller_role = caller_role_for_tools_list(scheduler, caller_slot_id).await; let tools = all_tool_descriptors_for_role(caller_role); + let tools_for_dump = tool_descriptors_to_mcp_json(&tools); + if let Err(error) = dump_team_tools_list( + prompt_dump, + TeamToolsListDump { + team_id, + caller_slot_id, + caller_role, + tools: &tools_for_dump, + }, + ) { + warn!( + team_id, + caller_slot_id, + error = %error, + "team tools/list prompt dump failed" + ); + } JsonRpcResponse::success(id, json!({ "tools": tools })) } +fn tool_descriptors_to_mcp_json(tools: &[super::tools::ToolDescriptor]) -> Vec { + tools + .iter() + .map(|d| { + json!({ + "name": d.name, + "description": d.description, + "inputSchema": d.input_schema, + }) + }) + .collect() +} + // --------------------------------------------------------------------------- // tools/call dispatcher // --------------------------------------------------------------------------- @@ -491,14 +523,13 @@ pub(crate) async fn dispatch_tool( "team_spawn_agent" => exec_spawn_agent(arguments, service, team_id, caller_slot_id, caller_role).await, "team_task_create" => exec_task_create(arguments, scheduler).await, "team_task_update" => exec_task_update(arguments, scheduler).await, - "team_task_list" => exec_task_list(scheduler).await, + "team_task_list" => exec_task_list(arguments, scheduler).await, "team_members" => exec_members(scheduler).await, "team_rename_agent" => exec_rename_agent(arguments, scheduler, service, team_id).await, "team_shutdown_agent" => { exec_shutdown_agent(arguments, scheduler, service, team_id, caller_slot_id, caller_role).await } "team_list_assistants" => exec_list_assistants(arguments, service).await, - "team_list_models" => exec_list_models(arguments, service).await, "team_describe_assistant" => exec_describe_assistant(arguments, service).await, _ => Err(ToolCallError::from_message(format!("Unknown tool: {tool_name}"))), } @@ -519,29 +550,6 @@ async fn exec_list_assistants(args: &Value, service: &Weak) serde_json::to_string_pretty(&value).map_err(|e| ToolCallError::from_message(format!("Serialization error: {e}"))) } -async fn exec_list_models(args: &Value, service: &Weak) -> Result { - if args.get("backend").is_some() { - return Err(ToolCallError::from_message( - "backend is no longer accepted; use assistant_id", - )); - } - if args.get("agent_type").is_some() { - return Err(ToolCallError::from_message( - "agent_type is no longer accepted; use assistant_id", - )); - } - let assistant_id_filter = args.get("assistant_id").and_then(Value::as_str); - - let value = match service.upgrade() { - Some(svc) => svc - .list_models_from_db(assistant_id_filter) - .await - .map_err(|error| ToolCallError::from_message(error.to_string()))?, - None => handle_team_list_models(args), - }; - serde_json::to_string_pretty(&value).map_err(|e| ToolCallError::from_message(format!("Serialization error: {e}"))) -} - async fn exec_describe_assistant(args: &Value, service: &Weak) -> Result { if args.get("custom_agent_id").is_some() { return Err(ToolCallError::from_message( @@ -569,17 +577,23 @@ async fn exec_describe_assistant(args: &Value, service: &Weak Result { +async fn resolve_agent_target( + scheduler: &TeammateManager, + target: &str, + allow_broadcast: bool, +) -> Result { let agents = scheduler.list_agents().await; if agents.iter().any(|a| a.slot_id == target) { return Ok(target.to_owned()); } - let query = target.to_lowercase(); - let hits: Vec<_> = agents.iter().filter(|a| a.name.to_lowercase() == query).collect(); - match hits.len() { - 0 => Err(format!("No agent matches '{target}'")), - 1 => Ok(hits[0].slot_id.clone()), - _ => Err(format!("Multiple agents match '{target}'")), + if allow_broadcast { + Err(format!( + "Invalid agent target '{target}': expected slot_id or \"*\". Call team_members to get slot_id." + )) + } else { + Err(format!( + "Invalid agent target '{target}': expected slot_id. Call team_members to get slot_id." + )) } } @@ -625,25 +639,30 @@ async fn exec_send_message( .map_err(|e| ToolCallError::from_message(e.to_string()))?; if let Some(svc) = service.upgrade() && let Err(e) = svc - .wake_leader_after_recovery_message(team_id, caller_slot_id, TeamWakeSource::ShutdownRejected) + .wake_leader_after_recovery_message(team_id, caller_slot_id, WorkSource::ShutdownRejected) .await { warn!( team_id, slot_id = %caller_slot_id, - wake_source = %TeamWakeSource::ShutdownRejected, + wake_source = %WorkSource::ShutdownRejected, error = %e, "failed to apply shutdown_rejected wake policy" ); } debug!(from = caller_slot_id, reason, "shutdown_rejected handled"); - return Ok(format!("shutdown_rejected: {reason}")); + return Ok(json!({ + "status": "ok", + "action": "shutdown_rejected", + "reason": reason, + }) + .to_string()); } let resolved_to = if input.to == "*" { "*".to_owned() } else { - resolve_agent_target(scheduler, &input.to) + resolve_agent_target(scheduler, &input.to, true) .await .map_err(ToolCallError::from_message)? }; @@ -684,14 +703,86 @@ fn build_send_message_queued_response( .first() .ok_or_else(|| "No message targets resolved".to_string())?; Ok(TeamSendMessageQueuedResponse { - status: TeamSendMessageStatus::Queued, - delivery: first.delivery.clone(), - reason: first.target.queue_state.clone(), team_run_id: first.team_run_id.clone(), - targets: target_results.into_iter().map(|result| result.target).collect(), + target: first.target.clone(), + }) +} + +fn agent_json(agent: &TeamAgent) -> Value { + let status = agent.status.unwrap_or(TeammateStatus::Idle); + json!({ + "slot_id": agent.slot_id, + "name": agent.name, + "role": agent.role, + "status": status, + "assistant_id": agent.assistant_id, + "model": agent.model, + }) +} + +fn task_json(task: &TeamTask) -> Value { + json!({ + "task_id": task.id, + "subject": task.subject, + "status": task.status, + "owner": task.owner, + "blocked_by": task.blocked_by, + }) +} + +const MAX_TASK_LIST_LIMIT: usize = 200; + +#[derive(Debug)] +struct TaskListFilters { + owner: Option, + statuses: Option>, + include_deleted: bool, + limit: Option, +} + +fn parse_task_list_filters(args: &Value) -> Result { + let input: TaskListInput = serde_json::from_value(args.clone()) + .map_err(|e| ToolCallError::from_message(format!("Invalid params: {e}")))?; + let statuses = match input.status { + Some(TaskListStatusInput::Single(status)) => Some(vec![parse_task_status_arg(&status)?]), + Some(TaskListStatusInput::Many(statuses)) => { + if statuses.is_empty() { + return Err(ToolCallError::from_message("Invalid params: status must not be empty")); + } + let parsed = statuses + .iter() + .map(|status| parse_task_status_arg(status)) + .collect::, _>>()?; + Some(parsed) + } + None => None, + }; + let limit = match input.limit { + Some(value) if value <= 0 => { + return Err(ToolCallError::from_message( + "Invalid params: limit must be greater than 0", + )); + } + Some(value) => Some((value as usize).min(MAX_TASK_LIST_LIMIT)), + None => None, + }; + Ok(TaskListFilters { + owner: input.owner, + statuses, + include_deleted: input.include_deleted.unwrap_or(true), + limit, }) } +fn parse_task_status_arg(status: &str) -> Result { + TaskStatus::parse(status) + .ok_or_else(|| ToolCallError::from_message(format!("Invalid params: unsupported task status '{status}'"))) +} + +fn json_text(value: &Value) -> Result { + serde_json::to_string_pretty(value).map_err(|e| ToolCallError::from_message(format!("Serialization error: {e}"))) +} + async fn exec_spawn_agent( args: &Value, service: &Weak, @@ -716,6 +807,11 @@ async fn exec_spawn_agent( "agent_type is no longer accepted; use assistant_id", )); } + if args.get("model").is_some() { + return Err(ToolCallError::from_message( + "model is no longer accepted; use the assistant configuration or UI model selector", + )); + } let input: SpawnAgentInput = serde_json::from_value(args.clone()) .map_err(|e| ToolCallError::from_message(format!("Invalid params: {e}")))?; @@ -734,7 +830,6 @@ async fn exec_spawn_agent( let req = SpawnAgentRequest { name: requested_name.clone(), assistant_id: Some(assistant_id), - model: input.model, }; let service = service @@ -744,7 +839,14 @@ async fn exec_spawn_agent( service .spawn_agent_in_session(team_id, caller_slot_id, req) .await - .map(|agent| format!("Agent '{}' spawned (slot_id={})", agent.name, agent.slot_id)) + .and_then(|agent| { + serde_json::to_string_pretty(&json!({ + "status": "ok", + "action": "agent_spawned", + "agent": agent_json(&agent), + })) + .map_err(TeamError::Json) + }) .map_err(|e| ToolCallError::from_message(e.to_string())) } @@ -752,46 +854,54 @@ async fn exec_task_create(args: &Value, scheduler: &TeammateManager) -> Result Result { let input: TaskUpdateInput = serde_json::from_value(args.clone()) .map_err(|e| ToolCallError::from_message(format!("Invalid params: {e}")))?; - let action = crate::scheduler::SchedulerAction::TaskUpdate { - task_id: input.task_id.clone(), - status: input.status, - description: input.description, - owner: input.owner, - blocked_by: input.blocked_by, - }; - scheduler - .execute_action("system", &action) + let task = scheduler + .update_task( + &input.task_id, + input.status.as_deref(), + input.description, + input.owner, + input.blocked_by, + ) .await .map_err(|e| ToolCallError::from_message(e.to_string()))?; - Ok(format!("Task '{}' updated", input.task_id)) + json_text(&json!({ "status": "ok", "task": task_json(&task) })) } -async fn exec_task_list(scheduler: &TeammateManager) -> Result { +async fn exec_task_list(args: &Value, scheduler: &TeammateManager) -> Result { + let filters = parse_task_list_filters(args)?; let tasks = scheduler .list_tasks() .await .map_err(|e| ToolCallError::from_message(e.to_string()))?; - let output: Vec = tasks + let mut output: Vec = tasks .iter() + .filter(|task| match filters.owner.as_deref() { + Some(owner) => task.owner.as_deref() == Some(owner), + None => true, + }) + .filter(|task| match filters.statuses.as_ref() { + Some(statuses) => statuses.contains(&task.status), + None if !filters.include_deleted => task.status != TaskStatus::Deleted, + None => true, + }) .map(|t| { json!({ "id": t.id, @@ -804,32 +914,19 @@ async fn exec_task_list(scheduler: &TeammateManager) -> Result Result { - let agents = scheduler.list_agents().await; - let output: Vec = agents - .iter() - .map(|a| { - // `TeamAgent::status` is `None` for cold-start agents that have not - // yet transitioned through `set_status` (e.g. the lead before its - // first wake). The scheduler already tracks them as `Idle` - // internally (see `TeammateManager::new`), and AionUi's - // TeammateManager exposes `'idle'` as the initial value. Mirror - // that here so MCP clients never see `null` and misread a live - // teammate as offline. - let status = a.status.unwrap_or(TeammateStatus::Idle); - json!({ - "slot_id": a.slot_id, - "name": a.name, - "role": a.role, - "status": status, - "backend": a.backend, - "model": a.model, - }) - }) - .collect(); + let mut agents = scheduler.list_agents().await; + agents.sort_by_key(|agent| match agent.role { + TeammateRole::Lead => 0, + TeammateRole::Teammate => 1, + }); + let output: Vec = agents.iter().map(agent_json).collect(); serde_json::to_string_pretty(&output).map_err(|e| ToolCallError::from_message(format!("Serialization error: {e}"))) } @@ -842,7 +939,7 @@ async fn exec_rename_agent( let input: RenameAgentInput = serde_json::from_value(args.clone()) .map_err(|e| ToolCallError::from_message(format!("Invalid params: {e}")))?; - let resolved_slot = resolve_agent_target(scheduler, &input.slot_id) + let resolved_slot = resolve_agent_target(scheduler, &input.slot_id, false) .await .map_err(ToolCallError::from_message)?; @@ -861,7 +958,15 @@ async fn exec_rename_agent( .map_err(|e| ToolCallError::from_message(e.to_string()))?; } - Ok(format!("Agent '{}' renamed to '{}'", input.slot_id, input.new_name)) + let agent = scheduler + .get_agent(&resolved_slot) + .await + .map_err(|e| ToolCallError::from_message(e.to_string()))?; + json_text(&json!({ + "status": "ok", + "action": "agent_renamed", + "agent": agent_json(&agent), + })) } async fn exec_shutdown_agent( @@ -878,18 +983,24 @@ async fn exec_shutdown_agent( let input: ShutdownAgentInput = serde_json::from_value(args.clone()) .map_err(|e| ToolCallError::from_message(format!("Invalid params: {e}")))?; - let target_slot_id = resolve_agent_target(scheduler, &input.slot_id) + let target_slot_id = resolve_agent_target(scheduler, &input.slot_id, false) .await .map_err(ToolCallError::from_message)?; let service = service .upgrade() .ok_or_else(|| ToolCallError::from_message("Team service not available; cannot wake shutdown target"))?; + let reason = input.reason.clone(); service .shutdown_agent_in_session(team_id, caller_slot_id, &target_slot_id, input.reason) .await .map_err(|e| ToolCallError::from_message(e.to_string()))?; - Ok(format!("Shutdown request sent to agent '{}'", target_slot_id)) + json_text(&json!({ + "status": "ok", + "action": "shutdown_requested", + "target_slot_id": target_slot_id, + "reason": reason, + })) } // --------------------------------------------------------------------------- @@ -902,6 +1013,7 @@ async fn http_mcp_loop( scheduler: Arc, service: Weak, team_id: String, + prompt_dump: TeamPromptDumpConfig, mut shutdown_rx: watch::Receiver, ) { use tokio::io::{AsyncReadExt, AsyncWriteExt}; @@ -915,6 +1027,7 @@ async fn http_mcp_loop( let sched = scheduler.clone(); let svc = service.clone(); let tid = team_id.clone(); + let dump = prompt_dump.clone(); tokio::spawn(async move { let mut buf = vec![0u8; 65536]; let n = match stream.read(&mut buf).await { @@ -973,14 +1086,24 @@ async fn http_mcp_loop( return; } "tools/list" => { - let tools: Vec = all_tool_descriptors_for_role(caller_role) - .iter() - .map(|d| json!({ - "name": d.name, - "description": d.description, - "inputSchema": d.input_schema, - })) - .collect(); + let descriptors = all_tool_descriptors_for_role(caller_role); + let tools: Vec = tool_descriptors_to_mcp_json(&descriptors); + if let Err(error) = dump_team_tools_list( + &dump, + TeamToolsListDump { + team_id: &tid, + caller_slot_id, + caller_role, + tools: &tools, + }, + ) { + warn!( + team_id = %tid, + caller_slot_id, + error = %error, + "team HTTP tools/list prompt dump failed" + ); + } json!({ "tools": tools }) } "tools/call" => { @@ -1061,34 +1184,35 @@ fn http_bearer_token(request: &str) -> Option<&str> { #[cfg(test)] mod tests { use super::*; - use aionui_api_types::{ - TeamRunTargetRole, TeamSendMessageDelivery, TeamSendMessageReason, TeamSendMessageTargetQueueState, - }; + use aionui_api_types::{TeamRunTargetRole, TeamSlotWorkPayload, TeamSlotWorkState}; #[test] fn build_send_message_queued_response_serializes_json_contract() { let response = build_send_message_queued_response(vec![AgentMessageQueueResult { - team_run_id: "run-1".into(), - delivery: TeamSendMessageDelivery::WakeRecorded, - target: TeamSendMessageTargetQueueState { + team_run_id: Some("run-1".into()), + target: TeamSlotWorkPayload { slot_id: "worker-1".into(), role: TeamRunTargetRole::Teammate, - queue_state: TeamSendMessageReason::QueuedForIdle, - pending_wake_count: 1, - starting_child_count: 0, + state: TeamSlotWorkState::Queued, + queued_foreground_count: 0, + queued_background_count: 1, active_turn_id: None, - suppressed_wake_count: 0, + active_turn_started_at_ms: None, + active_turn_elapsed_ms: None, + active_turn_slow: None, + active_turn_slow_threshold_ms: None, + blocked_reason: None, + team_run_id: Some("run-1".into()), }, }]) .unwrap(); let text = serde_json::to_string(&response).unwrap(); let payload: serde_json::Value = serde_json::from_str(&text).expect("team_send_message result must be JSON"); - assert_eq!(payload["status"], "queued"); - assert_eq!(payload["delivery"], "wake_recorded"); - assert_eq!(payload["reason"], "queued_for_idle"); - assert_eq!(payload["targets"][0]["slot_id"], "worker-1"); - assert_eq!(payload["targets"][0]["queue_state"], "queued_for_idle"); + assert_eq!(payload["team_run_id"], "run-1"); + assert_eq!(payload["target"]["slot_id"], "worker-1"); + assert_eq!(payload["target"]["state"], "queued"); + assert_eq!(payload["target"]["queued_background_count"], 1); } /// Non-Lead callers are rejected at the dispatch layer with the @@ -1130,8 +1254,7 @@ mod tests { let service: Weak = Weak::new(); let args = json!({ "name": "Helper", - "assistant_id": "word-creator", - "model": "claude-sonnet-4" + "assistant_id": "word-creator" }); let result = exec_spawn_agent(&args, &service, "team-1", "lead-1", TeammateRole::Lead).await; let err = result.expect_err("dead Weak must not succeed"); @@ -1141,6 +1264,22 @@ mod tests { ); } + #[tokio::test] + async fn exec_spawn_agent_rejects_model_override() { + let service: Weak = Weak::new(); + let args = json!({ + "name": "Helper", + "assistant_id": "word-creator", + "model": "claude-sonnet-4" + }); + let result = exec_spawn_agent(&args, &service, "team-1", "lead-1", TeammateRole::Lead).await; + let err = result.expect_err("model override must be rejected before service lookup"); + assert!( + err.message.contains("model is no longer accepted"), + "expected explicit model rejection, got {err:?}" + ); + } + #[tokio::test] async fn exec_spawn_agent_rejects_legacy_backend_alias() { let service: Weak = Weak::new(); @@ -1180,32 +1319,6 @@ mod tests { assert_eq!(err.domain_code, Some("TEAM_ASSISTANT_ID_REQUIRED")); } - #[tokio::test] - async fn exec_list_models_rejects_legacy_backend_alias() { - let service: Weak = Weak::new(); - let args = json!({ "backend": "claude" }); - let result = exec_list_models(&args, &service).await; - let err = result.expect_err("legacy backend alias must be rejected"); - assert!( - err.message.contains("backend is no longer accepted"), - "expected explicit backend rejection, got {err:?}" - ); - assert_eq!(err.domain_code, Some("TEAM_ASSISTANT_FIELD_UNSUPPORTED")); - } - - #[tokio::test] - async fn exec_list_models_rejects_legacy_agent_type_alias() { - let service: Weak = Weak::new(); - let args = json!({ "agent_type": "claude" }); - let result = exec_list_models(&args, &service).await; - let err = result.expect_err("legacy agent_type alias must be rejected"); - assert!( - err.message.contains("agent_type is no longer accepted"), - "expected explicit agent_type rejection, got {err:?}" - ); - assert_eq!(err.domain_code, Some("TEAM_ASSISTANT_FIELD_UNSUPPORTED")); - } - #[tokio::test] async fn exec_list_assistants_reports_service_unavailable_when_weak_dead() { let service: Weak = Weak::new(); diff --git a/crates/aionui-team/src/mcp/tools.rs b/crates/aionui-team/src/mcp/tools.rs index cbd0be288..440355795 100644 --- a/crates/aionui-team/src/mcp/tools.rs +++ b/crates/aionui-team/src/mcp/tools.rs @@ -1,13 +1,11 @@ -use aionui_db::models::AgentMetadataRow; use serde::{Deserialize, Serialize}; -use serde_json::{Value, json}; +use serde_json::Value; use crate::scheduler::SchedulerAction; use crate::types::TeammateRole; pub use aionui_team_prompts::tools::{ - TEAM_DESCRIBE_ASSISTANT_DESCRIPTION, TEAM_LIST_ASSISTANTS_DESCRIPTION, TEAM_LIST_MODELS_DESCRIPTION, - TEAM_SPAWN_AGENT_DESCRIPTION, + TEAM_DESCRIBE_ASSISTANT_DESCRIPTION, TEAM_LIST_ASSISTANTS_DESCRIPTION, TEAM_SPAWN_AGENT_DESCRIPTION, }; // --------------------------------------------------------------------------- @@ -52,18 +50,15 @@ pub struct SendMessageInput { /// Arguments for the `team_spawn_agent` MCP tool call. /// -/// Team spawning is assistant-first. The MCP tool only accepts -/// `assistant_id`, optional `model`, and optional `role`. +/// Team spawning is assistant-first. The MCP tool accepts an assistant identity; +/// model selection comes from the assistant configuration or UI model selector. #[derive(Debug, Default, Deserialize)] +#[serde(deny_unknown_fields)] pub struct SpawnAgentInput { pub name: String, #[serde(default)] - pub role: Option, - #[serde(default)] #[serde(alias = "assistantId")] pub assistant_id: Option, - #[serde(default)] - pub model: Option, } #[derive(Debug, Deserialize)] @@ -83,6 +78,26 @@ pub struct TaskUpdateInput { pub blocked_by: Option>, } +#[derive(Debug, Deserialize)] +#[serde(untagged)] +pub enum TaskListStatusInput { + Single(String), + Many(Vec), +} + +#[derive(Debug, Default, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct TaskListInput { + #[serde(default)] + pub owner: Option, + #[serde(default)] + pub status: Option, + #[serde(default)] + pub include_deleted: Option, + #[serde(default)] + pub limit: Option, +} + #[derive(Debug, Deserialize)] pub struct RenameAgentInput { pub slot_id: String, @@ -149,133 +164,11 @@ pub fn parse_tool_call( | "team_rename_agent" | "team_shutdown_agent" | "team_list_assistants" - | "team_list_models" | "team_describe_assistant" => Err("handled directly by server".into()), _ => Err(format!("Unknown tool: {tool_name}")), } } -// --------------------------------------------------------------------------- -// Phase-1 minimal handlers for `team_list_models` and `team_describe_assistant` -// --------------------------------------------------------------------------- - -/// Phase-1 minimal `team_list_models` handler. Returns a hard-coded -/// backend → models mapping. Used as fallback when DB is unavailable. -pub fn handle_team_list_models(_args: &Value) -> Value { - json!({ - "backends": [ - { - "backend": "claude", - "models": ["claude-sonnet-4", "claude-opus-4"] - }, - { - "backend": "codex", - "models": ["codex-mini-latest"] - } - ] - }) -} - -/// Build `team_list_models` response from DB rows. Reads each enabled, -/// team-capable backend's `available_models` column. Filters by -/// `backend` if provided. For internal agents (backend=NULL), -/// `provider_models` supplies the aggregated models from the providers table. -pub fn build_list_models_from_rows( - rows: &[AgentMetadataRow], - backend_filter: Option<&str>, - provider_models: &[String], -) -> Value { - use aionui_api_types::BehaviorPolicy; - use aionui_common::constants::is_team_capable; - - let mut backends: Vec = Vec::new(); - - for row in rows { - if !row.enabled { - continue; - } - // Use backend if present, otherwise agent_type as identifier (handles aionrs with backend=NULL) - let key = match row.backend.as_deref() { - Some(b) => b.to_owned(), - None => row.agent_type.clone(), - }; - let is_internal = row.backend.is_none(); - - // Check team capability: behavior_policy.supports_team OR legacy whitelist+MCP detection - let bp_supports = row - .behavior_policy - .as_deref() - .and_then(|s| serde_json::from_str::(s).ok()) - .is_some_and(|bp| bp.supports_team); - if !bp_supports { - let caps = row - .agent_capabilities - .as_deref() - .and_then(|s| serde_json::from_str::(s).ok()); - if !is_team_capable(&key, caps.as_ref()) { - continue; - } - } - - // Apply backend filter - if let Some(filter) = backend_filter - && key != filter - { - continue; - } - - // For internal agents (aionrs), use provider models - if is_internal && !provider_models.is_empty() { - backends.push(json!({ - "backend": key, - "models": provider_models, - })); - continue; - } - - // Parse available_models from DB. - // Format is either: - // {"current_model_id":"...", "available_models": [{"id":"...", "label":"..."}]} - // or legacy array: - // [{"id":"...", "name":"..."}] - let models: Vec = row - .available_models - .as_deref() - .and_then(|s| serde_json::from_str::(s).ok()) - .and_then(|v| { - // Try object with "available_models" key first (ModelInfoPayload format) - if let Some(arr) = v.get("available_models").and_then(Value::as_array) { - let ids: Vec = arr - .iter() - .filter_map(|e| e.get("id").and_then(Value::as_str).map(String::from)) - .collect(); - if !ids.is_empty() { - return Some(ids); - } - } - // Fallback: try parsing as direct array - if let Some(arr) = v.as_array() { - let ids: Vec = arr - .iter() - .filter_map(|e| e.get("id").and_then(Value::as_str).map(String::from)) - .collect(); - if !ids.is_empty() { - return Some(ids); - } - } - None - }) - .unwrap_or_default(); - - backends.push(json!({ - "backend": key, - "models": models, - })); - } - - json!({ "backends": backends }) -} - // --------------------------------------------------------------------------- // Tests // --------------------------------------------------------------------------- @@ -283,10 +176,11 @@ pub fn build_list_models_from_rows( #[cfg(test)] mod tests { use super::*; + use serde_json::json; #[test] fn all_descriptors_count() { - assert_eq!(all_tool_descriptors().len(), 11); + assert_eq!(all_tool_descriptors().len(), 10); } #[test] @@ -295,7 +189,7 @@ mod tests { let mut names: Vec<&str> = descs.iter().map(|d| d.name.as_str()).collect(); names.sort(); names.dedup(); - assert_eq!(names.len(), 11); + assert_eq!(names.len(), 10); } #[test] @@ -326,13 +220,16 @@ mod tests { } #[test] - fn team_spawn_agent_schema_exposes_model_and_assistant_id_only() { + fn team_spawn_agent_schema_exposes_assistant_id_without_model_override() { let desc = all_tool_descriptors() .into_iter() .find(|d| d.name == "team_spawn_agent") .unwrap(); let props = desc.input_schema["properties"].as_object().unwrap(); - assert!(props.contains_key("model"), "schema must expose 'model' field"); + assert!( + !props.contains_key("model"), + "team_spawn_agent must not expose model override; model comes from assistant configuration/UI" + ); assert!( props.contains_key("assistant_id"), "schema must expose 'assistant_id' field" @@ -341,6 +238,29 @@ mod tests { !props.contains_key("agent_type"), "assistant-first schema must not expose 'agent_type'" ); + assert!(!props.contains_key("role"), "spawn schema must not expose role"); + } + + #[test] + fn team_spawn_agent_rejects_role_argument() { + let result = serde_json::from_value::(json!({ + "name": "Helper", + "assistant_id": "word-creator", + "role": "teammate" + })); + + assert!(result.is_err(), "role must be denied as an unknown field"); + } + + #[test] + fn team_spawn_agent_rejects_model_argument() { + let result = serde_json::from_value::(json!({ + "name": "Helper", + "assistant_id": "word-creator", + "model": "claude-sonnet-4" + })); + + assert!(result.is_err(), "model must be denied as an unknown field"); } #[test] @@ -488,36 +408,6 @@ mod tests { // ---- D4 descriptor text remains aligned with assistant-first MCP contract ---- - #[test] - fn team_list_models_descriptor_text_matches() { - let desc = all_tool_descriptors() - .into_iter() - .find(|d| d.name == "team_list_models") - .expect("team_list_models descriptor missing"); - assert_eq!(desc.description, TEAM_LIST_MODELS_DESCRIPTION); - assert!( - desc.description - .starts_with("Query available models for assistant backends.") - ); - assert!( - desc.description.contains( - "Pass assistant_id to query models for a specific assistant, or omit it to see all backends." - ) - ); - } - - #[test] - fn team_list_models_schema_prefers_assistant_id() { - let desc = all_tool_descriptors() - .into_iter() - .find(|d| d.name == "team_list_models") - .unwrap(); - let props = desc.input_schema["properties"].as_object().unwrap(); - assert!(props.contains_key("assistant_id")); - assert!(!props.contains_key("agent_type")); - assert!(!props.contains_key("backend")); - } - #[test] fn team_describe_assistant_descriptor_text_matches() { let desc = all_tool_descriptors() @@ -595,6 +485,7 @@ mod tests { let props = desc.input_schema["properties"].as_object().unwrap(); let assistant_desc = props["assistant_id"]["description"].as_str().unwrap(); assert!(assistant_desc.starts_with("Assistant ID to spawn")); + assert!(!props.contains_key("model")); assert!(!props.contains_key("agent_type")); assert!(!props.contains_key("backend")); } @@ -605,10 +496,9 @@ mod tests { .into_iter() .find(|d| d.name == "team_spawn_agent") .unwrap(); - assert!( - desc.description - .contains("recommended assistant, and recommended model") - ); + assert!(desc.description.contains("recommended assistant")); + assert!(!desc.description.contains("recommended model")); + assert!(!desc.description.contains("model parameter")); assert!(!desc.description.contains("recommended assistant or backend")); } @@ -626,262 +516,4 @@ mod tests { assert!(assistant_desc.starts_with("The assistant ID from the available assistants catalog")); assert!(!assistant_desc.contains("preset assistant ID")); } - - // ---- D4 handlers return non-error payloads ---- - - #[test] - fn team_list_models_handler_returns_non_error() { - let value = handle_team_list_models(&json!({})); - let backends = value - .get("backends") - .and_then(|v| v.as_array()) - .expect("backends array missing"); - assert!(!backends.is_empty()); - let types: Vec<&str> = backends - .iter() - .filter_map(|e| e.get("backend").and_then(|v| v.as_str())) - .collect(); - assert!(types.contains(&"claude")); - assert!(types.contains(&"codex")); - } - - #[test] - fn build_list_models_from_rows_includes_enabled_team_capable_backends() { - let rows = vec![ - make_agent_row("claude", true, r#"[{"id":"claude-sonnet-4","name":"Sonnet 4"}]"#), - make_agent_row("codebuddy", true, r#"[{"id":"codebuddy-pro","name":"CodeBuddy Pro"}]"#), - make_agent_row("disabled-one", false, r#"[{"id":"m1","name":"M1"}]"#), - ]; - let value = build_list_models_from_rows(&rows, None, &[]); - let types: Vec<&str> = value["backends"] - .as_array() - .unwrap() - .iter() - .filter_map(|e| e["backend"].as_str()) - .collect(); - assert!(types.contains(&"claude")); - assert!(types.contains(&"codebuddy")); - assert!(!types.contains(&"disabled-one"), "disabled backends must be excluded"); - } - - #[test] - fn build_list_models_from_rows_uses_db_models_not_hardcoded() { - let rows = vec![make_agent_row( - "claude", - true, - r#"[{"id":"claude-opus-4","name":"Opus 4"},{"id":"claude-sonnet-4","name":"Sonnet 4"}]"#, - )]; - let value = build_list_models_from_rows(&rows, None, &[]); - let claude_entry = value["backends"] - .as_array() - .unwrap() - .iter() - .find(|e| e["backend"].as_str() == Some("claude")) - .expect("claude entry"); - let models: Vec<&str> = claude_entry["models"] - .as_array() - .unwrap() - .iter() - .filter_map(|v| v.as_str()) - .collect(); - assert_eq!(models, vec!["claude-opus-4", "claude-sonnet-4"]); - } - - #[test] - fn build_list_models_from_rows_filters_by_backend() { - let rows = vec![ - make_agent_row("claude", true, r#"[{"id":"claude-sonnet-4","name":"Sonnet 4"}]"#), - make_agent_row("codebuddy", true, r#"[{"id":"cb-pro","name":"Pro"}]"#), - ]; - let value = build_list_models_from_rows(&rows, Some("codebuddy"), &[]); - let types: Vec<&str> = value["backends"] - .as_array() - .unwrap() - .iter() - .filter_map(|e| e["backend"].as_str()) - .collect(); - assert_eq!(types, vec!["codebuddy"]); - } - - #[test] - fn build_list_models_from_rows_skips_null_available_models() { - let rows = vec![ - make_agent_row("claude", true, r#"[{"id":"claude-sonnet-4","name":"Sonnet 4"}]"#), - make_agent_row_no_models("gemini", true), - ]; - let value = build_list_models_from_rows(&rows, None, &[]); - let types: Vec<&str> = value["backends"] - .as_array() - .unwrap() - .iter() - .filter_map(|e| e["backend"].as_str()) - .collect(); - // gemini has no available_models in DB → should still appear but with empty models - assert!(types.contains(&"gemini")); - } - - fn make_agent_row(backend: &str, enabled: bool, available_models: &str) -> AgentMetadataRow { - AgentMetadataRow { - id: format!("id-{backend}"), - icon: None, - name: capitalize_first(backend), - name_i18n: None, - description: None, - description_i18n: None, - backend: Some(backend.to_owned()), - agent_type: "acp".to_owned(), - agent_source: "builtin".to_owned(), - agent_source_info: None, - enabled, - command: None, - args: None, - env: None, - native_skills_dirs: None, - behavior_policy: None, - yolo_id: None, - agent_capabilities: Some(r#"{"mcp":true}"#.to_owned()), - auth_methods: None, - config_options: None, - available_modes: None, - available_models: Some(available_models.to_owned()), - available_commands: None, - sort_order: 0, - last_check_status: None, - last_check_kind: None, - last_check_error_code: None, - last_check_error_message: None, - last_check_guidance: None, - last_check_latency_ms: None, - last_check_at: None, - last_success_at: None, - last_failure_at: None, - command_override: None, - env_override: None, - created_at: 0, - updated_at: 0, - } - } - - fn make_agent_row_no_models(backend: &str, enabled: bool) -> AgentMetadataRow { - let mut row = make_agent_row(backend, enabled, "[]"); - row.available_models = None; - row - } - - fn capitalize_first(s: &str) -> String { - let mut c = s.chars(); - match c.next() { - None => String::new(), - Some(f) => f.to_uppercase().collect::() + c.as_str(), - } - } - - #[test] - fn build_list_models_from_rows_includes_null_backend_with_supports_team() { - let mut aionrs_row = make_agent_row("aionrs", true, r#"[{"id":"aionrs-default","name":"AionRS"}]"#); - aionrs_row.backend = None; - aionrs_row.agent_type = "aionrs".to_owned(); - aionrs_row.agent_source = "internal".to_owned(); - aionrs_row.agent_capabilities = None; - aionrs_row.behavior_policy = Some(r#"{"supports_team":true}"#.to_owned()); - - let rows = vec![ - make_agent_row("claude", true, r#"[{"id":"claude-sonnet-4","name":"Sonnet 4"}]"#), - aionrs_row, - ]; - let value = build_list_models_from_rows(&rows, None, &[]); - let types: Vec<&str> = value["backends"] - .as_array() - .unwrap() - .iter() - .filter_map(|e| e["backend"].as_str()) - .collect(); - assert!(types.contains(&"claude")); - assert!( - types.contains(&"aionrs"), - "aionrs with backend=NULL but supports_team=true must be included" - ); - } - - #[test] - fn build_list_models_from_rows_filters_null_backend_by_backend() { - let mut aionrs_row = make_agent_row("aionrs", true, r#"[{"id":"aionrs-default","name":"AionRS"}]"#); - aionrs_row.backend = None; - aionrs_row.agent_type = "aionrs".to_owned(); - aionrs_row.agent_capabilities = None; - aionrs_row.behavior_policy = Some(r#"{"supports_team":true}"#.to_owned()); - - let rows = vec![ - make_agent_row("claude", true, r#"[{"id":"claude-sonnet-4","name":"Sonnet 4"}]"#), - aionrs_row, - ]; - // Filter by "aionrs" should only return aionrs - let value = build_list_models_from_rows(&rows, Some("aionrs"), &[]); - let types: Vec<&str> = value["backends"] - .as_array() - .unwrap() - .iter() - .filter_map(|e| e["backend"].as_str()) - .collect(); - assert_eq!(types, vec!["aionrs"]); - } - - #[test] - fn build_list_models_from_rows_parses_model_info_payload_format() { - let model_info_json = r#"{"current_model_id":"DeepSeek-V3.2","current_model_label":"DeepSeek-V3.2","available_models":[{"id":"GLM-5.0","label":"GLM-5.0"},{"id":"GLM-5.0-Turbo","label":"GLM-5.0-Turbo"},{"id":"DeepSeek-V3.2","label":"DeepSeek-V3.2"}]}"#; - let rows = vec![make_agent_row("codebuddy", true, model_info_json)]; - let value = build_list_models_from_rows(&rows, None, &[]); - let cb_entry = value["backends"] - .as_array() - .unwrap() - .iter() - .find(|e| e["backend"].as_str() == Some("codebuddy")) - .expect("codebuddy entry"); - let models: Vec<&str> = cb_entry["models"] - .as_array() - .unwrap() - .iter() - .filter_map(|v| v.as_str()) - .collect(); - assert_eq!(models, vec!["GLM-5.0", "GLM-5.0-Turbo", "DeepSeek-V3.2"]); - } - - #[test] - fn build_list_models_from_rows_uses_provider_models_for_internal_agents() { - let mut aionrs_row = make_agent_row("aionrs", true, "[]"); - aionrs_row.backend = None; - aionrs_row.agent_type = "aionrs".to_owned(); - aionrs_row.agent_source = "internal".to_owned(); - aionrs_row.agent_capabilities = None; - aionrs_row.available_models = None; - aionrs_row.behavior_policy = Some(r#"{"supports_team":true}"#.to_owned()); - - let provider_models = vec![ - "gemini-3.1-pro-preview".to_owned(), - "gpt-5.4".to_owned(), - "gpt-5.2".to_owned(), - ]; - let rows = vec![ - make_agent_row( - "claude", - true, - r#"{"available_models":[{"id":"claude-sonnet-4","label":"Sonnet 4"}]}"#, - ), - aionrs_row, - ]; - let value = build_list_models_from_rows(&rows, None, &provider_models); - let aionrs_entry = value["backends"] - .as_array() - .unwrap() - .iter() - .find(|e| e["backend"].as_str() == Some("aionrs")) - .expect("aionrs entry"); - let models: Vec<&str> = aionrs_entry["models"] - .as_array() - .unwrap() - .iter() - .filter_map(|v| v.as_str()) - .collect(); - assert_eq!(models, vec!["gemini-3.1-pro-preview", "gpt-5.4", "gpt-5.2"]); - } } diff --git a/crates/aionui-team/src/member_runtime.rs b/crates/aionui-team/src/member_runtime.rs new file mode 100644 index 000000000..2f42b99f4 --- /dev/null +++ b/crates/aionui-team/src/member_runtime.rs @@ -0,0 +1,893 @@ +use std::collections::HashMap; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::{Mutex, MutexGuard}; + +use tokio::sync::watch; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct MemberRuntimeFailure { + pub(crate) classification: &'static str, + pub(crate) public_reason: String, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum AttachOutcome { + Ready, + Failed(MemberRuntimeFailure), + Removed, + SessionStopped, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +enum AttachSignal { + Pending, + Terminal(AttachOutcome), +} + +#[derive(Debug, Clone)] +pub(crate) struct AttachWaiter { + operation_id: u64, + outcome_rx: watch::Receiver, +} + +impl AttachWaiter { + pub(crate) fn operation_id(&self) -> u64 { + self.operation_id + } + + pub(crate) async fn wait(mut self) -> AttachOutcome { + loop { + let terminal = match &*self.outcome_rx.borrow() { + AttachSignal::Pending => None, + AttachSignal::Terminal(outcome) => Some(outcome.clone()), + }; + if let Some(outcome) = terminal { + return outcome; + } + if self.outcome_rx.changed().await.is_err() { + return AttachOutcome::SessionStopped; + } + } + } +} + +#[derive(Debug)] +pub(crate) struct AttachLease { + session_generation: String, + agent_id: String, + waiter: AttachWaiter, +} + +impl AttachLease { + pub(crate) fn operation_id(&self) -> u64 { + self.waiter.operation_id() + } + + pub(crate) fn waiter(&self) -> AttachWaiter { + self.waiter.clone() + } +} + +#[derive(Debug)] +pub(crate) struct RemoveLease { + session_generation: String, + agent_id: String, + operation_id: u64, +} + +impl RemoveLease { + pub(crate) fn operation_id(&self) -> u64 { + self.operation_id + } +} + +#[derive(Debug)] +pub(crate) enum ReserveAttach { + Start(AttachLease), + Join(AttachWaiter), + AlreadyReady, + Removing(AttachWaiter), + SessionStopped, +} + +#[derive(Debug)] +pub(crate) enum BeginRemove { + Start(RemoveLease), + Join(AttachWaiter), + Absent, + SessionStopped, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum MemberRuntimeSnapshot { + Absent, + Attaching { + operation_id: u64, + }, + Ready, + Failed { + operation_id: u64, + failure: MemberRuntimeFailure, + }, + Removing { + operation_id: u64, + }, + SessionStopped, +} + +#[derive(Debug)] +enum MemberRuntimeEntry { + Attaching { + operation_id: u64, + outcome_tx: watch::Sender, + }, + Ready, + Failed { + operation_id: u64, + failure: MemberRuntimeFailure, + outcome_tx: watch::Sender, + }, + Removing { + operation_id: u64, + outcome_tx: watch::Sender, + }, +} + +#[derive(Debug)] +pub(crate) struct MemberRuntimeRegistry { + session_generation: String, + next_operation_id: AtomicU64, + entries: Mutex>, + stopped: AtomicBool, +} + +impl MemberRuntimeRegistry { + pub(crate) fn new(session_generation: impl ToString) -> Self { + Self { + session_generation: session_generation.to_string(), + next_operation_id: AtomicU64::new(1), + entries: Mutex::new(HashMap::new()), + stopped: AtomicBool::new(false), + } + } + + pub(crate) fn generation(&self) -> &str { + &self.session_generation + } + + pub(crate) fn seed_ready(&self, agent_id: impl Into) -> bool { + let agent_id = agent_id.into(); + let mut entries = self.lock_entries(); + if self.stopped.load(Ordering::Acquire) { + return false; + } + match entries.get(&agent_id) { + None => { + entries.insert(agent_id, MemberRuntimeEntry::Ready); + true + } + Some(MemberRuntimeEntry::Ready) => true, + Some( + MemberRuntimeEntry::Attaching { .. } + | MemberRuntimeEntry::Failed { .. } + | MemberRuntimeEntry::Removing { .. }, + ) => false, + } + } + + pub(crate) fn reserve_attach(&self, agent_id: &str, retry_failed: bool) -> ReserveAttach { + let mut entries = self.lock_entries(); + if self.stopped.load(Ordering::Acquire) { + return ReserveAttach::SessionStopped; + } + + match entries.get(agent_id) { + Some(MemberRuntimeEntry::Attaching { + operation_id, + outcome_tx, + }) => ReserveAttach::Join(waiter(*operation_id, outcome_tx.subscribe())), + Some(MemberRuntimeEntry::Ready) => ReserveAttach::AlreadyReady, + Some(MemberRuntimeEntry::Failed { + operation_id, + outcome_tx, + .. + }) if !retry_failed => ReserveAttach::Join(waiter(*operation_id, outcome_tx.subscribe())), + Some(MemberRuntimeEntry::Removing { + operation_id, + outcome_tx, + }) => ReserveAttach::Removing(waiter(*operation_id, outcome_tx.subscribe())), + Some(MemberRuntimeEntry::Failed { .. }) | None => { + let operation_id = self.next_operation_id(); + let (outcome_tx, outcome_rx) = watch::channel(AttachSignal::Pending); + entries.insert( + agent_id.to_owned(), + MemberRuntimeEntry::Attaching { + operation_id, + outcome_tx, + }, + ); + ReserveAttach::Start(AttachLease { + session_generation: self.session_generation.clone(), + agent_id: agent_id.to_owned(), + waiter: waiter(operation_id, outcome_rx), + }) + } + } + } + + /// Atomically converts a previously-ready slot into a repair attach. + /// + /// This is only for reconciliation after the task manager confirms the + /// underlying runtime disappeared. Ordinary attaches must use + /// [`Self::reserve_attach`]. + pub(crate) fn reserve_repair(&self, agent_id: &str) -> ReserveAttach { + let mut entries = self.lock_entries(); + if self.stopped.load(Ordering::Acquire) { + return ReserveAttach::SessionStopped; + } + + match entries.get(agent_id) { + Some(MemberRuntimeEntry::Attaching { + operation_id, + outcome_tx, + }) => ReserveAttach::Join(waiter(*operation_id, outcome_tx.subscribe())), + Some(MemberRuntimeEntry::Ready) => { + let operation_id = self.next_operation_id(); + let (outcome_tx, outcome_rx) = watch::channel(AttachSignal::Pending); + entries.insert( + agent_id.to_owned(), + MemberRuntimeEntry::Attaching { + operation_id, + outcome_tx, + }, + ); + ReserveAttach::Start(AttachLease { + session_generation: self.session_generation.clone(), + agent_id: agent_id.to_owned(), + waiter: waiter(operation_id, outcome_rx), + }) + } + Some(MemberRuntimeEntry::Failed { + operation_id, + outcome_tx, + .. + }) => ReserveAttach::Join(waiter(*operation_id, outcome_tx.subscribe())), + Some(MemberRuntimeEntry::Removing { + operation_id, + outcome_tx, + }) => ReserveAttach::Removing(waiter(*operation_id, outcome_tx.subscribe())), + None => { + let operation_id = self.next_operation_id(); + let (outcome_tx, outcome_rx) = watch::channel(AttachSignal::Pending); + entries.insert( + agent_id.to_owned(), + MemberRuntimeEntry::Attaching { + operation_id, + outcome_tx, + }, + ); + ReserveAttach::Start(AttachLease { + session_generation: self.session_generation.clone(), + agent_id: agent_id.to_owned(), + waiter: waiter(operation_id, outcome_rx), + }) + } + } + } + + pub(crate) fn owns_attach(&self, lease: &AttachLease) -> bool { + if lease.session_generation != self.session_generation || self.stopped.load(Ordering::Acquire) { + return false; + } + matches!( + self.lock_entries().get(&lease.agent_id), + Some(MemberRuntimeEntry::Attaching { operation_id, .. }) if *operation_id == lease.operation_id() + ) + } + + pub(crate) fn commit_ready(&self, lease: &AttachLease) -> bool { + if lease.session_generation != self.session_generation { + return false; + } + let mut entries = self.lock_entries(); + if self.stopped.load(Ordering::Acquire) { + return false; + } + let Some(MemberRuntimeEntry::Attaching { + operation_id: current_operation_id, + outcome_tx, + }) = entries.get(&lease.agent_id) + else { + return false; + }; + if *current_operation_id != lease.operation_id() { + return false; + } + let outcome_tx = outcome_tx.clone(); + entries.insert(lease.agent_id.clone(), MemberRuntimeEntry::Ready); + outcome_tx.send_replace(AttachSignal::Terminal(AttachOutcome::Ready)); + true + } + + pub(crate) fn commit_failed(&self, lease: &AttachLease, failure: MemberRuntimeFailure) -> bool { + if lease.session_generation != self.session_generation { + return false; + } + let mut entries = self.lock_entries(); + if self.stopped.load(Ordering::Acquire) { + return false; + } + let Some(MemberRuntimeEntry::Attaching { + operation_id: current_operation_id, + outcome_tx, + }) = entries.get(&lease.agent_id) + else { + return false; + }; + if *current_operation_id != lease.operation_id() { + return false; + } + let outcome_tx = outcome_tx.clone(); + entries.insert( + lease.agent_id.clone(), + MemberRuntimeEntry::Failed { + operation_id: lease.operation_id(), + failure: failure.clone(), + outcome_tx: outcome_tx.clone(), + }, + ); + outcome_tx.send_replace(AttachSignal::Terminal(AttachOutcome::Failed(failure))); + true + } + + pub(crate) fn begin_remove(&self, agent_id: &str) -> BeginRemove { + let mut entries = self.lock_entries(); + if self.stopped.load(Ordering::Acquire) { + return BeginRemove::SessionStopped; + } + if let Some(MemberRuntimeEntry::Removing { + operation_id, + outcome_tx, + }) = entries.get(agent_id) + { + return BeginRemove::Join(waiter(*operation_id, outcome_tx.subscribe())); + } + + let previous = entries.remove(agent_id); + let cancelled_outcome_tx = match previous { + Some(MemberRuntimeEntry::Attaching { + operation_id: _, + outcome_tx, + }) => Some(outcome_tx), + Some(MemberRuntimeEntry::Failed { .. }) => None, + Some(MemberRuntimeEntry::Ready) => None, + Some(MemberRuntimeEntry::Removing { .. }) => unreachable!("handled above"), + None => return BeginRemove::Absent, + }; + + let operation_id = self.next_operation_id(); + let (outcome_tx, _outcome_rx) = watch::channel(AttachSignal::Pending); + entries.insert( + agent_id.to_owned(), + MemberRuntimeEntry::Removing { + operation_id, + outcome_tx, + }, + ); + if let Some(cancelled_outcome_tx) = cancelled_outcome_tx { + cancelled_outcome_tx.send_replace(AttachSignal::Terminal(AttachOutcome::Removed)); + } + BeginRemove::Start(RemoveLease { + session_generation: self.session_generation.clone(), + agent_id: agent_id.to_owned(), + operation_id, + }) + } + + pub(crate) fn finish_remove(&self, lease: &RemoveLease) -> bool { + if lease.session_generation != self.session_generation { + return false; + } + let mut entries = self.lock_entries(); + if self.stopped.load(Ordering::Acquire) { + return false; + } + let Some(MemberRuntimeEntry::Removing { + operation_id: current_operation_id, + outcome_tx, + }) = entries.get(&lease.agent_id) + else { + return false; + }; + if *current_operation_id != lease.operation_id() { + return false; + } + let outcome_tx = outcome_tx.clone(); + entries.remove(&lease.agent_id); + outcome_tx.send_replace(AttachSignal::Terminal(AttachOutcome::Removed)); + true + } + + /// Marks a removed runtime as requiring a fresh attach after persistence fails. + /// + /// Call this only after runtime cancellation and join cleanup have completed, + /// so no live runtime remains. A later reconcile with `retry_failed = true` + /// starts a new attach operation. + pub(crate) fn restore_attach_required_after_remove_persist_error( + &self, + lease: &RemoveLease, + failure: MemberRuntimeFailure, + ) -> bool { + if lease.session_generation != self.session_generation { + return false; + } + let mut entries = self.lock_entries(); + if self.stopped.load(Ordering::Acquire) { + return false; + } + let Some(MemberRuntimeEntry::Removing { + operation_id: current_operation_id, + outcome_tx, + }) = entries.get(&lease.agent_id) + else { + return false; + }; + if *current_operation_id != lease.operation_id() { + return false; + } + let outcome_tx = outcome_tx.clone(); + entries.insert( + lease.agent_id.clone(), + MemberRuntimeEntry::Failed { + operation_id: lease.operation_id(), + failure: failure.clone(), + outcome_tx: outcome_tx.clone(), + }, + ); + outcome_tx.send_replace(AttachSignal::Terminal(AttachOutcome::Failed(failure))); + true + } + + pub(crate) fn snapshot(&self, agent_id: &str) -> MemberRuntimeSnapshot { + let entries = self.lock_entries(); + if self.stopped.load(Ordering::Acquire) { + return MemberRuntimeSnapshot::SessionStopped; + } + match entries.get(agent_id) { + None => MemberRuntimeSnapshot::Absent, + Some(MemberRuntimeEntry::Attaching { operation_id, .. }) => MemberRuntimeSnapshot::Attaching { + operation_id: *operation_id, + }, + Some(MemberRuntimeEntry::Ready) => MemberRuntimeSnapshot::Ready, + Some(MemberRuntimeEntry::Failed { + operation_id, failure, .. + }) => MemberRuntimeSnapshot::Failed { + operation_id: *operation_id, + failure: failure.clone(), + }, + Some(MemberRuntimeEntry::Removing { operation_id, .. }) => MemberRuntimeSnapshot::Removing { + operation_id: *operation_id, + }, + } + } + + pub(crate) fn stop(&self) -> bool { + let mut entries = self.lock_entries(); + if self.stopped.swap(true, Ordering::AcqRel) { + return false; + } + for entry in entries.drain().map(|(_, entry)| entry) { + match entry { + MemberRuntimeEntry::Attaching { outcome_tx, .. } | MemberRuntimeEntry::Removing { outcome_tx, .. } => { + outcome_tx.send_replace(AttachSignal::Terminal(AttachOutcome::SessionStopped)); + } + MemberRuntimeEntry::Ready | MemberRuntimeEntry::Failed { .. } => {} + } + } + true + } + + fn next_operation_id(&self) -> u64 { + self.next_operation_id.fetch_add(1, Ordering::Relaxed) + } + + /// Critical sections contain only short in-memory transitions with no await, + /// I/O, or user callbacks, so retaining the map after poisoning is safe. + fn lock_entries(&self) -> MutexGuard<'_, HashMap> { + self.entries.lock().unwrap_or_else(|poisoned| poisoned.into_inner()) + } +} + +fn waiter(operation_id: u64, outcome_rx: watch::Receiver) -> AttachWaiter { + AttachWaiter { + operation_id, + outcome_rx, + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use tokio::sync::{mpsc, oneshot}; + + use super::*; + + fn failure(classification: &'static str, public_reason: &str) -> MemberRuntimeFailure { + MemberRuntimeFailure { + classification, + public_reason: public_reason.to_owned(), + } + } + + #[tokio::test] + async fn concurrent_reservations_join_the_same_attach_operation() { + let registry = Arc::new(MemberRuntimeRegistry::new(41)); + let (ready_tx, mut ready_rx) = mpsc::channel(2); + let (release_first_tx, release_first_rx) = oneshot::channel(); + let (release_second_tx, release_second_rx) = oneshot::channel(); + + let first_registry = Arc::clone(®istry); + let first_ready_tx = ready_tx.clone(); + let first = tokio::spawn(async move { + first_ready_tx.send(()).await.unwrap(); + release_first_rx.await.unwrap(); + first_registry.reserve_attach("worker-1", false) + }); + let second_registry = Arc::clone(®istry); + let second = tokio::spawn(async move { + ready_tx.send(()).await.unwrap(); + release_second_rx.await.unwrap(); + second_registry.reserve_attach("worker-1", false) + }); + + ready_rx.recv().await.unwrap(); + ready_rx.recv().await.unwrap(); + release_first_tx.send(()).unwrap(); + release_second_tx.send(()).unwrap(); + + let first = first.await.unwrap(); + let second = second.await.unwrap(); + let (lease, waiter) = match (first, second) { + (ReserveAttach::Start(lease), ReserveAttach::Join(waiter)) + | (ReserveAttach::Join(waiter), ReserveAttach::Start(lease)) => (lease, waiter), + _ => panic!("expected one starter and one joiner"), + }; + + assert_eq!(lease.operation_id(), 1); + assert_eq!(waiter.operation_id(), 1); + assert!(registry.commit_ready(&lease)); + assert_eq!(lease.waiter().wait().await, AttachOutcome::Ready); + assert_eq!(waiter.wait().await, AttachOutcome::Ready); + assert_eq!(registry.snapshot("worker-1"), MemberRuntimeSnapshot::Ready); + } + + #[tokio::test] + async fn failed_operation_is_retried_only_by_a_later_reconcile() { + let registry = MemberRuntimeRegistry::new(42); + assert_eq!(registry.generation(), "42"); + assert!(registry.seed_ready("lead-1")); + assert!(matches!( + registry.reserve_attach("lead-1", false), + ReserveAttach::AlreadyReady + )); + + let first = match registry.reserve_attach("worker-1", false) { + ReserveAttach::Start(lease) => lease, + _ => panic!("first reservation must start"), + }; + assert_eq!(first.operation_id(), 1); + let first_failure = failure("provider_unavailable", "Agent is temporarily unavailable"); + assert!(registry.commit_failed(&first, first_failure.clone())); + assert_eq!( + first.waiter().wait().await, + AttachOutcome::Failed(first_failure.clone()) + ); + + let failed_waiter = match registry.reserve_attach("worker-1", false) { + ReserveAttach::Join(waiter) => waiter, + _ => panic!("failed operation must not retry without reconciliation"), + }; + assert_eq!(failed_waiter.operation_id(), 1); + assert_eq!(failed_waiter.wait().await, AttachOutcome::Failed(first_failure)); + + let retry = match registry.reserve_attach("worker-1", true) { + ReserveAttach::Start(lease) => lease, + _ => panic!("later reconciliation must start a retry"), + }; + assert_eq!(retry.operation_id(), 2); + assert_eq!( + registry.snapshot("worker-1"), + MemberRuntimeSnapshot::Attaching { operation_id: 2 } + ); + assert!(registry.commit_ready(&retry)); + assert_eq!(retry.waiter().wait().await, AttachOutcome::Ready); + } + + #[tokio::test] + async fn removing_cancels_attach_and_rejects_late_ready_commit() { + let registry = MemberRuntimeRegistry::new(43); + let attach = match registry.reserve_attach("worker-1", false) { + ReserveAttach::Start(lease) => lease, + _ => panic!("first reservation must start"), + }; + assert_eq!(attach.operation_id(), 1); + + let remove = match registry.begin_remove("worker-1") { + BeginRemove::Start(lease) => lease, + _ => panic!("remove must start"), + }; + assert_eq!(remove.operation_id(), 2); + assert_eq!( + registry.snapshot("worker-1"), + MemberRuntimeSnapshot::Removing { operation_id: 2 } + ); + assert_eq!(attach.waiter().wait().await, AttachOutcome::Removed); + assert!(!registry.commit_ready(&attach)); + assert!(!registry.commit_failed(&attach, failure("late_provider_error", "Agent attach did not complete"))); + + let removing_waiter = match registry.reserve_attach("worker-1", false) { + ReserveAttach::Removing(waiter) => waiter, + _ => panic!("reservation must wait while removal is active"), + }; + assert_eq!(removing_waiter.operation_id(), 2); + assert!(registry.finish_remove(&remove)); + assert_eq!(removing_waiter.wait().await, AttachOutcome::Removed); + assert_eq!(registry.snapshot("worker-1"), MemberRuntimeSnapshot::Absent); + + assert!(registry.seed_ready("worker-2")); + let second_remove = match registry.begin_remove("worker-2") { + BeginRemove::Start(lease) => lease, + _ => panic!("seeded runtime removal must start"), + }; + assert_eq!(second_remove.operation_id(), 3); + let second_remove_waiter = match registry.reserve_attach("worker-2", false) { + ReserveAttach::Removing(waiter) => waiter, + _ => panic!("attach must wait for the active removal"), + }; + let remove_failure = failure("remove_rejected", "Agent removal could not be completed"); + assert!(registry.restore_attach_required_after_remove_persist_error(&second_remove, remove_failure.clone())); + assert_eq!( + second_remove_waiter.wait().await, + AttachOutcome::Failed(remove_failure.clone()) + ); + assert_eq!( + registry.snapshot("worker-2"), + MemberRuntimeSnapshot::Failed { + operation_id: 3, + failure: remove_failure.clone(), + } + ); + + let retry = match registry.reserve_attach("worker-2", true) { + ReserveAttach::Start(lease) => lease, + _ => panic!("persistence failure must require a fresh attach"), + }; + assert_eq!(retry.operation_id(), 4); + assert!(registry.commit_ready(&retry)); + } + + #[tokio::test] + async fn session_stop_cancels_all_attaches_and_unblocks_waiters() { + let registry = Arc::new(MemberRuntimeRegistry::new(44)); + let first = match registry.reserve_attach("worker-1", false) { + ReserveAttach::Start(lease) => lease, + _ => panic!("first reservation must start"), + }; + let second = match registry.reserve_attach("worker-2", false) { + ReserveAttach::Start(lease) => lease, + _ => panic!("second reservation must start"), + }; + assert_eq!(first.operation_id(), 1); + assert_eq!(second.operation_id(), 2); + + let (waiting_tx, mut waiting_rx) = mpsc::channel(2); + let first_wait = first.waiter(); + let first_waiting_tx = waiting_tx.clone(); + let first_task = tokio::spawn(async move { + first_waiting_tx.send(()).await.unwrap(); + first_wait.wait().await + }); + let second_wait = second.waiter(); + let second_task = tokio::spawn(async move { + waiting_tx.send(()).await.unwrap(); + second_wait.wait().await + }); + waiting_rx.recv().await.unwrap(); + waiting_rx.recv().await.unwrap(); + + assert!(registry.stop()); + assert!(!registry.stop()); + assert_eq!(first_task.await.unwrap(), AttachOutcome::SessionStopped); + assert_eq!(second_task.await.unwrap(), AttachOutcome::SessionStopped); + assert!(!registry.commit_ready(&first)); + assert!(!registry.commit_failed(&second, failure("provider_stopped", "Agent session stopped"))); + assert_eq!(registry.snapshot("worker-1"), MemberRuntimeSnapshot::SessionStopped); + + assert!(matches!( + registry.reserve_attach("worker-3", false), + ReserveAttach::SessionStopped + )); + } + + #[tokio::test] + async fn stale_operation_id_cannot_commit_over_a_newer_operation() { + let registry = MemberRuntimeRegistry::new(45); + let first = match registry.reserve_attach("worker-1", false) { + ReserveAttach::Start(lease) => lease, + _ => panic!("first reservation must start"), + }; + let first_failure = failure("transport", "Agent connection failed"); + assert!(registry.commit_failed(&first, first_failure.clone())); + assert_eq!(first.waiter().wait().await, AttachOutcome::Failed(first_failure)); + + let second = match registry.reserve_attach("worker-1", true) { + ReserveAttach::Start(lease) => lease, + _ => panic!("retry reservation must start"), + }; + assert_eq!(second.operation_id(), 2); + assert!(!registry.commit_ready(&first)); + assert!(!registry.commit_failed(&first, failure("stale", "Stale attach failed"))); + assert_eq!( + registry.snapshot("worker-1"), + MemberRuntimeSnapshot::Attaching { operation_id: 2 } + ); + assert!(registry.commit_ready(&second)); + assert_eq!(second.waiter().wait().await, AttachOutcome::Ready); + assert_eq!(registry.snapshot("worker-1"), MemberRuntimeSnapshot::Ready); + } + + #[tokio::test] + async fn failed_attach_outcome_is_not_rewritten_when_remove_begins() { + let registry = MemberRuntimeRegistry::new(46); + let attach = match registry.reserve_attach("worker-1", false) { + ReserveAttach::Start(lease) => lease, + _ => panic!("first reservation must start"), + }; + let attach_failure = failure("transport", "Agent connection failed"); + assert!(registry.commit_failed(&attach, attach_failure.clone())); + + let remove = match registry.begin_remove("worker-1") { + BeginRemove::Start(lease) => lease, + _ => panic!("failed runtime removal must start"), + }; + + assert_eq!(remove.operation_id(), 2); + assert_eq!(attach.waiter().wait().await, AttachOutcome::Failed(attach_failure)); + } + + #[tokio::test] + async fn failed_attach_outcome_is_not_rewritten_when_session_stops() { + let registry = MemberRuntimeRegistry::new(47); + let attach = match registry.reserve_attach("worker-1", false) { + ReserveAttach::Start(lease) => lease, + _ => panic!("first reservation must start"), + }; + let attach_failure = failure("transport", "Agent connection failed"); + assert!(registry.commit_failed(&attach, attach_failure.clone())); + + assert!(registry.stop()); + + assert_eq!(attach.waiter().wait().await, AttachOutcome::Failed(attach_failure)); + } + + #[tokio::test] + async fn concurrent_remove_requests_join_the_same_remove_operation() { + let registry = MemberRuntimeRegistry::new(48); + assert!(registry.seed_ready("worker-1")); + let owner = match registry.begin_remove("worker-1") { + BeginRemove::Start(lease) => lease, + _ => panic!("first remove must start"), + }; + let joiner = match registry.begin_remove("worker-1") { + BeginRemove::Join(waiter) => waiter, + _ => panic!("second remove must join"), + }; + + assert_eq!(owner.operation_id(), 1); + assert_eq!(joiner.operation_id(), 1); + assert!(registry.finish_remove(&owner)); + assert_eq!(joiner.wait().await, AttachOutcome::Removed); + } + + #[tokio::test] + async fn session_stop_unblocks_waiters_for_an_active_remove() { + let registry = MemberRuntimeRegistry::new(49); + assert!(registry.seed_ready("worker-1")); + let owner = match registry.begin_remove("worker-1") { + BeginRemove::Start(lease) => lease, + _ => panic!("remove must start"), + }; + let reserve_waiter = match registry.reserve_attach("worker-1", false) { + ReserveAttach::Removing(waiter) => waiter, + _ => panic!("attach must wait for active remove"), + }; + + assert!(registry.stop()); + + assert_eq!(reserve_waiter.wait().await, AttachOutcome::SessionStopped); + assert!(!registry.finish_remove(&owner)); + } + + #[test] + fn stale_remove_owner_cannot_finish_or_restore_a_newer_remove() { + let registry = MemberRuntimeRegistry::new(50); + assert!(registry.seed_ready("worker-1")); + let first = match registry.begin_remove("worker-1") { + BeginRemove::Start(lease) => lease, + _ => panic!("first remove must start"), + }; + assert!(registry.finish_remove(&first)); + + assert!(registry.seed_ready("worker-1")); + let second = match registry.begin_remove("worker-1") { + BeginRemove::Start(lease) => lease, + _ => panic!("second remove must start"), + }; + + assert_eq!(first.operation_id(), 1); + assert_eq!(second.operation_id(), 2); + assert!(!registry.finish_remove(&first)); + assert!( + !registry + .restore_attach_required_after_remove_persist_error(&first, failure("stale", "Stale remove failed")) + ); + assert_eq!( + registry.snapshot("worker-1"), + MemberRuntimeSnapshot::Removing { operation_id: 2 } + ); + assert!(registry.finish_remove(&second)); + } + + #[test] + fn attach_owner_from_another_session_generation_cannot_commit() { + let old_registry = MemberRuntimeRegistry::new(51); + let old_owner = match old_registry.reserve_attach("worker-1", false) { + ReserveAttach::Start(lease) => lease, + _ => panic!("old session attach must start"), + }; + let current_registry = MemberRuntimeRegistry::new(52); + let current_owner = match current_registry.reserve_attach("worker-1", false) { + ReserveAttach::Start(lease) => lease, + _ => panic!("current session attach must start"), + }; + + assert_eq!(old_owner.operation_id(), 1); + assert_eq!(current_owner.operation_id(), 1); + assert!(!current_registry.commit_ready(&old_owner)); + assert_eq!( + current_registry.snapshot("worker-1"), + MemberRuntimeSnapshot::Attaching { operation_id: 1 } + ); + assert!(current_registry.commit_ready(¤t_owner)); + } + + #[test] + fn remove_owner_from_another_session_generation_cannot_finish_or_restore() { + let old_registry = MemberRuntimeRegistry::new(53); + assert!(old_registry.seed_ready("worker-1")); + let old_owner = match old_registry.begin_remove("worker-1") { + BeginRemove::Start(lease) => lease, + _ => panic!("old session remove must start"), + }; + let current_registry = MemberRuntimeRegistry::new(54); + assert!(current_registry.seed_ready("worker-1")); + let current_owner = match current_registry.begin_remove("worker-1") { + BeginRemove::Start(lease) => lease, + _ => panic!("current session remove must start"), + }; + + assert_eq!(old_owner.operation_id(), 1); + assert_eq!(current_owner.operation_id(), 1); + assert!(!current_registry.finish_remove(&old_owner)); + assert!(!current_registry.restore_attach_required_after_remove_persist_error( + &old_owner, + failure("stale_session", "Stale session remove failed") + )); + assert_eq!( + current_registry.snapshot("worker-1"), + MemberRuntimeSnapshot::Removing { operation_id: 1 } + ); + assert!(current_registry.finish_remove(¤t_owner)); + } +} diff --git a/crates/aionui-team/src/message_projection.rs b/crates/aionui-team/src/message_projection.rs index e6b40f548..fbe87e316 100644 --- a/crates/aionui-team/src/message_projection.rs +++ b/crates/aionui-team/src/message_projection.rs @@ -15,6 +15,7 @@ const TEXT_MESSAGE_TYPE: &str = "text"; #[derive(Debug, Clone, PartialEq, Eq)] pub enum TeamProjectionSource { User, + TeamSystem, Teammate { from_slot_id: String, from_name: String, @@ -84,9 +85,32 @@ impl TeamProjectionRequest { } } + pub fn team_system_visible( + team_id: impl Into, + slot_id: impl Into, + conversation_id: impl Into, + content: impl Into, + mailbox_message_id: impl Into, + ) -> Self { + let team_id = team_id.into(); + let conversation_id = conversation_id.into(); + let mailbox_message_id = mailbox_message_id.into(); + Self { + dedupe_key: Some(teammate_dedupe_key(&team_id, &mailbox_message_id, &conversation_id)), + team_id, + slot_id: slot_id.into(), + conversation_id, + source: TeamProjectionSource::TeamSystem, + content: content.into(), + files: Vec::new(), + visibility: TeamVisibilityPolicy::teammate_message(), + } + } + fn should_insert_visible_bubble(&self) -> bool { match self.source { TeamProjectionSource::User => self.visibility.insert_user_visible_bubble, + TeamProjectionSource::TeamSystem => self.visibility.insert_teammate_visible_bubble, TeamProjectionSource::Teammate { .. } => self.visibility.insert_teammate_visible_bubble, } } @@ -171,14 +195,13 @@ where let row = Self::build_message_row(&request, &msg_id, aionui_common::now_ms())?; self.store.insert_projected_message(&row).await?; - if let TeamProjectionSource::Teammate { - from_slot_id, - from_name, - sender_backend, - sender_conversation_id, - } = &request.source - { - let payload = serde_json::json!({ + let teammate_event_payload = match &request.source { + TeamProjectionSource::Teammate { + from_slot_id, + from_name, + sender_backend, + sender_conversation_id, + } => Some(serde_json::json!({ "team_id": request.team_id, "slot_id": request.slot_id, "conversation_id": request.conversation_id, @@ -189,7 +212,22 @@ where "teammate_message": true, "sender_backend": sender_backend, "sender_conversation_id": sender_conversation_id, - }); + })), + TeamProjectionSource::TeamSystem => Some(serde_json::json!({ + "team_id": request.team_id, + "slot_id": request.slot_id, + "conversation_id": request.conversation_id, + "msg_id": msg_id, + "content": request.content, + "from_slot_id": "team_system", + "from_name": "team_system", + "teammate_message": true, + "sender_backend": null, + "sender_conversation_id": null, + })), + TeamProjectionSource::User => None, + }; + if let Some(payload) = teammate_event_payload { self.broadcaster .broadcast(WebSocketMessage::new(TEAMMATE_MESSAGE_EVENT, payload)); } @@ -200,6 +238,7 @@ where conversation_id = %request.conversation_id, event_name = match request.source { TeamProjectionSource::User => "message.stream", + TeamProjectionSource::TeamSystem => TEAMMATE_MESSAGE_EVENT, TeamProjectionSource::Teammate { .. } => TEAMMATE_MESSAGE_EVENT, }, outcome = "inserted", @@ -228,6 +267,16 @@ where }), ) } + TeamProjectionSource::TeamSystem => ( + "left", + serde_json::json!({ + "content": request.content, + "teammate_message": true, + "sender_name": "team_system", + "sender_backend": null, + "sender_conversation_id": null, + }), + ), TeamProjectionSource::Teammate { from_slot_id: _, from_name, diff --git a/crates/aionui-team/src/ports.rs b/crates/aionui-team/src/ports.rs index e50c3d9fc..98d625d09 100644 --- a/crates/aionui-team/src/ports.rs +++ b/crates/aionui-team/src/ports.rs @@ -28,6 +28,17 @@ pub trait TeamConversationLookupPort: Send + Sync { #[async_trait] pub trait TeamAssistantCatalogPort: Send + Sync { async fn list_team_selectable_assistants(&self) -> Result, TeamError>; + + async fn resolve_team_selectable_assistant( + &self, + assistant_id: &str, + ) -> Result, TeamError> { + Ok(self + .list_team_selectable_assistants() + .await? + .into_iter() + .find(|assistant| assistant.assistant_id == assistant_id)) + } } #[derive(Debug, Clone, PartialEq, Eq)] @@ -82,7 +93,7 @@ pub type AgentTurnStartedCallback = #[derive(Debug, Clone, PartialEq, Eq)] pub struct AgentTurnStarted { - pub team_run_id: String, + pub team_run_id: Option, pub slot_id: String, pub role: TeamRunTargetRole, pub conversation_id: String, diff --git a/crates/aionui-team/src/prompt_dump.rs b/crates/aionui-team/src/prompt_dump.rs new file mode 100644 index 000000000..af95025a1 --- /dev/null +++ b/crates/aionui-team/src/prompt_dump.rs @@ -0,0 +1,151 @@ +use std::fs; +use std::io; +use std::path::{Path, PathBuf}; +use std::time::{SystemTime, UNIX_EPOCH}; + +use serde_json::Value; + +use crate::types::TeammateRole; + +#[derive(Clone, Debug, Default)] +pub struct TeamPromptDumpConfig { + dump_dir: Option, +} + +impl TeamPromptDumpConfig { + pub fn disabled() -> Self { + Self { dump_dir: None } + } + + pub fn enabled(dump_dir: impl AsRef) -> Self { + Self { + dump_dir: Some(dump_dir.as_ref().to_path_buf()), + } + } + + pub fn from_data_dir(data_dir: impl AsRef, enabled: bool) -> Self { + if enabled { + Self::enabled(data_dir.as_ref().join("prompt-dumps")) + } else { + Self::disabled() + } + } + + fn dump_dir(&self) -> Option<&Path> { + self.dump_dir.as_deref() + } +} + +pub(crate) struct TeamWakePromptDump<'a> { + pub team_id: &'a str, + pub slot_id: &'a str, + pub conversation_id: &'a str, + pub role: TeammateRole, + pub needs_role_prompt: bool, + pub unread_count: usize, + pub prompt: &'a str, +} + +pub(crate) struct TeamToolsListDump<'a> { + pub team_id: &'a str, + pub caller_slot_id: &'a str, + pub caller_role: TeammateRole, + pub tools: &'a [Value], +} + +pub(crate) fn dump_team_wake_prompt( + config: &TeamPromptDumpConfig, + dump: TeamWakePromptDump<'_>, +) -> io::Result> { + let Some(dump_dir) = config.dump_dir() else { + return Ok(None); + }; + fs::create_dir_all(dump_dir)?; + + let created_at_ms = current_time_ms(); + let path = dump_dir.join(format!( + "{}-team-wake-prompt-{}-{}.txt", + current_time_nanos(), + sanitize_segment(dump.team_id), + sanitize_segment(dump.slot_id) + )); + + let body = format!( + "kind: team-wake-prompt\nscope: team-wake-content-only\nnot_final_agent_input: true\nteam_id: {}\nslot_id: {}\nconversation_id: {}\nrole: {}\nneeds_role_prompt: {}\nunread_count: {}\ncreated_at_ms: {}\n\n---- prompt ----\n{}\n", + dump.team_id, + dump.slot_id, + dump.conversation_id, + role_label(dump.role), + dump.needs_role_prompt, + dump.unread_count, + created_at_ms, + dump.prompt + ); + fs::write(&path, body)?; + Ok(Some(path)) +} + +pub(crate) fn dump_team_tools_list( + config: &TeamPromptDumpConfig, + dump: TeamToolsListDump<'_>, +) -> io::Result> { + let Some(dump_dir) = config.dump_dir() else { + return Ok(None); + }; + fs::create_dir_all(dump_dir)?; + + let path = dump_dir.join(format!( + "{}-team-tools-list-{}-{}.json", + current_time_nanos(), + sanitize_segment(dump.team_id), + sanitize_segment(dump.caller_slot_id) + )); + let body = serde_json::to_vec_pretty(&serde_json::json!({ + "kind": "team-tools-list", + "scope": "team-mcp-server-tools-only", + "not_final_agent_tools": true, + "team_id": dump.team_id, + "caller_slot_id": dump.caller_slot_id, + "caller_role": role_label(dump.caller_role), + "created_at_ms": current_time_ms(), + "tools": dump.tools, + }))?; + fs::write(&path, body)?; + Ok(Some(path)) +} + +fn role_label(role: TeammateRole) -> &'static str { + match role { + TeammateRole::Lead => "lead", + TeammateRole::Teammate => "teammate", + } +} + +fn current_time_ms() -> u128 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_millis()) + .unwrap_or(0) +} + +fn current_time_nanos() -> u128 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_nanos()) + .unwrap_or(0) +} + +fn sanitize_segment(value: &str) -> String { + let segment: String = value + .chars() + .map(|ch| { + if ch.is_ascii_alphanumeric() || matches!(ch, '-' | '_' | '.') { + ch + } else { + '_' + } + }) + .take(96) + .collect(); + if segment.is_empty() { "none".to_owned() } else { segment } +} diff --git a/crates/aionui-team/src/prompts/mod.rs b/crates/aionui-team/src/prompts/mod.rs index ce87be415..26ce89c28 100644 --- a/crates/aionui-team/src/prompts/mod.rs +++ b/crates/aionui-team/src/prompts/mod.rs @@ -1,8 +1,10 @@ -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; + +mod wake_summary; pub use aionui_team_prompts::AvailableAssistant; -use crate::types::{MailboxMessage, MailboxMessageType, TaskStatus, TeamAgent, TeamTask}; +use crate::types::{MailboxMessage, MailboxMessageType, TeamAgent, TeamTask}; fn to_prompt_role(role: crate::types::TeammateRole) -> aionui_team_prompts::TeamPromptRole { match role { @@ -32,14 +34,17 @@ fn to_prompt_agent(agent: &TeamAgent) -> aionui_team_prompts::TeamPromptAgent { /// staffing the team. Callers should only include assistants that are both /// enabled and team-selectable. pub fn build_lead_prompt( + agent: &TeamAgent, team_name: &str, members: &[TeamAgent], available_assistants: &[AvailableAssistant], ) -> String { + let prompt_agent = to_prompt_agent(agent); let prompt_members: Vec<_> = members.iter().map(to_prompt_agent).collect(); let renamed: HashMap = HashMap::new(); let body = aionui_team_prompts::build_lead_prompt(&aionui_team_prompts::LeadPromptParams { + agent: &prompt_agent, team_name, teammates: &prompt_members, available_agent_types: &[], @@ -82,7 +87,12 @@ pub fn build_teammate_prompt(agent: &TeamAgent, team_name: &str, members: &[Team }) } -pub fn build_wake_payload(agent: &TeamAgent, tasks: &[TeamTask], unread_messages: &[MailboxMessage]) -> String { +pub fn build_wake_payload( + agent: &TeamAgent, + tasks: &[TeamTask], + unread_messages: &[MailboxMessage], + current_slot_ids: &HashSet, +) -> String { let mut payload = String::with_capacity(2048); if !unread_messages.is_empty() { @@ -106,33 +116,7 @@ pub fn build_wake_payload(agent: &TeamAgent, tasks: &[TeamTask], unread_messages payload.push_str("## New Messages\n\nNo new messages.\n\n"); } - if !tasks.is_empty() { - payload.push_str("## Current Task Board\n\n"); - payload.push_str("| ID | Subject | Status | Owner | Blocked By |\n"); - payload.push_str("|---|---|---|---|---|\n"); - for task in tasks { - let status = match task.status { - TaskStatus::Pending => "pending", - TaskStatus::InProgress => "in_progress", - TaskStatus::Completed => "completed", - TaskStatus::Deleted => "deleted", - }; - let owner = task.owner.as_deref().unwrap_or("-"); - let blocked = if task.blocked_by.is_empty() { - "-".to_owned() - } else { - task.blocked_by.join(", ") - }; - let short_id = if task.id.len() > 8 { &task.id[..8] } else { &task.id }; - payload.push_str(&format!( - "| {short_id}… | {} | {status} | {owner} | {blocked} |\n", - task.subject, - )); - } - payload.push('\n'); - } else { - payload.push_str("## Current Task Board\n\nNo tasks on the board.\n\n"); - } + payload.push_str(&wake_summary::render_task_board_summary(agent, tasks, current_slot_ids)); payload.push_str(&format!( "You are **{}** (role: {}). Proceed with your work.\n", @@ -145,7 +129,7 @@ pub fn build_wake_payload(agent: &TeamAgent, tasks: &[TeamTask], unread_messages #[cfg(test)] mod tests { use super::*; - use crate::types::TeammateRole; + use crate::types::{TaskStatus, TeammateRole}; fn make_lead() -> TeamAgent { TeamAgent { @@ -208,6 +192,10 @@ mod tests { } } + fn roster(ids: &[&str]) -> HashSet { + ids.iter().map(|id| (*id).to_owned()).collect() + } + // -- Lead prompt ---------------------------------------------------------- fn default_assistants() -> Vec { @@ -223,36 +211,37 @@ mod tests { #[test] fn lead_prompt_contains_team_name() { let assistants = default_assistants(); - let prompt = build_lead_prompt("Alpha", &[], &assistants); + let prompt = build_lead_prompt(&make_lead(), "Alpha", &[], &assistants); assert!(prompt.contains("\"Alpha\"")); } #[test] - fn lead_prompt_contains_member_list() { + fn lead_prompt_uses_tool_for_member_list() { let assistants = default_assistants(); let members = vec![make_lead(), make_teammate("w1", "Worker1")]; - let prompt = build_lead_prompt("Alpha", &members, &assistants); + let prompt = build_lead_prompt(&make_lead(), "Alpha", &members, &assistants); - // AionUi bullet format: `- {name} ({backend}, status: {status})` - assert!(prompt.contains("- Lead (acp, status:")); - assert!(prompt.contains("- Worker1 (acp, status:")); + assert!(!prompt.contains("- Lead (acp, status:")); + assert!(!prompt.contains("- Worker1 (acp, status:")); + assert!(prompt.contains("team_members")); } #[test] fn lead_prompt_contains_core_sections() { let assistants = default_assistants(); - let prompt = build_lead_prompt("Alpha", &[], &assistants); + let prompt = build_lead_prompt(&make_lead(), "Alpha", &[], &assistants); - // Workflow — 15-step procedure with model listing at step 3 + // Workflow uses tools for dynamic state. assert!(prompt.contains("## Workflow")); - assert!(prompt.contains("FIRST call `team_list_assistants`")); + assert!(prompt.contains("FIRST call `team_members`")); + assert!(prompt.contains("call `team_list_assistants`")); assert!(prompt.contains("Wait for explicit confirmation before using team_spawn_agent")); assert!(prompt.contains("End your turn after the proposal")); - // Model Selection Guidelines - assert!(prompt.contains("## Model Selection Guidelines")); - assert!(prompt.contains("exact model ID strings")); - assert!(prompt.contains("omit the model parameter")); + // Assistant selection, not model selection. + assert!(prompt.contains("## Assistant Selection Guidelines")); + assert!(!prompt.contains("team_list_models")); + assert!(prompt.contains("Do not pass a model to `team_spawn_agent`")); // Conversation Style — don't pitch proposals up-front assert!(prompt.contains("## Conversation Style")); @@ -274,42 +263,47 @@ mod tests { } #[test] - fn lead_prompt_includes_available_assistants_section() { + fn lead_prompt_omits_dynamic_available_assistants_section() { let assistants = default_assistants(); - let prompt = build_lead_prompt("Alpha", &[], &assistants); + let prompt = build_lead_prompt(&make_lead(), "Alpha", &[], &assistants); - assert!(prompt.contains("## Available Assistants for Spawning")); - assert!(prompt.contains("- `word-creator` (Word Creator) — Drafts Word documents")); - assert!(prompt.contains("skills: docx, formatting")); - assert!(prompt.contains("Pass the assistant's ID as `assistant_id`")); - assert!(!prompt.contains("backend: claude")); + assert!(!prompt.contains("## Available Assistants for Spawning")); + assert!(!prompt.contains("- `word-creator` (Word Creator) — Drafts Word documents")); + assert!(prompt.contains("team_list_assistants")); } #[test] fn lead_prompt_omits_available_assistants_section_when_empty() { - let prompt = build_lead_prompt("Alpha", &[], &[]); + let prompt = build_lead_prompt(&make_lead(), "Alpha", &[], &[]); assert!(!prompt.contains("## Available Assistants for Spawning")); } #[test] - fn lead_prompt_no_members_shows_empty_lineup_copy() { + fn lead_prompt_does_not_include_dynamic_member_snapshot() { let assistants = default_assistants(); - let prompt = build_lead_prompt("Solo", &[], &assistants); - assert!(prompt.contains("(no teammates yet")); - assert!(prompt.contains("propose the lineup to the user first")); + let prompt = build_lead_prompt(&make_lead(), "Solo", &[], &assistants); + assert!(!prompt.contains("## Your Teammates")); + assert!(!prompt.contains("(no teammates yet")); + assert!(prompt.to_lowercase().contains("first team turn")); + assert!(prompt.contains("team_members")); } #[test] fn lead_prompt_has_no_unsubstituted_placeholders() { let assistants = default_assistants(); let members = vec![make_lead(), make_teammate("w1", "Worker1")]; - let prompt = build_lead_prompt("Alpha", &members, &assistants); + let prompt = build_lead_prompt(&make_lead(), "Alpha", &members, &assistants); assert!( !prompt.contains("${"), "unsubstituted template placeholder leaked:\n{prompt}" ); assert!(!prompt.contains("assistant or backend")); assert!(!prompt.contains("Available Generic Backends")); + assert!(!prompt.contains("## Your Teammates")); + assert!(!prompt.contains("## Available Assistants for Spawning")); + assert!(prompt.contains("Name: Lead")); + assert!(prompt.contains("Slot ID: lead-1")); + assert!(prompt.contains("Role: lead")); } // -- Teammate prompt ------------------------------------------------------ @@ -322,8 +316,10 @@ mod tests { assert!(prompt.contains("## Team Governance")); assert!(prompt.contains("Name: Worker1")); + assert!(prompt.contains("Slot ID: w1")); assert!(prompt.contains("Team: Alpha")); assert!(prompt.contains("Leader: Lead")); + assert!(!prompt.contains("Teammates:")); } #[test] @@ -355,7 +351,7 @@ mod tests { fn wake_payload_with_messages() { let agent = make_lead(); let msgs = vec![make_message("w1", "Task A done", MailboxMessageType::Message)]; - let payload = build_wake_payload(&agent, &[], &msgs); + let payload = build_wake_payload(&agent, &[], &msgs, &roster(&["lead-1", "w1"])); assert!(payload.contains("New Messages")); assert!(payload.contains("`w1`")); @@ -368,7 +364,7 @@ mod tests { let agent = make_lead(); let mut msg = make_message("w1", "idle", MailboxMessageType::IdleNotification); msg.summary = Some("Finished feature X".into()); - let payload = build_wake_payload(&agent, &[], &[msg]); + let payload = build_wake_payload(&agent, &[], &[msg], &roster(&["lead-1", "w1"])); assert!(payload.contains("[idle_notification]")); assert!(payload.contains("Summary: Finished feature X")); @@ -378,7 +374,7 @@ mod tests { fn wake_payload_with_shutdown_request() { let agent = make_teammate("w1", "W"); let msg = make_message("lead-1", "No longer needed", MailboxMessageType::ShutdownRequest); - let payload = build_wake_payload(&agent, &[], &[msg]); + let payload = build_wake_payload(&agent, &[], &[msg], &roster(&["lead-1", "w1"])); assert!(payload.contains("[shutdown_request]")); assert!(payload.contains("No longer needed")); @@ -395,9 +391,10 @@ mod tests { ), make_task("bbbbbbbb-1234-5678-9abc-def012345678", "Test Y", TaskStatus::Pending), ]; - let payload = build_wake_payload(&agent, &tasks, &[]); + let payload = build_wake_payload(&agent, &tasks, &[], &roster(&["lead-1", "worker-1"])); - assert!(payload.contains("Current Task Board")); + assert!(payload.contains("Current Task Board Summary")); + assert!(payload.contains("Showing 2 of 2 tasks.")); assert!(payload.contains("Implement X")); assert!(payload.contains("in_progress")); assert!(payload.contains("Test Y")); @@ -410,15 +407,17 @@ mod tests { let agent = make_lead(); let mut task = make_task("cccccccc-1234-5678-9abc-def012345678", "Deploy", TaskStatus::Pending); task.blocked_by = vec!["task-a".into(), "task-b".into()]; - let payload = build_wake_payload(&agent, &[task], &[]); + let payload = build_wake_payload(&agent, &[task], &[], &roster(&["lead-1", "worker-1"])); - assert!(payload.contains("task-a, task-b")); + assert!(payload.contains("task-a…")); + assert!(payload.contains("task-b…")); + assert!(!payload.contains("task-a, task-b")); } #[test] fn wake_payload_empty() { let agent = make_lead(); - let payload = build_wake_payload(&agent, &[], &[]); + let payload = build_wake_payload(&agent, &[], &[], &roster(&["lead-1"])); assert!(payload.contains("No new messages")); assert!(payload.contains("No tasks on the board")); @@ -428,7 +427,7 @@ mod tests { #[test] fn wake_payload_contains_agent_identity() { let agent = make_teammate("w1", "Worker1"); - let payload = build_wake_payload(&agent, &[], &[]); + let payload = build_wake_payload(&agent, &[], &[], &roster(&["w1"])); assert!(payload.contains("**Worker1**")); assert!(payload.contains("teammate")); @@ -438,7 +437,7 @@ mod tests { fn wake_payload_short_task_id_no_truncation() { let agent = make_lead(); let task = make_task("short", "Short ID Task", TaskStatus::Pending); - let payload = build_wake_payload(&agent, &[task], &[]); + let payload = build_wake_payload(&agent, &[task], &[], &roster(&["lead-1", "worker-1"])); assert!(payload.contains("short…")); } } diff --git a/crates/aionui-team/src/prompts/wake_summary.rs b/crates/aionui-team/src/prompts/wake_summary.rs new file mode 100644 index 000000000..d4593cafa --- /dev/null +++ b/crates/aionui-team/src/prompts/wake_summary.rs @@ -0,0 +1,980 @@ +use std::cmp::Ordering; +use std::collections::HashSet; + +use crate::types::{TaskStatus, TeamAgent, TeamTask, TeammateRole}; + +pub(super) const MAX_WAKE_TASK_ROWS: usize = 40; +pub(super) const MAX_RECENT_COMPLETED_ROWS: usize = 8; +pub(super) const MAX_TEAMMATE_RECENT_COMPLETED_ROWS: usize = 3; +pub(super) const MAX_SUBJECT_CHARS: usize = 120; +pub(super) const MAX_BLOCKED_BY_IDS: usize = 5; + +pub(super) fn render_task_board_summary( + agent: &TeamAgent, + tasks: &[TeamTask], + current_slot_ids: &HashSet, +) -> String { + let mut output = String::with_capacity(1024); + output.push_str("## Current Task Board Summary\n\n"); + + if tasks.is_empty() { + output.push_str("No tasks on the board.\n\n"); + return output; + } + + let visible_tasks = filter_summary_input(tasks, current_slot_ids); + let selection = select_summary_tasks(agent, &visible_tasks); + let selected = selection.tasks; + let displayed_ids: HashSet<&str> = selected.iter().map(|task| task.id.as_str()).collect(); + let hidden_active = tasks + .iter() + .filter(|task| is_active(task.status)) + .filter(|task| !displayed_ids.contains(task.id.as_str())) + .count(); + let displayed_completed = selected + .iter() + .filter(|task| task.status == TaskStatus::Completed) + .count(); + let total_completed = tasks.iter().filter(|task| task.status == TaskStatus::Completed).count(); + let total_deleted = tasks.iter().filter(|task| task.status == TaskStatus::Deleted).count(); + let hidden_completed = total_completed.saturating_sub(displayed_completed); + let hidden_deleted = total_deleted; + + output.push_str(&format!("Showing {} of {} tasks.\n", selected.len(), tasks.len())); + output.push_str(&format!( + "Hidden: active={hidden_active}, completed={hidden_completed}, deleted={hidden_deleted}.\n" + )); + if hidden_active + hidden_completed + hidden_deleted > 0 { + output.push_str("More tasks are available via `team_task_list`; use filters when needed.\n"); + } + output.push('\n'); + output.push_str("| ID | Subject | Status | Owner | Blocked By |\n"); + output.push_str("|---|---|---|---|---|\n"); + for task in selected { + output.push_str(&format!( + "| {} | {} | {} | {} | {} |\n", + short_task_id(&task.id), + truncate_subject(&task.subject), + task.status, + render_owner(task, current_slot_ids), + render_blocked_by(&task.blocked_by), + )); + } + output.push('\n'); + output +} + +struct SummarySelection<'a> { + tasks: Vec<&'a TeamTask>, +} + +fn filter_summary_input<'a>(tasks: &'a [TeamTask], current_slot_ids: &HashSet) -> Vec<&'a TeamTask> { + let visible_base_ids: HashSet<&str> = tasks + .iter() + .filter(|task| is_active(task.status)) + .filter(|task| is_current_or_ownerless(task, current_slot_ids)) + .map(|task| task.id.as_str()) + .collect(); + + let visible_blocker_ids: HashSet<&str> = tasks + .iter() + .filter(|task| visible_base_ids.contains(task.id.as_str())) + .flat_map(|task| task.blocked_by.iter().map(String::as_str)) + .collect(); + + tasks + .iter() + .filter(|task| { + if is_current_or_ownerless(task, current_slot_ids) { + return true; + } + + is_active(task.status) + && is_removed_owner(task, current_slot_ids) + && visible_blocker_ids.contains(task.id.as_str()) + }) + .collect() +} + +fn select_summary_tasks<'a>(agent: &TeamAgent, tasks: &[&'a TeamTask]) -> SummarySelection<'a> { + let own_active_ids: HashSet<&str> = tasks + .iter() + .copied() + .filter(|task| is_active(task.status)) + .filter(|task| task.owner.as_deref() == Some(agent.slot_id.as_str())) + .map(|task| task.id.as_str()) + .collect(); + + let blocker_ids: HashSet<&str> = tasks + .iter() + .copied() + .filter(|task| own_active_ids.contains(task.id.as_str())) + .flat_map(|task| task.blocked_by.iter().map(String::as_str)) + .collect(); + + let mut blocked_by_own_ids: HashSet<&str> = tasks + .iter() + .copied() + .filter(|task| own_active_ids.contains(task.id.as_str())) + .flat_map(|task| task.blocks.iter().map(String::as_str)) + .collect(); + blocked_by_own_ids.extend( + tasks + .iter() + .copied() + .filter(|task| is_active(task.status)) + .filter(|task| task.blocked_by.iter().any(|id| own_active_ids.contains(id.as_str()))) + .map(|task| task.id.as_str()), + ); + + if agent.role == TeammateRole::Teammate { + return select_teammate_summary_tasks(tasks, &own_active_ids, &blocker_ids, &blocked_by_own_ids, agent); + } + + select_leader_summary_tasks(tasks, &own_active_ids, &blocker_ids, &blocked_by_own_ids) +} + +fn select_leader_summary_tasks<'a>( + tasks: &[&'a TeamTask], + own_active_ids: &HashSet<&str>, + blocker_ids: &HashSet<&str>, + blocked_by_own_ids: &HashSet<&str>, +) -> SummarySelection<'a> { + let mut active: Vec<&TeamTask> = tasks.iter().copied().filter(|task| is_active(task.status)).collect(); + active.sort_by(|left, right| compare_active_tasks(left, right, own_active_ids, blocker_ids, blocked_by_own_ids)); + + let mut selected: Vec<&TeamTask> = active.into_iter().take(MAX_WAKE_TASK_ROWS).collect(); + if selected.len() < MAX_WAKE_TASK_ROWS { + let mut completed: Vec<&TeamTask> = tasks + .iter() + .copied() + .filter(|task| task.status == TaskStatus::Completed) + .collect(); + completed.sort_by(|left, right| { + right + .updated_at + .cmp(&left.updated_at) + .then_with(|| left.created_at.cmp(&right.created_at)) + .then_with(|| left.id.cmp(&right.id)) + }); + let remaining = MAX_WAKE_TASK_ROWS - selected.len(); + selected.extend(completed.into_iter().take(MAX_RECENT_COMPLETED_ROWS.min(remaining))); + } + + SummarySelection { tasks: selected } +} + +fn select_teammate_summary_tasks<'a>( + tasks: &[&'a TeamTask], + own_active_ids: &HashSet<&str>, + blocker_ids: &HashSet<&str>, + blocked_by_own_ids: &HashSet<&str>, + agent: &TeamAgent, +) -> SummarySelection<'a> { + let mut relevant_active: Vec<&TeamTask> = tasks + .iter() + .copied() + .filter(|task| is_active(task.status)) + .filter(|task| { + own_active_ids.contains(task.id.as_str()) + || blocker_ids.contains(task.id.as_str()) + || blocked_by_own_ids.contains(task.id.as_str()) + }) + .collect(); + relevant_active + .sort_by(|left, right| compare_active_tasks(left, right, own_active_ids, blocker_ids, blocked_by_own_ids)); + + let mut selected: Vec<&TeamTask> = relevant_active.into_iter().take(MAX_WAKE_TASK_ROWS).collect(); + if selected.len() < MAX_WAKE_TASK_ROWS { + let mut completed: Vec<&TeamTask> = tasks + .iter() + .copied() + .filter(|task| task.status == TaskStatus::Completed) + .filter(|task| task.owner.as_deref() == Some(agent.slot_id.as_str())) + .collect(); + completed.sort_by(|left, right| { + right + .updated_at + .cmp(&left.updated_at) + .then_with(|| left.created_at.cmp(&right.created_at)) + .then_with(|| left.id.cmp(&right.id)) + }); + let remaining = MAX_WAKE_TASK_ROWS - selected.len(); + selected.extend( + completed + .into_iter() + .take(MAX_TEAMMATE_RECENT_COMPLETED_ROWS.min(remaining)), + ); + } + + SummarySelection { tasks: selected } +} + +fn compare_active_tasks( + left: &TeamTask, + right: &TeamTask, + own_active_ids: &HashSet<&str>, + blocker_ids: &HashSet<&str>, + blocked_by_own_ids: &HashSet<&str>, +) -> Ordering { + active_rank(left, own_active_ids, blocker_ids, blocked_by_own_ids) + .cmp(&active_rank(right, own_active_ids, blocker_ids, blocked_by_own_ids)) + .then_with(|| status_rank(left.status).cmp(&status_rank(right.status))) + .then_with(|| compare_same_status_recency(left, right)) + .then_with(|| left.id.cmp(&right.id)) +} + +fn compare_same_status_recency(left: &TeamTask, right: &TeamTask) -> Ordering { + match left.status { + TaskStatus::InProgress => right + .updated_at + .cmp(&left.updated_at) + .then_with(|| left.created_at.cmp(&right.created_at)), + TaskStatus::Pending => right + .updated_at + .cmp(&left.updated_at) + .then_with(|| left.created_at.cmp(&right.created_at)), + TaskStatus::Completed | TaskStatus::Deleted => Ordering::Equal, + } +} + +fn active_rank( + task: &TeamTask, + own_active_ids: &HashSet<&str>, + blocker_ids: &HashSet<&str>, + blocked_by_own_ids: &HashSet<&str>, +) -> u8 { + if own_active_ids.contains(task.id.as_str()) { + 0 + } else if blocker_ids.contains(task.id.as_str()) { + 1 + } else if blocked_by_own_ids.contains(task.id.as_str()) { + 2 + } else if task.status == TaskStatus::InProgress { + 3 + } else { + 4 + } +} + +fn status_rank(status: TaskStatus) -> u8 { + match status { + TaskStatus::InProgress => 0, + TaskStatus::Pending => 1, + TaskStatus::Completed => 2, + TaskStatus::Deleted => 3, + } +} + +fn is_active(status: TaskStatus) -> bool { + matches!(status, TaskStatus::Pending | TaskStatus::InProgress) +} + +fn is_current_or_ownerless(task: &TeamTask, current_slot_ids: &HashSet) -> bool { + match task.owner.as_deref() { + Some(owner) => current_slot_ids.contains(owner), + None => true, + } +} + +fn is_removed_owner(task: &TeamTask, current_slot_ids: &HashSet) -> bool { + task.owner + .as_deref() + .is_some_and(|owner| !current_slot_ids.contains(owner)) +} + +fn render_owner(task: &TeamTask, current_slot_ids: &HashSet) -> String { + match task.owner.as_deref() { + Some(owner) if current_slot_ids.contains(owner) => owner.to_owned(), + Some(owner) => format!("{owner} (removed)"), + None => "-".to_owned(), + } +} + +fn short_task_id(id: &str) -> String { + format!("{}…", id.chars().take(8).collect::()) +} + +fn truncate_subject(subject: &str) -> String { + if subject.chars().count() <= MAX_SUBJECT_CHARS { + subject.to_owned() + } else { + format!("{}...", subject.chars().take(MAX_SUBJECT_CHARS).collect::()) + } +} + +fn render_blocked_by(blocked_by: &[String]) -> String { + if blocked_by.is_empty() { + return "-".to_owned(); + } + let mut parts: Vec = blocked_by + .iter() + .take(MAX_BLOCKED_BY_IDS) + .map(|id| short_task_id(id)) + .collect(); + if blocked_by.len() > MAX_BLOCKED_BY_IDS { + parts.push(format!("+{} more", blocked_by.len() - MAX_BLOCKED_BY_IDS)); + } + parts.join(", ") +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashSet; + + fn agent(slot_id: &str) -> TeamAgent { + agent_with_role(slot_id, TeammateRole::Lead) + } + + fn teammate(slot_id: &str) -> TeamAgent { + agent_with_role(slot_id, TeammateRole::Teammate) + } + + fn agent_with_role(slot_id: &str, role: TeammateRole) -> TeamAgent { + TeamAgent { + slot_id: slot_id.to_owned(), + name: "Worker".to_owned(), + role, + conversation_id: "conv-1".to_owned(), + backend: "acp".to_owned(), + model: "claude".to_owned(), + assistant_id: None, + status: None, + conversation_type: None, + cli_path: None, + } + } + + fn task( + id: &str, + subject: &str, + status: TaskStatus, + owner: Option<&str>, + created_at: i64, + updated_at: i64, + ) -> TeamTask { + TeamTask { + id: id.to_owned(), + team_id: "team-1".to_owned(), + subject: subject.to_owned(), + description: None, + status, + owner: owner.map(str::to_owned), + blocked_by: Vec::new(), + blocks: Vec::new(), + metadata: None, + created_at, + updated_at, + } + } + + fn row_count(payload: &str) -> usize { + payload + .lines() + .filter(|line| line.starts_with("| ") && !line.starts_with("| ID ") && !line.starts_with("|---")) + .count() + } + + fn roster(ids: &[&str]) -> HashSet { + ids.iter().map(|id| (*id).to_owned()).collect() + } + + fn default_roster() -> HashSet { + roster(&[ + "worker-1", + "worker-2", + "worker-3", + "worker-4", + "019f4056-fd18-7411-ab09-8868ff17cb36", + ]) + } + + #[test] + fn empty_board_renders_empty_summary() { + let payload = render_task_board_summary(&agent("worker-1"), &[], &roster(&["worker-1"])); + assert!(payload.contains("## Current Task Board Summary")); + assert!(payload.contains("No tasks on the board.")); + } + + #[test] + fn active_tasks_under_limit_all_render() { + let tasks = vec![ + task( + "aaaaaaaa-1111", + "Owned active", + TaskStatus::InProgress, + Some("worker-1"), + 1, + 10, + ), + task( + "bbbbbbbb-1111", + "Other pending", + TaskStatus::Pending, + Some("worker-2"), + 2, + 20, + ), + ]; + let payload = render_task_board_summary(&agent("worker-1"), &tasks, &default_roster()); + assert!(payload.contains("Showing 2 of 2 tasks.")); + assert!(payload.contains("Hidden: active=0, completed=0, deleted=0.")); + assert!(payload.contains("| aaaaaaaa… | Owned active | in_progress | worker-1 | - |")); + assert!(payload.contains("| bbbbbbbb… | Other pending | pending | worker-2 | - |")); + } + + #[test] + fn completed_tasks_are_limited_to_recent_rows() { + let tasks: Vec<_> = (0..10) + .map(|index| { + task( + &format!("done{index:04}-1111"), + &format!("Completed {index}"), + TaskStatus::Completed, + Some("worker-2"), + index, + index, + ) + }) + .collect(); + let payload = render_task_board_summary(&agent("worker-1"), &tasks, &default_roster()); + assert_eq!(row_count(&payload), MAX_RECENT_COMPLETED_ROWS); + assert!(payload.contains("Completed 9")); + assert!(payload.contains("Completed 2")); + assert!(!payload.contains("Completed 1")); + assert!(payload.contains("Hidden: active=0, completed=2, deleted=0.")); + assert!(payload.contains("More tasks are available via `team_task_list`; use filters when needed.")); + assert!(!payload.contains("team_task_list({})")); + assert!(!payload.contains(r#"team_task_list({"status":["pending","in_progress"]})"#)); + } + + #[test] + fn deleted_tasks_are_hidden_and_counted() { + let tasks = vec![ + task("aaaaaaaa-1111", "Active", TaskStatus::Pending, Some("worker-1"), 1, 1), + task( + "dddddddd-1111", + "Deleted history", + TaskStatus::Deleted, + Some("worker-1"), + 2, + 2, + ), + ]; + let payload = render_task_board_summary(&agent("worker-1"), &tasks, &default_roster()); + assert!(payload.contains("Active")); + assert!(!payload.contains("Deleted history")); + assert!(payload.contains("Hidden: active=0, completed=0, deleted=1.")); + } + + #[test] + fn active_tasks_over_limit_are_truncated_and_counted() { + let tasks: Vec<_> = (0..45) + .map(|index| { + task( + &format!("task{index:04}-1111"), + &format!("Active {index}"), + TaskStatus::Pending, + Some("worker-2"), + index, + index, + ) + }) + .collect(); + let payload = render_task_board_summary(&agent("worker-1"), &tasks, &default_roster()); + assert_eq!(row_count(&payload), MAX_WAKE_TASK_ROWS); + assert!(payload.contains("Showing 40 of 45 tasks.")); + assert!(payload.contains("Hidden: active=5, completed=0, deleted=0.")); + assert!(payload.contains("| task0044… | Active 44 | pending | worker-2 | - |")); + assert!(!payload.contains("| task0004… | Active 4 | pending | worker-2 | - |")); + } + + #[test] + fn owned_active_tasks_sort_before_other_active_tasks() { + let tasks = vec![ + task( + "other111-1111", + "Other new", + TaskStatus::InProgress, + Some("worker-2"), + 1, + 100, + ), + task( + "owned111-1111", + "Owned old", + TaskStatus::Pending, + Some("worker-1"), + 2, + 1, + ), + ]; + let payload = render_task_board_summary(&agent("worker-1"), &tasks, &default_roster()); + assert!(payload.find("Owned old").unwrap() < payload.find("Other new").unwrap()); + } + + #[test] + fn active_dependency_neighbors_sort_after_owned_tasks() { + let mut owned = task("ownedaaa-1111", "Owned", TaskStatus::InProgress, Some("worker-1"), 1, 1); + owned.blocked_by = vec!["blockera-1111".to_owned()]; + owned.blocks = vec!["downstrm-1111".to_owned()]; + let blocker = task("blockera-1111", "Blocker", TaskStatus::Pending, Some("worker-2"), 2, 20); + let downstream = task( + "downstrm-1111", + "Downstream", + TaskStatus::Pending, + Some("worker-3"), + 3, + 30, + ); + let other = task( + "otheraaa-1111", + "Other", + TaskStatus::InProgress, + Some("worker-4"), + 4, + 40, + ); + let payload = render_task_board_summary( + &agent("worker-1"), + &[other, downstream, blocker, owned], + &default_roster(), + ); + assert!(payload.find("Owned").unwrap() < payload.find("Blocker").unwrap()); + assert!(payload.find("Blocker").unwrap() < payload.find("Downstream").unwrap()); + assert!(payload.find("Downstream").unwrap() < payload.find("Other").unwrap()); + } + + #[test] + fn teammate_summary_keeps_owned_and_related_active_tasks_only() { + let mut owned = task( + "ownedaaa-1111", + "Owned active", + TaskStatus::InProgress, + Some("worker-1"), + 1, + 10, + ); + owned.blocked_by = vec!["blockera-1111".to_owned()]; + owned.blocks = vec!["downstrm-1111".to_owned()]; + let blocker = task( + "blockera-1111", + "Active blocker", + TaskStatus::Pending, + Some("worker-2"), + 2, + 20, + ); + let downstream = task( + "downstrm-1111", + "Active downstream", + TaskStatus::Pending, + Some("worker-3"), + 3, + 30, + ); + let unrelated_active = task( + "unrelact-1111", + "Unrelated active", + TaskStatus::InProgress, + Some("worker-4"), + 4, + 40, + ); + let own_completed = task( + "owndonea-1111", + "Own completed", + TaskStatus::Completed, + Some("worker-1"), + 5, + 50, + ); + let unrelated_completed = task( + "othdonea-1111", + "Other completed", + TaskStatus::Completed, + Some("worker-4"), + 6, + 60, + ); + + let payload = render_task_board_summary( + &teammate("worker-1"), + &[ + unrelated_active, + unrelated_completed, + own_completed, + downstream, + blocker, + owned, + ], + &default_roster(), + ); + + assert!(payload.contains("Owned active")); + assert!(payload.contains("Active blocker")); + assert!(payload.contains("Active downstream")); + assert!(payload.contains("Own completed")); + assert!(!payload.contains("Unrelated active")); + assert!(!payload.contains("Other completed")); + assert!(payload.contains("Hidden: active=1, completed=1, deleted=0.")); + } + + #[test] + fn teammate_recent_completed_rows_are_limited_to_three_owned_tasks() { + let mut tasks = vec![task( + "otherdon-1111", + "Other newest completed", + TaskStatus::Completed, + Some("worker-2"), + 10, + 100, + )]; + tasks.extend((0..4).map(|index| { + task( + &format!("owndone{index}-1111"), + &format!("Own completed {index}"), + TaskStatus::Completed, + Some("worker-1"), + index, + index, + ) + })); + + let payload = render_task_board_summary(&teammate("worker-1"), &tasks, &default_roster()); + + assert_eq!(row_count(&payload), 3); + assert!(payload.contains("Own completed 3")); + assert!(payload.contains("Own completed 1")); + assert!(!payload.contains("Own completed 0")); + assert!(!payload.contains("Other newest completed")); + } + + #[test] + fn reverse_blocked_by_downstream_sorts_after_owned_tasks() { + let owned = task("ownedaaa-1111", "Owned", TaskStatus::InProgress, Some("worker-1"), 1, 1); + let mut downstream = task( + "downstrm-1111", + "Downstream reverse", + TaskStatus::Pending, + Some("worker-3"), + 3, + 30, + ); + downstream.blocked_by = vec!["ownedaaa-1111".to_owned()]; + let other = task( + "otheraaa-1111", + "Other", + TaskStatus::InProgress, + Some("worker-4"), + 4, + 40, + ); + let payload = render_task_board_summary(&agent("worker-1"), &[other, downstream, owned], &default_roster()); + assert!(payload.find("Owned").unwrap() < payload.find("Downstream reverse").unwrap()); + assert!(payload.find("Downstream reverse").unwrap() < payload.find("Other").unwrap()); + } + + #[test] + fn completed_and_deleted_do_not_create_dependency_priority() { + let mut owned = task("ownedaaa-1111", "Owned", TaskStatus::InProgress, Some("worker-1"), 1, 1); + owned.blocked_by = vec!["doneaaaa-1111".to_owned(), "delddddd-1111".to_owned()]; + let completed = task( + "doneaaaa-1111", + "Completed blocker", + TaskStatus::Completed, + Some("worker-2"), + 2, + 200, + ); + let deleted = task( + "delddddd-1111", + "Deleted blocker", + TaskStatus::Deleted, + Some("worker-2"), + 3, + 300, + ); + let other = task( + "otheraaa-1111", + "Other active", + TaskStatus::InProgress, + Some("worker-3"), + 4, + 400, + ); + let payload = render_task_board_summary( + &agent("worker-1"), + &[completed, deleted, other, owned], + &default_roster(), + ); + assert!(payload.find("Owned").unwrap() < payload.find("Other active").unwrap()); + assert!(payload.find("Other active").unwrap() < payload.find("Completed blocker").unwrap()); + assert!(!payload.contains("Deleted blocker")); + } + + #[test] + fn same_priority_sort_is_stable_and_deterministic() { + let tasks = vec![ + task( + "cccccccc-1111", + "Pending newest", + TaskStatus::Pending, + Some("worker-2"), + 1, + 300, + ), + task( + "bbbbbbbb-1111", + "In progress old", + TaskStatus::InProgress, + Some("worker-2"), + 2, + 10, + ), + task( + "aaaaaaaa-1111", + "In progress new", + TaskStatus::InProgress, + Some("worker-2"), + 3, + 20, + ), + ]; + let payload = render_task_board_summary(&agent("worker-1"), &tasks, &default_roster()); + assert!(payload.find("In progress new").unwrap() < payload.find("In progress old").unwrap()); + assert!(payload.find("In progress old").unwrap() < payload.find("Pending newest").unwrap()); + } + + #[test] + fn pending_tasks_with_same_priority_sort_by_updated_at_desc() { + let tasks = vec![ + task( + "aaaaaaaa-1111", + "Pending older update", + TaskStatus::Pending, + Some("worker-2"), + 1, + 10, + ), + task( + "bbbbbbbb-1111", + "Pending newer update", + TaskStatus::Pending, + Some("worker-2"), + 2, + 20, + ), + ]; + let payload = render_task_board_summary(&agent("worker-1"), &tasks, &default_roster()); + assert!(payload.find("Pending newer update").unwrap() < payload.find("Pending older update").unwrap()); + } + + #[test] + fn hidden_counts_are_mutually_exclusive() { + let mut tasks: Vec = (0..42) + .map(|index| { + task( + &format!("act{index:05}-1111"), + &format!("Active {index}"), + TaskStatus::Pending, + Some("worker-2"), + index, + index, + ) + }) + .collect(); + tasks.push(task( + "doneaaaa-1111", + "Completed", + TaskStatus::Completed, + Some("worker-2"), + 100, + 100, + )); + tasks.push(task( + "delddddd-1111", + "Deleted", + TaskStatus::Deleted, + Some("worker-2"), + 101, + 101, + )); + let payload = render_task_board_summary(&agent("worker-1"), &tasks, &default_roster()); + assert!(payload.contains("Showing 40 of 44 tasks.")); + assert!(payload.contains("Hidden: active=2, completed=1, deleted=1.")); + } + + #[test] + fn subject_and_blocked_by_values_are_truncated() { + let mut active = task( + "aaaaaaaa-1111", + &"x".repeat(MAX_SUBJECT_CHARS + 5), + TaskStatus::Pending, + Some("worker-1"), + 1, + 1, + ); + active.blocked_by = vec![ + "block001-1111".to_owned(), + "block002-1111".to_owned(), + "block003-1111".to_owned(), + "block004-1111".to_owned(), + "block005-1111".to_owned(), + "block006-1111".to_owned(), + ]; + let payload = render_task_board_summary(&agent("worker-1"), &[active], &default_roster()); + assert!(payload.contains(&format!("{}...", "x".repeat(MAX_SUBJECT_CHARS)))); + assert!(payload.contains("block001…")); + assert!(payload.contains("+1 more")); + assert!(!payload.contains("block006")); + } + + #[test] + fn task_ids_are_short_but_owner_slot_id_is_full() { + let mut active = task( + "aaaaaaaa-1234-5678", + "Active", + TaskStatus::Pending, + Some("019f4056-fd18-7411-ab09-8868ff17cb36"), + 1, + 1, + ); + active.blocked_by = vec!["bbbbbbbb-1234-5678".to_owned()]; + let payload = render_task_board_summary(&agent("worker-1"), &[active], &default_roster()); + assert!( + payload.contains("| aaaaaaaa… | Active | pending | 019f4056-fd18-7411-ab09-8868ff17cb36 | bbbbbbbb… |") + ); + assert!(!payload.contains("bbbbbbbb-1234-5678")); + } + + #[test] + fn removed_owner_active_task_is_hidden_when_unrelated() { + let removed = task( + "removeda-1111", + "Removed member active", + TaskStatus::InProgress, + Some("worker-2"), + 1, + 10, + ); + let current = task( + "currenta-1111", + "Current member active", + TaskStatus::Pending, + Some("worker-1"), + 2, + 20, + ); + + let payload = + render_task_board_summary(&agent("lead-1"), &[removed, current], &roster(&["lead-1", "worker-1"])); + + assert!(payload.contains("Current member active")); + assert!(!payload.contains("Removed member active")); + assert!(!payload.contains("worker-2 (removed)")); + assert!(payload.contains("Showing 1 of 2 tasks.")); + assert!(payload.contains("Hidden: active=1, completed=0, deleted=0.")); + } + + #[test] + fn removed_owner_active_task_is_shown_when_it_blocks_visible_work() { + let blocker = task( + "blockera-1111", + "Removed member blocker", + TaskStatus::InProgress, + Some("worker-2"), + 1, + 10, + ); + let mut current = task( + "currenta-1111", + "Current member blocked", + TaskStatus::Pending, + Some("worker-1"), + 2, + 20, + ); + current.blocked_by = vec!["blockera-1111".to_owned()]; + + let payload = + render_task_board_summary(&agent("lead-1"), &[blocker, current], &roster(&["lead-1", "worker-1"])); + + assert!(payload.contains("Removed member blocker")); + assert!(payload.contains("Current member blocked")); + assert!(payload.contains("| blockera… | Removed member blocker | in_progress | worker-2 (removed) | - |")); + assert!(payload.contains("| currenta… | Current member blocked | pending | worker-1 | blockera… |")); + assert!(payload.contains("Hidden: active=0, completed=0, deleted=0.")); + } + + #[test] + fn removed_owner_active_task_is_shown_when_it_blocks_ownerless_visible_work() { + let blocker = task( + "blockera-1111", + "Removed owner blocker", + TaskStatus::Pending, + Some("worker-2"), + 1, + 10, + ); + let mut ownerless = task( + "ownerlsa-1111", + "Ownerless blocked task", + TaskStatus::Pending, + None, + 2, + 20, + ); + ownerless.blocked_by = vec!["blockera-1111".to_owned()]; + + let payload = render_task_board_summary( + &agent("lead-1"), + &[blocker, ownerless], + &roster(&["lead-1", "worker-1"]), + ); + + assert!(payload.contains("Removed owner blocker")); + assert!(payload.contains("Ownerless blocked task")); + assert!(payload.contains("worker-2 (removed)")); + } + + #[test] + fn removed_owner_completed_and_deleted_tasks_are_hidden() { + let completed = task( + "doneaaaa-1111", + "Removed completed", + TaskStatus::Completed, + Some("worker-2"), + 1, + 10, + ); + let deleted = task( + "delddddd-1111", + "Removed deleted", + TaskStatus::Deleted, + Some("worker-2"), + 2, + 20, + ); + let current = task( + "currenta-1111", + "Current active", + TaskStatus::Pending, + Some("worker-1"), + 3, + 30, + ); + + let payload = render_task_board_summary( + &agent("lead-1"), + &[completed, deleted, current], + &roster(&["lead-1", "worker-1"]), + ); + + assert!(payload.contains("Current active")); + assert!(!payload.contains("Removed completed")); + assert!(!payload.contains("Removed deleted")); + assert!(payload.contains("Hidden: active=0, completed=1, deleted=1.")); + } +} diff --git a/crates/aionui-team/src/provisioning.rs b/crates/aionui-team/src/provisioning.rs index 9c83c4828..41d53a6a5 100644 --- a/crates/aionui-team/src/provisioning.rs +++ b/crates/aionui-team/src/provisioning.rs @@ -1,22 +1,18 @@ use std::sync::Arc; use aionui_ai_agent::IWorkerTaskManager; -use aionui_api_types::{AddAgentRequest, TeamAgentInput}; +use aionui_api_types::{AddAgentRequest, GetConfigOptionsResponse, TeamAgentInput}; use aionui_common::{AgentKillReason, AgentType, ProviderWithModel, generate_id}; use aionui_db::models::{AgentMetadataRow, TeamRow}; -use aionui_db::{ - IAgentMetadataRepository, IAssistantDefinitionRepository, IAssistantOverlayRepository, IProviderRepository, - ITeamRepository, UpdateTeamParams, -}; +use aionui_db::{IAgentMetadataRepository, IProviderRepository, ITeamRepository, UpdateTeamParams}; use async_trait::async_trait; use tracing::{info, warn}; use crate::error::TeamError; use crate::mcp::TeamMcpStdioConfig; +use crate::ports::TeamAssistantCatalogPort; use crate::service::inherit_team_workspace; -use crate::service::spawn_support::{ - acp_backend_metadata, parse_agent_type, resolve_runtime_backend, session_mode_for_backend, -}; +use crate::service::spawn_support::{acp_backend_metadata, parse_agent_type, session_mode_for_backend}; use crate::types::{Team, TeamAgent, TeammateRole}; use crate::workspace::TeamWorkspaceResolver; @@ -24,8 +20,7 @@ use crate::workspace::TeamWorkspaceResolver; pub struct TeamAgentProvisioner { repo: Arc, agent_metadata_repo: Arc, - assistant_definition_repo: Arc, - assistant_overlay_repo: Arc, + assistant_catalog: Arc, provider_repo: Arc, conversation_port: Arc, } @@ -51,6 +46,7 @@ struct NewAgentProvisioning { model: String, assistant_id: Option, workspace: Option, + session_mode: Option, } pub(crate) struct PersistSpawnedAgentRequest { @@ -95,6 +91,8 @@ pub trait TeamConversationProvisioningPort: Send + Sync { async fn save_acp_runtime_mode(&self, conversation_id: &str, mode: &str) -> Result<(), TeamError>; + async fn get_config_options(&self, conversation_id: &str) -> Result; + async fn warmup_agent_process( &self, user_id: &str, @@ -106,6 +104,11 @@ pub trait TeamConversationProvisioningPort: Send + Sync { } impl TeamAgentProvisioner { + fn normalized_role(input: &TeamAgentInput) -> Result { + TeammateRole::parse(input.role.trim()) + .ok_or_else(|| TeamError::InvalidRequest(format!("invalid team agent role: {}", input.role))) + } + fn effective_assistant_id(assistant_id: Option<&str>) -> Option { assistant_id .map(str::trim) @@ -116,16 +119,14 @@ impl TeamAgentProvisioner { pub(crate) fn new( repo: Arc, agent_metadata_repo: Arc, - assistant_definition_repo: Arc, - assistant_overlay_repo: Arc, + assistant_catalog: Arc, provider_repo: Arc, conversation_port: Arc, ) -> Self { Self { repo, agent_metadata_repo, - assistant_definition_repo, - assistant_overlay_repo, + assistant_catalog, provider_repo, conversation_port, } @@ -142,10 +143,26 @@ impl TeamAgentProvisioner { inputs: &[TeamAgentInput], shared_workspace: Option<&str>, ) -> Result { - let Some((leader_input, teammate_inputs)) = inputs.split_first() else { + if inputs.is_empty() { return Err(TeamError::InvalidRequest("at least one agent is required".into())); }; + let roles = inputs + .iter() + .map(Self::normalized_role) + .collect::, _>>()?; + let leaders = roles + .iter() + .enumerate() + .filter_map(|(idx, role)| (*role == TeammateRole::Lead).then_some(idx)) + .collect::>(); + let [leader_idx] = leaders.as_slice() else { + return Err(TeamError::InvalidRequest( + "exactly one team agent must have role lead".into(), + )); + }; + + let leader_input = &inputs[*leader_idx]; let leader_slot_id = generate_id(); let leader_role = TeammateRole::Lead; let leader_assistant_id = Self::effective_assistant_id(leader_input.assistant_id.as_deref()); @@ -163,6 +180,7 @@ impl TeamAgentProvisioner { &leader_input.model, leader_assistant_id.as_deref(), shared_workspace, + None, ) .await?; @@ -192,9 +210,12 @@ impl TeamAgentProvisioner { cli_path: None, }); - for input in teammate_inputs { + for (input, role) in inputs + .iter() + .zip(roles.iter()) + .filter(|(_, role)| **role == TeammateRole::Teammate) + { let slot_id = generate_id(); - let role = TeammateRole::parse(&input.role).unwrap_or(TeammateRole::Teammate); let assistant_id = Self::effective_assistant_id(input.assistant_id.as_deref()); let backend = self .resolve_requested_backend(input.backend.as_deref(), assistant_id.as_deref()) @@ -204,18 +225,19 @@ impl TeamAgentProvisioner { user_id, team_id, &slot_id, - role, + *role, &input.name, &backend, &input.model, assistant_id.as_deref(), Some(&team_workspace), + None, ) .await?; agents.push(TeamAgent { slot_id, name: input.name.clone(), - role, + role: *role, conversation_id: conversation.conversation_id, backend, model: input.model.clone(), @@ -251,7 +273,13 @@ impl TeamAgentProvisioner { team: &mut Team, req: AddAgentRequest, ) -> Result { - let role = TeammateRole::parse(&req.role).unwrap_or(TeammateRole::Teammate); + let role = TeammateRole::parse(req.role.trim()) + .ok_or_else(|| TeamError::InvalidRequest(format!("invalid team agent role: {}", req.role)))?; + if role != TeammateRole::Teammate { + return Err(TeamError::InvalidRequest( + "add_agent only supports teammate role".into(), + )); + } let workspace = self.workspace_resolver().resolve_for_new_agent(row, team).await?; let assistant_id = Self::effective_assistant_id(req.assistant_id.as_deref()); let backend = self @@ -268,6 +296,7 @@ impl TeamAgentProvisioner { model: req.model, assistant_id, workspace: Some(workspace), + session_mode: row.session_mode.clone(), }) .await?; team.agents.push(agent.clone()); @@ -282,17 +311,14 @@ impl TeamAgentProvisioner { ) -> Result { let assistant_id = assistant_id.map(str::trim).filter(|value| !value.is_empty()); if let Some(assistant_id) = assistant_id { - let definition = self - .assistant_definition_repo - .get_by_assistant_id(assistant_id) + return self + .assistant_catalog + .resolve_team_selectable_assistant(assistant_id) .await? - .ok_or_else(|| TeamError::InvalidRequest(format!("Preset assistant not found: {assistant_id}")))?; - let overlay = self.assistant_overlay_repo.get(&definition.id).await?; - let effective_agent_id = overlay - .and_then(|row| row.agent_id_override) - .filter(|value| !value.trim().is_empty()) - .unwrap_or(definition.agent_id); - return resolve_runtime_backend(&self.agent_metadata_repo, &effective_agent_id).await; + .map(|assistant| assistant.backend) + .ok_or_else(|| { + TeamError::InvalidRequest(format!("Assistant is not available for team mode: {assistant_id}")) + }); } let Some(requested_backend) = requested_backend.map(str::trim).filter(|value| !value.is_empty()) else { @@ -322,6 +348,7 @@ impl TeamAgentProvisioner { model: req.model, assistant_id: req.assistant_id, workspace: Some(workspace), + session_mode: row.session_mode.clone(), }) .await?; team.agents.push(agent.clone()); @@ -412,6 +439,7 @@ impl TeamAgentProvisioner { &input.model, input.assistant_id.as_deref(), input.workspace.as_deref(), + input.session_mode.as_deref(), ) .await?; Ok(TeamAgent { @@ -440,6 +468,7 @@ impl TeamAgentProvisioner { model: &str, assistant_id: Option<&str>, workspace: Option<&str>, + session_mode: Option<&str>, ) -> Result { let acp_metadata = acp_backend_metadata(&self.agent_metadata_repo, backend).await?; let agent_type = if acp_metadata.is_some() { @@ -457,6 +486,7 @@ impl TeamAgentProvisioner { workspace, agent_type, acp_metadata.as_ref(), + session_mode, ); let provider_id = if agent_type == AgentType::Aionrs { self.resolve_provider_for_model(model) @@ -558,8 +588,13 @@ impl TeamAgentProvisioner { workspace: Option<&str>, agent_type: AgentType, acp_metadata: Option<&AgentMetadataRow>, + session_mode: Option<&str>, ) -> serde_json::Value { - let session_mode = session_mode_for_backend(backend, agent_type, acp_metadata); + let session_mode = session_mode + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_owned) + .unwrap_or_else(|| session_mode_for_backend(backend, agent_type, acp_metadata)); let mut extra = serde_json::json!({ "teamId": team_id, "slot_id": slot_id, @@ -614,9 +649,8 @@ mod tests { use aionui_ai_agent::types::BuildTaskOptions; use aionui_ai_agent::{AgentError, AgentInstance}; use aionui_db::models::{ - AgentMetadataRow, AssistantDefinitionRow, AssistantOverlayRow, Provider, UpdateAgentAvailabilitySnapshotParams, - UpdateAgentHandshakeParams, UpsertAgentMetadataParams, UpsertAssistantDefinitionParams, - UpsertAssistantOverlayParams, + AgentMetadataRow, Provider, UpdateAgentAvailabilitySnapshotParams, UpdateAgentHandshakeParams, + UpsertAgentMetadataParams, }; use aionui_db::{CreateProviderParams, DbError, UpdateProviderParams}; use std::sync::Mutex; @@ -660,6 +694,12 @@ mod tests { Ok(()) } + async fn get_config_options(&self, _conversation_id: &str) -> Result { + Ok(GetConfigOptionsResponse { + config_options: Vec::new(), + }) + } + async fn warmup_agent_process( &self, _user_id: &str, @@ -734,6 +774,17 @@ mod tests { struct UnusedAgentMetadataRepo; + struct EmptyTeamAssistantCatalog; + + #[async_trait] + impl TeamAssistantCatalogPort for EmptyTeamAssistantCatalog { + async fn list_team_selectable_assistants( + &self, + ) -> Result, TeamError> { + Ok(Vec::new()) + } + } + #[async_trait] impl IAgentMetadataRepository for UnusedAgentMetadataRepo { async fn list_all(&self) -> Result, DbError> { @@ -785,55 +836,6 @@ mod tests { } } - struct UnusedAssistantDefinitionRepo; - - #[async_trait] - impl IAssistantDefinitionRepository for UnusedAssistantDefinitionRepo { - async fn list(&self) -> Result, DbError> { - Ok(Vec::new()) - } - async fn get_by_assistant_id(&self, _assistant_id: &str) -> Result, DbError> { - Ok(None) - } - async fn get_by_id(&self, _id: &str) -> Result, DbError> { - Ok(None) - } - async fn get_by_source_ref( - &self, - _source: &str, - _source_ref: &str, - ) -> Result, DbError> { - Ok(None) - } - async fn upsert( - &self, - _params: &UpsertAssistantDefinitionParams<'_>, - ) -> Result { - Err(DbError::Init("unused".into())) - } - async fn soft_delete(&self, _id: &str, _deleted_at: i64) -> Result { - Ok(false) - } - } - - struct UnusedAssistantOverlayRepo; - - #[async_trait] - impl IAssistantOverlayRepository for UnusedAssistantOverlayRepo { - async fn get(&self, _assistant_definition_id: &str) -> Result, DbError> { - Ok(None) - } - async fn list(&self) -> Result, DbError> { - Ok(Vec::new()) - } - async fn upsert(&self, _params: &UpsertAssistantOverlayParams<'_>) -> Result { - Err(DbError::Init("unused".into())) - } - async fn delete(&self, _assistant_definition_id: &str) -> Result { - Ok(false) - } - } - struct EmptyProviderRepo; #[async_trait] @@ -859,8 +861,7 @@ mod tests { TeamAgentProvisioner::new( Arc::new(crate::test_utils::MockTeamRepo::new()), Arc::new(UnusedAgentMetadataRepo), - Arc::new(UnusedAssistantDefinitionRepo), - Arc::new(UnusedAssistantOverlayRepo), + Arc::new(EmptyTeamAssistantCatalog), Arc::new(EmptyProviderRepo), Arc::new(RecordingProvisioningPort { events }), ) diff --git a/crates/aionui-team/src/routes.rs b/crates/aionui-team/src/routes.rs index 1a906ff2b..0bab888ff 100644 --- a/crates/aionui-team/src/routes.rs +++ b/crates/aionui-team/src/routes.rs @@ -11,8 +11,9 @@ use axum::routing::{get, post}; use aionui_ai_agent::ActiveLeaseRegistry; use aionui_api_types::{ AddAgentRequest, ApiResponse, CancelTeamChildTurnRequest, CancelTeamRunRequest, CreateTeamRequest, - PauseTeamSlotRequest, RenameAgentRequest, RenameTeamRequest, SendAgentMessageRequest, SendTeamMessageRequest, - SetModeRequest, TeamAgentResponse, TeamListResponse, TeamResponse, TeamRunAckResponse, TeamRunStateResponse, + GetConfigOptionsResponse, PauseTeamSlotRequest, RenameAgentRequest, RenameTeamRequest, SendAgentMessageRequest, + SendTeamMessageRequest, SetModeRequest, TeamAgentResponse, TeamListResponse, TeamResponse, TeamRunAckResponse, + TeamRunStateResponse, }; use aionui_auth::CurrentUser; use aionui_common::ApiError; @@ -50,13 +51,34 @@ impl From for ApiError { ApiError::BadRequest(msg) } } - TeamError::SlotBusy(msg) => ApiError::Conflict(format!("Team slot is busy: {msg}")), TeamError::LeaderOnly(msg) => ApiError::Forbidden(msg), TeamError::Forbidden(msg) => ApiError::Forbidden(msg), TeamError::SessionNotFound(msg) => ApiError::NotFound(msg), TeamError::BlockedTaskNotFound(msg) => ApiError::BadRequest(msg), TeamError::BackendNotAllowed(msg) => ApiError::BadRequest(msg), TeamError::DuplicateAgentName(msg) => ApiError::BadRequest(format!("Agent name already taken: {msg}")), + TeamError::RuntimeNotReady { conversation_id } => ApiError::coded( + StatusCode::CONFLICT, + "TEAM_RUNTIME_NOT_READY", + format!("Team agent runtime is not ready for conversation: {conversation_id}"), + Some(serde_json::json!({ "conversation_id": conversation_id })), + ), + TeamError::MemberRuntimeFailed { + team_id, + slot_id, + conversation_id, + public_reason, + } => ApiError::coded( + StatusCode::CONFLICT, + "TEAM_MEMBER_RUNTIME_FAILED", + "A team member runtime failed to start", + Some(serde_json::json!({ + "team_id": team_id, + "slot_id": slot_id, + "conversation_id": conversation_id, + "reason": public_reason, + })), + ), TeamError::WorkspacePathUnavailable(path) => ApiError::WorkspacePathUnavailable(path), TeamError::WorkspacePathRuntimeUnavailable(path) => ApiError::WorkspacePathRuntimeUnavailable(path), TeamError::Database(db_err) => db_error_to_api_error(db_err), @@ -79,6 +101,10 @@ pub fn team_routes(state: TeamRouterState) -> Router { ) .route("/api/teams/{id}/messages", post(send_message)) .route("/api/teams/{id}/agents/{slot_id}/messages", post(send_message_to_agent)) + .route( + "/api/teams/{id}/conversations/{conversation_id}/config-options", + get(get_conversation_config_options), + ) .route("/api/teams/{id}/runs/{team_run_id}/cancel", post(cancel_run)) .route( "/api/teams/{id}/runs/{team_run_id}/agents/{slot_id}/cancel", @@ -314,6 +340,19 @@ async fn ensure_session( Ok(Json(ApiResponse::success())) } +async fn get_conversation_config_options( + State(state): State, + Extension(user): Extension, + Path((id, conversation_id)): Path<(String, String)>, +) -> Result>, ApiError> { + Ok(Json(ApiResponse::ok( + state + .service + .get_conversation_config_options(&user.id, &id, &conversation_id) + .await?, + ))) +} + async fn stop_session( State(state): State, Extension(user): Extension, @@ -390,12 +429,6 @@ mod tests { ); } - #[test] - fn slot_busy_maps_to_conflict() { - let api_error: ApiError = TeamError::SlotBusy("lead-1".into()).into(); - assert_eq!(api_error.status_code(), StatusCode::CONFLICT); - } - #[test] fn leader_only_maps_to_forbidden() { let err: ApiError = TeamError::LeaderOnly("spawn_agent".into()).into(); @@ -426,6 +459,40 @@ mod tests { assert!(matches!(err, ApiError::BadRequest(msg) if msg.contains("alice"))); } + #[test] + fn runtime_not_ready_maps_to_coded_conflict() { + let err: ApiError = TeamError::RuntimeNotReady { + conversation_id: "conv-1".into(), + } + .into(); + assert_eq!(err.status_code(), StatusCode::CONFLICT); + assert_eq!(err.error_code(), "TEAM_RUNTIME_NOT_READY"); + assert_eq!(err.error_details(), Some(json!({ "conversation_id": "conv-1" }))); + } + + #[test] + fn member_runtime_failure_maps_to_sanitized_coded_conflict() { + let err: ApiError = TeamError::MemberRuntimeFailed { + team_id: "team-1".into(), + slot_id: "slot-2".into(), + conversation_id: "conv-2".into(), + public_reason: "Agent runtime failed to start".into(), + } + .into(); + assert_eq!(err.status_code(), StatusCode::CONFLICT); + assert_eq!(err.error_code(), "TEAM_MEMBER_RUNTIME_FAILED"); + assert_eq!( + err.error_details(), + Some(json!({ + "team_id": "team-1", + "slot_id": "slot-2", + "conversation_id": "conv-2", + "reason": "Agent runtime failed to start", + })) + ); + assert!(!format!("{err:?}").contains("provider-secret")); + } + #[test] fn workspace_error_preserves_code() { let err: ApiError = TeamError::WorkspacePathUnavailable("/tmp/a b".into()).into(); diff --git a/crates/aionui-team/src/scheduler/actions.rs b/crates/aionui-team/src/scheduler/actions.rs index edd540991..26e5ec402 100644 --- a/crates/aionui-team/src/scheduler/actions.rs +++ b/crates/aionui-team/src/scheduler/actions.rs @@ -42,6 +42,39 @@ pub enum SchedulerAction { } impl TeammateManager { + pub async fn create_task( + &self, + subject: &str, + description: Option<&str>, + owner: Option<&str>, + blocked_by: &[String], + ) -> Result { + self.task_board + .create_task(&self.team_id, subject, description, owner, blocked_by) + .await + } + + pub async fn update_task( + &self, + task_id: &str, + status: Option<&str>, + description: Option, + owner: Option, + blocked_by: Option>, + ) -> Result { + use crate::task_board::TaskUpdate; + use crate::types::TaskStatus; + + let update = TaskUpdate { + status: status.and_then(TaskStatus::parse), + description, + owner, + blocked_by, + ..Default::default() + }; + self.task_board.update_task(&self.team_id, task_id, &update).await + } + pub async fn execute_action( &self, from_slot_id: &str, @@ -58,14 +91,7 @@ impl TeammateManager { owner, blocked_by, } => { - self.task_board - .create_task( - &self.team_id, - subject, - description.as_deref(), - owner.as_deref(), - blocked_by, - ) + self.create_task(subject, description.as_deref(), owner.as_deref(), blocked_by) .await?; Ok(None) } @@ -76,17 +102,14 @@ impl TeammateManager { owner, blocked_by, } => { - use crate::task_board::TaskUpdate; - use crate::types::TaskStatus; - - let update = TaskUpdate { - status: status.as_deref().and_then(TaskStatus::parse), - description: description.clone(), - owner: owner.clone(), - blocked_by: blocked_by.clone(), - ..Default::default() - }; - self.task_board.update_task(&self.team_id, task_id, &update).await?; + self.update_task( + task_id, + status.as_deref(), + description.clone(), + owner.clone(), + blocked_by.clone(), + ) + .await?; Ok(None) } SchedulerAction::IdleNotification { summary } => { diff --git a/crates/aionui-team/src/service.rs b/crates/aionui-team/src/service.rs index a2432aaae..858b38b83 100644 --- a/crates/aionui-team/src/service.rs +++ b/crates/aionui-team/src/service.rs @@ -2,15 +2,18 @@ mod describe_support; mod response_builder; pub(crate) mod spawn_support; +use std::collections::HashSet; use std::path::PathBuf; use std::sync::{Arc, Weak}; +use std::time::Instant; -use aionui_ai_agent::{ActiveLeaseRegistry, AgentError, AgentInstance, IWorkerTaskManager}; +use aionui_ai_agent::{ActiveLeaseRegistry, AgentError, AgentInstance, IWorkerTaskManager, IdleCleanupCoordinator}; use aionui_api_types::{ - AddAgentRequest, CreateTeamRequest, TeamAgentResponse, TeamMcpPhase, TeamMcpStatusPayload, TeamResponse, - TeamRunAckResponse, TeamRunStateResponse, TeamRunTargetRole, WebSocketMessage, + AddAgentRequest, CreateTeamRequest, GetConfigOptionsResponse, TeamAgentResponse, TeamAgentRuntimeStatus, + TeamResponse, TeamRunAckResponse, TeamRunStateResponse, TeamSessionPhase, TeamSessionStatus, + TeamSessionStatusPayload, WebSocketMessage, }; -use aionui_common::{AgentKillReason, generate_id, now_ms}; +use aionui_common::{AgentKillReason, ConversationStatus, TimestampMs, generate_id, now_ms}; use aionui_db::models::TeamRow; use aionui_db::{ IAgentMetadataRepository, IAssistantDefinitionRepository, IAssistantOverlayRepository, IProviderRepository, @@ -18,17 +21,27 @@ use aionui_db::{ }; use aionui_realtime::EventBroadcaster; use dashmap::DashMap; +use tokio::task::JoinSet; use tracing::{debug, info, warn}; use crate::error::TeamError; -use crate::event_loop::AgentLoopContext; -use crate::events::{TEAM_CREATED_EVENT, TEAM_MCP_STATUS_EVENT, TEAM_REMOVED_EVENT, TEAM_RENAMED_EVENT}; +use crate::event_loop::{AgentLoopContext, EventLoopRegistrationError}; +use crate::events::{ + TEAM_CREATED_EVENT, TEAM_REMOVED_EVENT, TEAM_RENAMED_EVENT, TEAM_SESSION_STATUS_CHANGED_EVENT, TeamEventEmitter, +}; +use crate::mcp::TeamMcpStdioConfig; +use crate::member_runtime::{ + AttachLease, AttachOutcome, AttachWaiter, BeginRemove, MemberRuntimeFailure, MemberRuntimeSnapshot, ReserveAttach, +}; use crate::message_projection::TeamProjectionMessageStore; use crate::ports::{AgentTurnCancellationPort, AgentTurnExecutionPort, TeamAssistantCatalogPort}; +use crate::prompt_dump::TeamPromptDumpConfig; use crate::provisioning::{TeamAgentProvisioner, TeamConversationProvisioningPort}; -use crate::session::{AgentMessageQueueResult, TeamSession}; -use crate::types::{Team, TeamAgent}; -use crate::wake::TeamWakeSource; +use crate::session::{AgentMessageQueueResult, TeamSession, attach_member_runtime, spawn_attach_agent_process_bg}; +use crate::team_run::TeamRunManager; +use crate::types::{Team, TeamAgent, TeammateRole}; +use crate::work_coordinator::RuntimeConstraint; +use crate::work_source::WorkSource; use crate::workspace::validate_create_workspace_path; pub(crate) fn inherit_team_workspace(extra: &mut serde_json::Value, workspace: &str) { @@ -42,6 +55,122 @@ struct SessionEntry { slow_monitor_handle: tokio::task::JoinHandle<()>, } +pub struct TeamIdleCleanupCoordinator { + service: Arc, + active_leases: Arc, +} + +impl TeamIdleCleanupCoordinator { + pub fn new(service: Arc, active_leases: Arc) -> Self { + Self { service, active_leases } + } +} + +#[async_trait::async_trait] +impl IdleCleanupCoordinator for TeamIdleCleanupCoordinator { + async fn cleanup_idle_conversations( + &self, + idle_conversation_ids: Vec, + idle_threshold_ms: TimestampMs, + ) -> Vec { + self.service + .cleanup_idle_team_runtime_tasks(idle_conversation_ids, &self.active_leases, idle_threshold_ms) + .await + } +} + +struct TeamAgentRebuildOutcome { + agent: TeamAgent, + duration_ms: u128, + result: Result<(), TeamError>, +} + +struct MemberRuntimeReconcileWork { + agent: TeamAgent, + waiter: AttachWaiter, + owner: Option, +} + +const TEAM_REBUILD_MAX_CONCURRENCY: usize = 3; +const TEAM_REBUILD_START_STAGGER: std::time::Duration = std::time::Duration::from_secs(3); + +fn format_rebuild_agent_identity(agent: &TeamAgent) -> String { + format!( + "{} (backend={}, model={}, role={}, slot_id={}, conversation_id={})", + agent.name, agent.backend, agent.model, agent.role, agent.slot_id, agent.conversation_id + ) +} + +fn spawn_rebuild_agent_process( + jobs: &mut JoinSet, + provisioner: TeamAgentProvisioner, + task_manager: Arc, + user_id: String, + agent: TeamAgent, + cfg: TeamMcpStdioConfig, +) { + jobs.spawn(async move { + let team_id = cfg.team_id.clone(); + info!( + team_id = %team_id, + slot_id = %agent.slot_id, + agent_name = %agent.name, + conversation_id = %agent.conversation_id, + backend = %agent.backend, + model = %agent.model, + role = %agent.role, + "team agent rebuild attach started" + ); + let attach_started_at = Instant::now(); + let result = provisioner + .attach_agent_process(&user_id, &agent, cfg, &task_manager) + .await; + let duration_ms = attach_started_at.elapsed().as_millis(); + match &result { + Ok(()) => info!( + team_id = %team_id, + slot_id = %agent.slot_id, + agent_name = %agent.name, + conversation_id = %agent.conversation_id, + backend = %agent.backend, + model = %agent.model, + role = %agent.role, + duration_ms, + "team agent rebuild attach finished" + ), + Err(error) => warn!( + team_id = %team_id, + slot_id = %agent.slot_id, + agent_name = %agent.name, + conversation_id = %agent.conversation_id, + backend = %agent.backend, + model = %agent.model, + role = %agent.role, + duration_ms, + error = %error, + "team agent rebuild attach failed" + ), + } + TeamAgentRebuildOutcome { + agent, + duration_ms, + result, + } + }); +} + +async fn join_next_rebuild_outcome( + jobs: &mut JoinSet, +) -> Result, TeamError> { + match jobs.join_next().await { + Some(Ok(outcome)) => Ok(Some(outcome)), + Some(Err(error)) => Err(TeamError::InvalidRequest(format!( + "team agent rebuild task failed: {error}" + ))), + None => Ok(None), + } +} + pub struct TeamSessionService { repo: Arc, agent_metadata_repo: Arc, @@ -56,14 +185,14 @@ pub struct TeamSessionService { turn_port: Arc, cancellation_port: Arc, backend_binary_path: Arc, + prompt_dump: TeamPromptDumpConfig, sessions: Arc>, - /// Per-team mutex serializing `add_agent` so concurrent callers cannot - /// read-modify-write the `agents` JSON with stale state (last-writer-wins - /// would otherwise drop entries). + /// Per-team mutex serializing membership mutations with session startup so + /// callers cannot read-modify-write the `agents` JSON or rebuild a runtime + /// session from a stale roster snapshot. add_agent_locks: Arc>>>, - /// Per-team mutex serializing `ensure_session` so concurrent callers - /// (e.g. create_team + frontend POST /session) cannot race and start - /// two sessions for the same team. + /// Per-team mutex serializing `ensure_session` so concurrent callers cannot + /// race and start two sessions for the same team. ensure_session_locks: Arc>>>, /// Back-pointer used by [`TeamSession::spawn_agent`] to reach DB-facing /// orchestration without threading the service through every session method. @@ -89,6 +218,41 @@ impl TeamSessionService { turn_port: Arc, cancellation_port: Arc, backend_binary_path: Arc, + ) -> Arc { + Self::new_with_prompt_dump( + repo, + agent_metadata_repo, + assistant_catalog, + assistant_definition_repo, + assistant_overlay_repo, + provider_repo, + conversation_port, + projection_store, + broadcaster, + task_manager, + turn_port, + cancellation_port, + backend_binary_path, + TeamPromptDumpConfig::disabled(), + ) + } + + #[allow(clippy::too_many_arguments)] + pub fn new_with_prompt_dump( + repo: Arc, + agent_metadata_repo: Arc, + assistant_catalog: Arc, + assistant_definition_repo: Arc, + assistant_overlay_repo: Arc, + provider_repo: Arc, + conversation_port: Arc, + projection_store: Arc, + broadcaster: Arc, + task_manager: Arc, + turn_port: Arc, + cancellation_port: Arc, + backend_binary_path: Arc, + prompt_dump: TeamPromptDumpConfig, ) -> Arc { Arc::new_cyclic(|weak| Self { repo, @@ -104,6 +268,7 @@ impl TeamSessionService { turn_port, cancellation_port, backend_binary_path, + prompt_dump, sessions: Arc::new(DashMap::new()), add_agent_locks: Arc::new(DashMap::new()), ensure_session_locks: Arc::new(DashMap::new()), @@ -115,8 +280,7 @@ impl TeamSessionService { TeamAgentProvisioner::new( self.repo.clone(), self.agent_metadata_repo.clone(), - self.assistant_definition_repo.clone(), - self.assistant_overlay_repo.clone(), + self.assistant_catalog.clone(), self.provider_repo.clone(), self.conversation_port.clone(), ) @@ -271,13 +435,6 @@ impl TeamSessionService { self.broadcast_team_created(&team.id, &team.name); - // Auto-start session so MCP is injected immediately after team creation. - // Failure only logs — the team is persisted and frontend can retry - // via POST /api/teams/{id}/session if needed. - if let Err(e) = self.ensure_session_inner(&team.id).await { - warn!(team_id = %team.id, error = %e, "auto ensure_session after create_team failed"); - } - self.build_team_response(&team).await } @@ -386,48 +543,212 @@ impl TeamSessionService { let agent = self.provisioner().add_agent(user_id, &row, &mut team, req).await?; if let Some(session) = self.sessions.get(team_id).map(|e| Arc::clone(&e.session)) { - session.add_agent(&agent).await; - self.register_event_loop(team_id, &agent.slot_id); + let reservation = session.reserve_dynamic_member_attach(&agent); + session.add_manual_agent(&agent).await?; + let service = self + .self_ref + .upgrade() + .ok_or_else(|| TeamError::InvalidRequest("add_agent requires a live TeamSessionService".into()))?; + self.broadcast_agent_runtime_status(team_id, &agent, TeamAgentRuntimeStatus::Pending, None); + spawn_attach_agent_process_bg( + service, + session, + user_id.to_owned(), + agent.clone(), + self.task_manager.clone(), + reservation, + ); + info!( + team_id = %team_id, + slot_id = %agent.slot_id, + assistant_id = %agent.assistant_id.as_deref().unwrap_or(""), + role = %agent.role, + notification_written = true, + wake_requested = true, + "manual teammate added" + ); + } else { + TeamEventEmitter::new(team_id.to_owned(), self.broadcaster.clone()).broadcast_agent_spawned(&agent); + info!( + team_id = %team_id, + slot_id = %agent.slot_id, + assistant_id = %agent.assistant_id.as_deref().unwrap_or(""), + role = %agent.role, + notification_written = false, + wake_requested = false, + "manual teammate added" + ); } self.build_agent_response(&agent).await } pub async fn remove_agent(&self, user_id: &str, team_id: &str, slot_id: &str) -> Result<(), TeamError> { - let mut team = self.load_owned_team(user_id, team_id).await?; + let lock = self + .add_agent_locks + .entry(team_id.to_owned()) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone(); + let (removed, session, removal_lease) = { + let _guard = lock.lock().await; + let team = self.load_owned_team(user_id, team_id).await?; + let removed = team + .agents + .iter() + .find(|agent| agent.slot_id == slot_id) + .cloned() + .ok_or_else(|| TeamError::AgentNotFound(slot_id.into()))?; + if removed.role == crate::types::TeammateRole::Lead { + return Err(TeamError::InvalidRequest("cannot remove the team lead".into())); + } + let session = self.sessions.get(team_id).map(|entry| Arc::clone(&entry.session)); + let removal = session + .as_ref() + .map(|session| session.member_runtimes().begin_remove(slot_id)); + let removal_lease = match removal { + Some(BeginRemove::Start(lease)) => Some(lease), + Some(BeginRemove::Join(waiter)) => { + drop(_guard); + return match waiter.wait().await { + AttachOutcome::Removed => Ok(()), + AttachOutcome::Failed(failure) => Err(TeamError::MemberRuntimeFailed { + team_id: team_id.to_owned(), + slot_id: removed.slot_id, + conversation_id: removed.conversation_id, + public_reason: failure.public_reason, + }), + AttachOutcome::Ready | AttachOutcome::SessionStopped => { + Err(TeamError::SessionNotFound(team_id.to_owned())) + } + }; + } + Some(BeginRemove::Absent | BeginRemove::SessionStopped) | None => None, + }; + (removed, session, removal_lease) + }; - let idx = team - .agents - .iter() - .position(|a| a.slot_id == slot_id) - .ok_or_else(|| TeamError::AgentNotFound(slot_id.into()))?; + // Cancellation and process cleanup intentionally happen without the + // membership lock. Concurrent ensure calls observe Removing and join + // the same registry operation instead of starting a replacement. + if let Some(session) = &session { + session.event_loops().remove(slot_id); + } + self.task_manager + .kill_and_wait(&removed.conversation_id, Some(AgentKillReason::TeamDeleted)) + .await; + + let persist_result = { + let _guard = lock.lock().await; + let mut current = self.load_owned_team(user_id, team_id).await?; + current.agents.retain(|agent| agent.slot_id != slot_id); + let agents_json = serde_json::to_string(¤t.agents)?; + self.repo + .update_team( + team_id, + &UpdateTeamParams { + agents: Some(agents_json), + ..Default::default() + }, + ) + .await + }; + + if let Err(error) = persist_result { + if let (Some(session), Some(lease)) = (&session, removal_lease.as_ref()) { + session + .member_runtimes() + .restore_attach_required_after_remove_persist_error( + lease, + MemberRuntimeFailure { + classification: "membership_persist_failed", + public_reason: "Agent runtime needs to restart after membership update failed".to_owned(), + }, + ); + self.refresh_member_runtime_status(session).await; + } + return Err(error.into()); + } - let removed = team.agents.remove(idx); + let published_session = self.sessions.get(team_id).map(|entry| Arc::clone(&entry.session)); + let active_session = if let Some(current) = published_session { + let current_removal_lease = if session.as_ref().is_some_and(|captured| Arc::ptr_eq(captured, ¤t)) { + removal_lease + } else { + match current.member_runtimes().begin_remove(slot_id) { + BeginRemove::Start(lease) => Some(lease), + BeginRemove::Join(waiter) => { + let _ = waiter.wait().await; + None + } + BeginRemove::Absent | BeginRemove::SessionStopped => None, + } + }; + current.event_loops().remove(slot_id); + self.task_manager + .kill_and_wait(&removed.conversation_id, Some(AgentKillReason::TeamDeleted)) + .await; + match current.scheduler().remove_agent(slot_id).await { + Ok(_) | Err(TeamError::AgentNotFound(_)) => {} + Err(error) => return Err(error), + } + if let Some(lease) = current_removal_lease.as_ref() { + current.member_runtimes().finish_remove(lease); + } + Some(current) + } else { + None + }; - let _ = self + if let Err(error) = self .conversation_port .delete_team_conversation(user_id, &removed.conversation_id) - .await; - - let agents_json = serde_json::to_string(&team.agents)?; - self.repo - .update_team( + .await + { + warn!( team_id, - &UpdateTeamParams { - agents: Some(agents_json), - ..Default::default() - }, - ) - .await?; + slot_id, + conversation_id = %removed.conversation_id, + error = %error, + "removed team member conversation cleanup failed" + ); + } - if let Some(session) = self.sessions.get(team_id).map(|e| Arc::clone(&e.session)) { - let _ = session.remove_agent(slot_id).await; + if let Some(session) = active_session.filter(|session| self.capture_published_session(session).is_some()) { + session.notify_leader_membership_removed(&removed).await?; + self.refresh_member_runtime_status(&session).await; + info!( + team_id = %team_id, + slot_id = %removed.slot_id, + assistant_id = %removed.assistant_id.as_deref().unwrap_or(""), + role = %removed.role, + notification_written = true, + wake_requested = true, + "manual teammate removed" + ); + } else { + TeamEventEmitter::new(team_id.to_owned(), self.broadcaster.clone()).broadcast_agent_removed(slot_id); + info!( + team_id = %team_id, + slot_id = %removed.slot_id, + assistant_id = %removed.assistant_id.as_deref().unwrap_or(""), + role = %removed.role, + notification_written = false, + wake_requested = false, + "manual teammate removed" + ); } Ok(()) } pub async fn rename_agent(&self, user_id: &str, team_id: &str, slot_id: &str, name: &str) -> Result<(), TeamError> { + let lock = self + .add_agent_locks + .entry(team_id.to_owned()) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone(); + let _guard = lock.lock().await; + let mut team = self.load_owned_team(user_id, team_id).await?; let normalized = crate::scheduler::normalize_name(name); @@ -499,34 +820,35 @@ impl TeamSessionService { } async fn ensure_session_inner(&self, team_id: &str) -> Result<(), TeamError> { - if self.sessions.contains_key(team_id) { - return Ok(()); - } - - let lock = self - .ensure_session_locks + let membership_lock = self + .add_agent_locks .entry(team_id.to_owned()) .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) .clone(); - let _guard = lock.lock().await; - - // Re-check after acquiring lock (another caller may have completed). - if self.sessions.contains_key(team_id) { - return Ok(()); - } + let membership_guard = membership_lock.lock().await; let row = match self.repo.get_team(team_id).await { Ok(Some(row)) => row, Ok(None) => { - self.broadcast_mcp_phase(team_id, "", TeamMcpPhase::LoadFailed, None, |p| { - p.error = Some(format!("team not found: {team_id}")); - }); + self.broadcast_session_status( + team_id, + TeamSessionStatus::Failed, + Some(TeamSessionPhase::LoadingTeam), + |p| { + p.error = Some(format!("team not found: {team_id}")); + }, + ); return Err(TeamError::TeamNotFound(team_id.into())); } Err(e) => { - self.broadcast_mcp_phase(team_id, "", TeamMcpPhase::LoadFailed, None, |p| { - p.error = Some(e.to_string()); - }); + self.broadcast_session_status( + team_id, + TeamSessionStatus::Failed, + Some(TeamSessionPhase::LoadingTeam), + |p| { + p.error = Some(e.to_string()); + }, + ); return Err(e.into()); } }; @@ -534,7 +856,49 @@ impl TeamSessionService { let team = Team::from_row(&row)?; let agents_snapshot: Vec = team.agents.clone(); - let session = match TeamSession::start( + if let Some(session) = self.sessions.get(team_id).map(|entry| Arc::clone(&entry.session)) { + let work = self + .reserve_member_runtime_reconciliation(&session, &agents_snapshot) + .await?; + drop(membership_guard); + return self + .complete_member_runtime_reconciliation(team_id, &user_id, session, work) + .await; + } + + let lock = self + .ensure_session_locks + .entry(team_id.to_owned()) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone(); + let ensure_guard = lock.lock().await; + + if let Some(session) = self.sessions.get(team_id).map(|entry| Arc::clone(&entry.session)) { + let work = self + .reserve_member_runtime_reconciliation(&session, &agents_snapshot) + .await?; + drop(membership_guard); + drop(ensure_guard); + return self + .complete_member_runtime_reconciliation(team_id, &user_id, session, work) + .await; + } + + self.broadcast_session_status( + team_id, + TeamSessionStatus::Starting, + Some(TeamSessionPhase::LoadingTeam), + |_| {}, + ); + + self.broadcast_session_status( + team_id, + TeamSessionStatus::Starting, + Some(TeamSessionPhase::StartingBridge), + |_| {}, + ); + + let session = match TeamSession::start_with_prompt_dump( team, self.repo.clone(), self.broadcaster.clone(), @@ -545,32 +909,80 @@ impl TeamSessionService { self.projection_store.clone(), user_id.clone(), self.self_ref.clone(), + self.prompt_dump.clone(), ) .await { - Ok(session) => session, + Ok(session) => Arc::new(session), Err(e) => { - self.broadcast_mcp_phase(team_id, "", TeamMcpPhase::SessionError, None, |p| { - p.error = Some(e.to_string()); - }); + self.broadcast_session_status( + team_id, + TeamSessionStatus::Failed, + Some(TeamSessionPhase::StartingBridge), + |p| { + p.error = Some(e.to_string()); + }, + ); return Err(e); } }; - self.broadcast_mcp_phase(team_id, "", TeamMcpPhase::SessionInjecting, None, |_| {}); + self.broadcast_session_status( + team_id, + TeamSessionStatus::Starting, + Some(TeamSessionPhase::AttachingAgents), + |_| {}, + ); if let Err(e) = self .rebuild_agent_processes(team_id, &session, &user_id, &agents_snapshot) .await { + self.broadcast_session_status( + team_id, + TeamSessionStatus::Failed, + Some(TeamSessionPhase::AttachingAgents), + |p| { + p.error = Some(e.to_string()); + }, + ); session.stop(); + self.cleanup_bootstrap_runtime_tasks(&agents_snapshot).await; return Err(e); } - let session = Arc::new(session); + if let Err(error) = self.spawn_event_loops(&session, &user_id, &agents_snapshot) { + warn!( + team_id, + generation = session.generation(), + error = %error, + "team bootstrap event loop registration failed" + ); + session.stop(); + self.cleanup_bootstrap_runtime_tasks(&agents_snapshot).await; + self.broadcast_session_status( + team_id, + TeamSessionStatus::Failed, + Some(TeamSessionPhase::AttachingAgents), + |payload| payload.error = Some(error.to_string()), + ); + return Err(error); + } - // Spawn per-agent event loops - self.spawn_event_loops(&session, &user_id, &agents_snapshot); + if agents_snapshot + .iter() + .any(|agent| !session.member_runtimes().seed_ready(agent.slot_id.clone())) + { + let error = TeamError::InvalidRequest("failed to seed bootstrap member runtime registry".to_owned()); + session.stop(); + self.cleanup_bootstrap_runtime_tasks(&agents_snapshot).await; + return Err(error); + } + for agent in &agents_snapshot { + session + .work_coordinator() + .set_runtime_constraint(&agent.slot_id, RuntimeConstraint::Ready); + } let slow_monitor_handle = Self::spawn_slow_monitor(session.clone()); let entry = SessionEntry { @@ -578,6 +990,18 @@ impl TeamSessionService { slow_monitor_handle, }; self.sessions.insert(team_id.to_owned(), entry); + drop(membership_guard); + + for agent in &agents_snapshot { + self.broadcast_agent_runtime_status(team_id, agent, TeamAgentRuntimeStatus::Ready, None); + } + + self.broadcast_session_status( + team_id, + TeamSessionStatus::Starting, + Some(TeamSessionPhase::Recovering), + |_| {}, + ); if let Err(err) = session.try_start_recovery_drain("ensure_session_ready").await { warn!( @@ -587,114 +1011,484 @@ impl TeamSessionService { ); } - self.broadcast_mcp_phase(team_id, "", TeamMcpPhase::SessionReady, None, |p| { + self.broadcast_session_status(team_id, TeamSessionStatus::Ready, None, |p| { p.server_count = Some(agents_snapshot.len()); }); Ok(()) } - fn broadcast_mcp_phase(&self, team_id: &str, slot_id: &str, phase: TeamMcpPhase, port: Option, customize: F) - where - F: FnOnce(&mut TeamMcpStatusPayload), - { - let mut payload = TeamMcpStatusPayload { - team_id: team_id.to_owned(), - slot_id: slot_id.to_owned(), - phase, - port, - server_count: None, - error: None, - }; - customize(&mut payload); - let event = WebSocketMessage::new( - TEAM_MCP_STATUS_EVENT, - serde_json::to_value(payload).expect("serialize mcp status payload"), - ); - self.broadcaster.broadcast(event); - } + async fn reserve_member_runtime_reconciliation( + &self, + session: &Arc, + agents: &[TeamAgent], + ) -> Result, TeamError> { + let scheduler_slots = session + .scheduler() + .list_agents() + .await + .into_iter() + .map(|agent| agent.slot_id) + .collect::>(); + let mut work = Vec::new(); - fn spawn_slow_monitor(session: Arc) -> tokio::task::JoinHandle<()> { - tokio::spawn(async move { - let mut interval = tokio::time::interval(std::time::Duration::from_secs(60)); - interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - loop { - interval.tick().await; - session.team_run_manager().observe_slow_child_turns(now_ms()).await; + for agent in agents { + if !scheduler_slots.contains(&agent.slot_id) { + session.scheduler().add_agent(agent).await; } - }) - } - - fn broadcast_team_created(&self, team_id: &str, team_name: &str) { - info!(team_id = %team_id, event_name = TEAM_CREATED_EVENT, "team event broadcast"); - self.broadcaster.broadcast(WebSocketMessage::new( - TEAM_CREATED_EVENT, - serde_json::json!({ "team_id": team_id, "team_name": team_name }), - )); - self.broadcast_team_list_changed(team_id, "created"); - } - - fn broadcast_team_removed(&self, team_id: &str) { - info!(team_id = %team_id, event_name = TEAM_REMOVED_EVENT, "team event broadcast"); - self.broadcaster.broadcast(WebSocketMessage::new( - TEAM_REMOVED_EVENT, - serde_json::json!({ "team_id": team_id }), - )); - self.broadcast_team_list_changed(team_id, "removed"); - } - - fn broadcast_team_renamed(&self, team_id: &str, team_name: &str) { - info!(team_id = %team_id, event_name = TEAM_RENAMED_EVENT, "team event broadcast"); - self.broadcaster.broadcast(WebSocketMessage::new( - TEAM_RENAMED_EVENT, - serde_json::json!({ "team_id": team_id, "team_name": team_name }), - )); - self.broadcast_team_list_changed(team_id, "renamed"); - } + let reservation = match session.member_runtimes().snapshot(&agent.slot_id) { + MemberRuntimeSnapshot::Ready if self.task_manager.get_task(&agent.conversation_id).is_none() => { + session.member_runtimes().reserve_repair(&agent.slot_id) + } + MemberRuntimeSnapshot::Ready => ReserveAttach::AlreadyReady, + _ => session.member_runtimes().reserve_attach(&agent.slot_id, true), + }; - fn broadcast_team_list_changed(&self, team_id: &str, action: &str) { - info!(team_id = %team_id, event_name = crate::events::TEAM_LIST_CHANGED_EVENT, action, "team event broadcast"); - self.broadcaster.broadcast(WebSocketMessage::new( - crate::events::TEAM_LIST_CHANGED_EVENT, - serde_json::json!({ "team_id": team_id, "action": action }), - )); + match reservation { + ReserveAttach::Start(owner) => { + self.broadcast_agent_runtime_status( + session.team_id(), + agent, + TeamAgentRuntimeStatus::Pending, + None, + ); + work.push(MemberRuntimeReconcileWork { + agent: agent.clone(), + waiter: owner.waiter(), + owner: Some(owner), + }); + } + ReserveAttach::Join(waiter) | ReserveAttach::Removing(waiter) => { + info!( + team_id = session.team_id(), + slot_id = agent.slot_id, + conversation_id = agent.conversation_id, + operation_id = waiter.operation_id(), + generation = session.generation(), + duration_ms = 0, + error_classification = "none", + "team member runtime reconciliation waiting" + ); + work.push(MemberRuntimeReconcileWork { + agent: agent.clone(), + waiter, + owner: None, + }); + } + ReserveAttach::AlreadyReady => {} + ReserveAttach::SessionStopped => { + return Err(TeamError::InvalidRequest( + "team session stopped during reconciliation reservation".to_owned(), + )); + } + } + } + Ok(work) } - async fn rebuild_agent_processes( + async fn complete_member_runtime_reconciliation( &self, team_id: &str, - session: &TeamSession, user_id: &str, - agents: &[TeamAgent], + session: Arc, + work: Vec, ) -> Result<(), TeamError> { - let provisioner = self.provisioner(); - for agent in agents { - let cfg = session.mcp_stdio_config(&agent.slot_id); + if !work.is_empty() { + self.publish_member_runtime_starting_if_current(&session); + } + let mut waiters = Vec::with_capacity(work.len()); + for item in work { + let waiter = item.waiter; + if let Some(owner) = item.owner { + tokio::spawn(attach_member_runtime( + self.self_ref + .upgrade() + .ok_or_else(|| TeamError::InvalidRequest("team service is shutting down".to_owned()))?, + Arc::clone(&session), + user_id.to_owned(), + item.agent.clone(), + self.task_manager.clone(), + owner, + )); + } + waiters.push((item.agent, waiter)); + } - if let Err(e) = provisioner - .attach_agent_process(user_id, agent, cfg, &self.task_manager) - .await - { - let msg = format!("failed to attach rebuilt agent {}: {e}", agent.slot_id); - self.broadcast_mcp_phase(team_id, &agent.slot_id, TeamMcpPhase::SessionError, None, |p| { - p.error = Some(msg.clone()); - }); - warn!( - team_id, - slot_id = %agent.slot_id, - conversation_id = %agent.conversation_id, - error = %e, - "warmup failed during rebuild" - ); - return Err(TeamError::InvalidRequest(msg)); + let outcomes = futures_util::future::join_all( + waiters + .into_iter() + .map(|(agent, waiter)| async move { (agent, waiter.wait().await) }), + ) + .await; + + let membership_lock = self + .add_agent_locks + .entry(team_id.to_owned()) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone(); + let _membership_guard = membership_lock.lock().await; + let current_agents = match self.repo.get_team(team_id).await? { + Some(row) => Team::from_row(&row)?.agents, + None => return Err(TeamError::TeamNotFound(team_id.to_owned())), + }; + let current_slots = current_agents + .iter() + .map(|agent| agent.slot_id.as_str()) + .collect::>(); + let current_session = self + .sessions + .get(team_id) + .ok_or_else(|| TeamError::SessionNotFound(team_id.to_owned()))?; + if !Arc::ptr_eq(¤t_session.session, &session) { + return Err(TeamError::SessionNotFound(team_id.to_owned())); + } + + for (agent, outcome) in outcomes { + if !current_slots.contains(agent.slot_id.as_str()) { + continue; + } + match outcome { + AttachOutcome::Ready | AttachOutcome::Removed => {} + AttachOutcome::Failed(failure) => { + self.broadcast_session_status( + team_id, + TeamSessionStatus::Failed, + Some(TeamSessionPhase::AttachingAgents), + |payload| payload.error = Some(failure.public_reason.clone()), + ); + return Err(TeamError::MemberRuntimeFailed { + team_id: team_id.to_owned(), + slot_id: agent.slot_id, + conversation_id: agent.conversation_id, + public_reason: failure.public_reason, + }); + } + AttachOutcome::SessionStopped => { + return Err(TeamError::InvalidRequest( + "team session stopped during reconciliation".to_owned(), + )); + } } } + + self.broadcast_session_status(team_id, TeamSessionStatus::Ready, None, |payload| { + payload.server_count = Some(current_agents.len()); + }); Ok(()) } + async fn cleanup_bootstrap_runtime_tasks(&self, agents: &[TeamAgent]) { + for agent in agents { + self.task_manager + .kill_and_wait(&agent.conversation_id, Some(AgentKillReason::TeamMcpRebuild)) + .await; + } + } + + pub(crate) async fn cleanup_stale_member_runtime_task( + &self, + captured_session: &TeamSession, + conversation_id: &str, + ) { + let lock = self + .ensure_session_locks + .entry(captured_session.team_id().to_owned()) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone(); + let _guard = lock.lock().await; + if self + .sessions + .get(captured_session.team_id()) + .is_some_and(|entry| !std::ptr::eq(entry.session.as_ref(), captured_session)) + { + return; + } + self.task_manager + .kill_and_wait(conversation_id, Some(AgentKillReason::TeamMcpRebuild)) + .await; + } + + pub async fn get_conversation_config_options( + &self, + user_id: &str, + team_id: &str, + conversation_id: &str, + ) -> Result { + let row = self + .repo + .get_team(team_id) + .await? + .ok_or_else(|| TeamError::TeamNotFound(team_id.to_owned()))?; + if row.user_id != user_id { + return Err(TeamError::Forbidden(format!( + "team {team_id} is not owned by current user" + ))); + } + + let team = Team::from_row(&row)?; + let member = team.agents.iter().any(|agent| agent.conversation_id == conversation_id); + if !member { + return Err(TeamError::AgentNotFound(conversation_id.to_owned())); + } + + self.conversation_port.get_config_options(conversation_id).await + } + + fn broadcast_session_status( + &self, + team_id: &str, + status: TeamSessionStatus, + phase: Option, + customize: F, + ) where + F: FnOnce(&mut TeamSessionStatusPayload), + { + let mut payload = TeamSessionStatusPayload { + team_id: team_id.to_owned(), + status, + phase, + server_count: None, + error: None, + }; + customize(&mut payload); + let event = WebSocketMessage::new( + TEAM_SESSION_STATUS_CHANGED_EVENT, + serde_json::to_value(payload).expect("serialize team session status payload"), + ); + self.broadcaster.broadcast(event); + } + + fn spawn_slow_monitor(session: Arc) -> tokio::task::JoinHandle<()> { + tokio::spawn(async move { + let mut interval = tokio::time::interval(std::time::Duration::from_secs(60)); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + interval.tick().await; + let snapshot = session.work_coordinator().snapshot(); + session.team_run_manager().publish_snapshot_update(&snapshot); + } + }) + } + + fn broadcast_team_created(&self, team_id: &str, team_name: &str) { + info!(team_id = %team_id, event_name = TEAM_CREATED_EVENT, "team event broadcast"); + self.broadcaster.broadcast(WebSocketMessage::new( + TEAM_CREATED_EVENT, + serde_json::json!({ "team_id": team_id, "team_name": team_name }), + )); + self.broadcast_team_list_changed(team_id, "created"); + } + + fn broadcast_team_removed(&self, team_id: &str) { + info!(team_id = %team_id, event_name = TEAM_REMOVED_EVENT, "team event broadcast"); + self.broadcaster.broadcast(WebSocketMessage::new( + TEAM_REMOVED_EVENT, + serde_json::json!({ "team_id": team_id }), + )); + self.broadcast_team_list_changed(team_id, "removed"); + } + + fn broadcast_team_renamed(&self, team_id: &str, team_name: &str) { + info!(team_id = %team_id, event_name = TEAM_RENAMED_EVENT, "team event broadcast"); + self.broadcaster.broadcast(WebSocketMessage::new( + TEAM_RENAMED_EVENT, + serde_json::json!({ "team_id": team_id, "team_name": team_name }), + )); + self.broadcast_team_list_changed(team_id, "renamed"); + } + + fn broadcast_team_list_changed(&self, team_id: &str, action: &str) { + info!(team_id = %team_id, event_name = crate::events::TEAM_LIST_CHANGED_EVENT, action, "team event broadcast"); + self.broadcaster.broadcast(WebSocketMessage::new( + crate::events::TEAM_LIST_CHANGED_EVENT, + serde_json::json!({ "team_id": team_id, "action": action }), + )); + } + + pub(crate) fn broadcast_agent_runtime_status( + &self, + team_id: &str, + agent: &TeamAgent, + status: TeamAgentRuntimeStatus, + error: Option, + ) { + TeamEventEmitter::new(team_id.to_owned(), self.broadcaster.clone()) + .broadcast_agent_runtime_status(agent, status, error); + } + + async fn rebuild_agent_processes( + &self, + team_id: &str, + session: &TeamSession, + user_id: &str, + agents: &[TeamAgent], + ) -> Result<(), TeamError> { + let provisioner = self.provisioner(); + let task_manager = self.task_manager.clone(); + let started_at = Instant::now(); + let mut rebuild_jobs: Vec = agents.to_vec(); + rebuild_jobs.sort_by_key(|agent| match agent.role { + TeammateRole::Lead => 0, + TeammateRole::Teammate => 1, + }); + + info!( + team_id, + agent_count = agents.len(), + max_concurrency = TEAM_REBUILD_MAX_CONCURRENCY, + start_stagger_ms = TEAM_REBUILD_START_STAGGER.as_millis(), + "team agent rebuild started" + ); + + let mut outcomes = Vec::new(); + let mut jobs = JoinSet::new(); + let mut failed = false; + + for (launched_count, agent) in rebuild_jobs.into_iter().enumerate() { + while jobs.len() >= TEAM_REBUILD_MAX_CONCURRENCY { + if let Some(outcome) = join_next_rebuild_outcome(&mut jobs).await? { + failed = outcome.result.is_err(); + outcomes.push(outcome); + } + if failed { + break; + } + } + if failed { + break; + } + + if launched_count > 0 { + let stagger = tokio::time::sleep(TEAM_REBUILD_START_STAGGER); + tokio::pin!(stagger); + loop { + tokio::select! { + _ = &mut stagger => break, + outcome = join_next_rebuild_outcome(&mut jobs), if !jobs.is_empty() => { + if let Some(outcome) = outcome? { + failed = outcome.result.is_err(); + outcomes.push(outcome); + } + if failed { + break; + } + } + } + } + if failed { + break; + } + } + + let cfg = session.mcp_stdio_config(&agent.slot_id); + self.broadcast_agent_runtime_status(team_id, &agent, TeamAgentRuntimeStatus::Pending, None); + spawn_rebuild_agent_process( + &mut jobs, + provisioner.clone(), + task_manager.clone(), + user_id.to_owned(), + agent, + cfg, + ); + } + + while let Some(outcome) = join_next_rebuild_outcome(&mut jobs).await? { + outcomes.push(outcome); + } + + let mut success_count = 0usize; + let mut failures: Vec<&TeamAgentRebuildOutcome> = Vec::new(); + for outcome in &outcomes { + match &outcome.result { + Ok(()) => success_count += 1, + Err(_) => failures.push(outcome), + } + } + + info!( + team_id, + agent_count = agents.len(), + success_count, + failure_count = failures.len(), + duration_ms = started_at.elapsed().as_millis(), + max_concurrency = TEAM_REBUILD_MAX_CONCURRENCY, + start_stagger_ms = TEAM_REBUILD_START_STAGGER.as_millis(), + "team agent rebuild completed" + ); + + if failures.is_empty() { + return Ok(()); + } + + let first_error = failures + .first() + .map(|outcome| { + let error = outcome + .result + .as_ref() + .err() + .map(ToString::to_string) + .unwrap_or_else(|| "unknown rebuild failure".to_owned()); + format!("{}: {error}", format_rebuild_agent_identity(&outcome.agent)) + }) + .unwrap_or_else(|| "unknown rebuild failure".to_owned()); + + for failure in &failures { + let error = failure + .result + .as_ref() + .err() + .map(ToString::to_string) + .unwrap_or_else(|| "unknown rebuild failure".to_owned()); + self.broadcast_agent_runtime_status( + team_id, + &failure.agent, + TeamAgentRuntimeStatus::Failed, + Some(error.clone()), + ); + warn!( + team_id, + slot_id = %failure.agent.slot_id, + agent_name = %failure.agent.name, + conversation_id = %failure.agent.conversation_id, + backend = %failure.agent.backend, + model = %failure.agent.model, + role = %failure.agent.role, + duration_ms = failure.duration_ms, + error = %error, + "warmup failed during rebuild" + ); + } + + for success in outcomes.iter().filter(|outcome| outcome.result.is_ok()) { + info!( + team_id, + slot_id = %success.agent.slot_id, + agent_name = %success.agent.name, + conversation_id = %success.agent.conversation_id, + backend = %success.agent.backend, + model = %success.agent.model, + role = %success.agent.role, + "cleaning up successfully attached agent after rebuild failure" + ); + self.task_manager + .kill_and_wait(&success.agent.conversation_id, Some(AgentKillReason::TeamMcpRebuild)) + .await; + } + + Err(TeamError::InvalidRequest(format!( + "failed to attach rebuilt agent: {first_error}" + ))) + } + /// Spawn per-agent event loops that drain the mailbox whenever notified. /// Each agent gets its own tokio task that runs until the session shuts down. - fn spawn_event_loops(&self, session: &Arc, user_id: &str, agents: &[TeamAgent]) { + fn spawn_event_loops( + &self, + session: &Arc, + user_id: &str, + agents: &[TeamAgent], + ) -> Result<(), TeamError> { let registry = session.event_loops(); for agent in agents { @@ -708,8 +1502,14 @@ impl TeamSessionService { turn_port: self.turn_port.clone(), registry: registry.clone(), }; - let _ = registry.spawn(&agent.slot_id, ctx); + registry.spawn(&agent.slot_id, ctx).map_err(|error| { + TeamError::InvalidRequest(format!( + "failed to register bootstrap event loop for slot {}: {error:?}", + agent.slot_id + )) + })?; } + Ok(()) } /// Register an event loop for a dynamically spawned agent. @@ -717,15 +1517,15 @@ impl TeamSessionService { /// Called by [`TeamSession::spawn_agent`] after `attach_spawned_agent_process` /// succeeds so the newly booted agent gets its own drain loop — exactly as /// `spawn_event_loops` does for the initial members during `ensure_session`. - pub(crate) fn register_event_loop(&self, team_id: &str, slot_id: &str) { - let Some(entry) = self.sessions.get(team_id) else { - return; - }; - let session = Arc::clone(&entry.session); + pub(crate) fn register_event_loop( + &self, + session: &Arc, + slot_id: &str, + ) -> Result { let registry = session.event_loops(); let ctx = AgentLoopContext { - team_id: team_id.to_owned(), + team_id: session.team_id().to_owned(), slot_id: slot_id.to_owned(), user_id: session.user_id().to_owned(), session: session.clone(), @@ -734,9 +1534,18 @@ impl TeamSessionService { turn_port: self.turn_port.clone(), registry: registry.clone(), }; - let registered = registry.spawn(slot_id, ctx); - if registered { - info!(team_id, slot_id, "agent event loop registered"); + match registry.spawn(slot_id, ctx) { + Ok(()) => { + info!( + team_id = session.team_id(), + slot_id, + generation = session.generation(), + "agent event loop registered" + ); + Ok(true) + } + Err(EventLoopRegistrationError::Duplicate) => Ok(false), + Err(error) => Err(error), } } @@ -744,15 +1553,122 @@ impl TeamSessionService { self.sessions.get(team_id).map(|e| e.session.user_id().to_owned()) } + pub(crate) fn capture_published_session(&self, expected: &TeamSession) -> Option> { + self.sessions + .get(expected.team_id()) + .and_then(|entry| std::ptr::eq(entry.session.as_ref(), expected).then(|| Arc::clone(&entry.session))) + } + + /// Run a synchronous side effect only while `expected` is still the + /// published session. Keeping the map guard alive through `action` + /// serializes the effect with session removal/replacement. + pub(crate) fn with_published_session( + &self, + expected: &TeamSession, + action: impl FnOnce(&TeamSession) -> R, + ) -> Option { + let entry = self.sessions.get(expected.team_id())?; + std::ptr::eq(entry.session.as_ref(), expected).then(|| action(&entry.session)) + } + + pub(crate) fn publish_member_runtime_ready_if_current(&self, expected: &TeamSession, agent: &TeamAgent) -> bool { + self.with_published_session(expected, |_| { + self.broadcast_agent_runtime_status(expected.team_id(), agent, TeamAgentRuntimeStatus::Ready, None); + }) + .is_some() + } + + pub(crate) fn publish_member_runtime_starting_if_current(&self, expected: &TeamSession) -> bool { + self.with_published_session(expected, |_| { + self.broadcast_session_status( + expected.team_id(), + TeamSessionStatus::Starting, + Some(TeamSessionPhase::AttachingAgents), + |_| {}, + ); + }) + .is_some() + } + + pub(crate) fn publish_member_runtime_failed_if_current(&self, expected: &TeamSession, reason: &str) -> bool { + self.with_published_session(expected, |_| { + self.broadcast_session_status( + expected.team_id(), + TeamSessionStatus::Failed, + Some(TeamSessionPhase::AttachingAgents), + |payload| payload.error = Some(reason.to_owned()), + ); + }) + .is_some() + } + + pub(crate) async fn refresh_member_runtime_status(&self, expected: &TeamSession) { + let membership_lock = self + .add_agent_locks + .entry(expected.team_id().to_owned()) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone(); + let _membership_guard = membership_lock.lock().await; + let Ok(Some(row)) = self.repo.get_team(expected.team_id()).await else { + return; + }; + let Ok(team) = Team::from_row(&row) else { + return; + }; + + let mut failed_reason = None; + let mut pending = false; + for agent in &team.agents { + match expected.member_runtimes().snapshot(&agent.slot_id) { + MemberRuntimeSnapshot::Ready => {} + MemberRuntimeSnapshot::Failed { failure, .. } => { + failed_reason.get_or_insert(failure.public_reason); + } + MemberRuntimeSnapshot::Absent + | MemberRuntimeSnapshot::Attaching { .. } + | MemberRuntimeSnapshot::Removing { .. } => pending = true, + MemberRuntimeSnapshot::SessionStopped => return, + } + } + + if let Some(reason) = failed_reason { + self.publish_member_runtime_failed_if_current(expected, &reason); + } else if pending { + self.publish_member_runtime_starting_if_current(expected); + } else { + let _ = self.with_published_session(expected, |_| { + self.broadcast_session_status(expected.team_id(), TeamSessionStatus::Ready, None, |payload| { + payload.server_count = Some(team.agents.len()); + }); + }); + } + } + pub async fn get_run_state(&self, user_id: &str, team_id: &str) -> Result { self.load_owned_team(user_id, team_id).await?; let session = self.sessions.get(team_id).map(|entry| Arc::clone(&entry.session)); - let active_run = match session { - Some(session) => session.team_run_manager().current_payload().await, - None => None, + let Some(session) = session else { + return Ok(TeamRunStateResponse { + session_generation: None, + active_run: None, + slot_work: Vec::new(), + }); }; - - Ok(TeamRunStateResponse { active_run }) + let snapshot = session.work_coordinator().snapshot(); + let active_run = session.team_run_manager().current_payload(&snapshot).filter(|run| { + matches!( + run.status, + aionui_api_types::TeamRunStatus::Accepted + | aionui_api_types::TeamRunStatus::Running + | aionui_api_types::TeamRunStatus::Cancelling + ) + }); + let slot_work = snapshot.slots.iter().map(TeamRunManager::slot_payload).collect(); + Ok(TeamRunStateResponse { + session_generation: Some(snapshot.session_generation), + active_run, + slot_work, + }) } pub fn get_session_scheduler(&self, team_id: &str) -> Option> { @@ -781,37 +1697,123 @@ impl TeamSessionService { fn stop_session_unchecked(&self, team_id: &str) { if let Some((_, entry)) = self.sessions.remove(team_id) { entry.slow_monitor_handle.abort(); - entry.session.event_loops().shutdown(); entry.session.stop(); } } - pub async fn send_message( + pub async fn cleanup_idle_team_runtime_tasks( &self, - user_id: &str, - team_id: &str, - content: &str, - files: Option>, - ) -> Result { - self.load_owned_team(user_id, team_id).await?; - self.ensure_session_inner(team_id).await?; - let session = { - let entry = self - .sessions - .get(team_id) - .ok_or_else(|| TeamError::SessionNotFound(team_id.into()))?; - Arc::clone(&entry.session) - }; - session.send_message(content, files).await - } + idle_conversation_ids: Vec, + active_leases: &ActiveLeaseRegistry, + idle_threshold_ms: TimestampMs, + ) -> Vec { + if idle_conversation_ids.is_empty() { + return Vec::new(); + } - pub async fn send_message_to_agent( - &self, - user_id: &str, - team_id: &str, - slot_id: &str, - content: &str, - files: Option>, + let idle_conversation_set: HashSet = idle_conversation_ids.iter().cloned().collect(); + let now = now_ms(); + let mut handled_conversations = HashSet::new(); + let mut cleanup_teams = Vec::new(); + + for entry in self.sessions.iter() { + let team_id = entry.key().clone(); + let session = Arc::clone(&entry.session); + let agents = session.scheduler().list_agents().await; + let matched_idle_count = agents + .iter() + .filter(|agent| idle_conversation_set.contains(&agent.conversation_id)) + .count(); + if matched_idle_count == 0 { + continue; + } + + for agent in &agents { + handled_conversations.insert(agent.conversation_id.clone()); + } + + if session.team_run_manager().current_active_run_id().is_some() { + debug!( + team_id, + matched_idle_count, "team idle cleanup skipped because team run is active" + ); + continue; + } + + if agents + .iter() + .any(|agent| active_leases.active_until(&agent.conversation_id).is_some()) + { + debug!( + team_id, + matched_idle_count, "team idle cleanup skipped because at least one member has an active lease" + ); + continue; + } + + if !agents.iter().all(|agent| { + self.task_manager + .get_task(&agent.conversation_id) + .map(|task| is_idle_collectable_team_member(&task, now, idle_threshold_ms)) + .unwrap_or(true) + }) { + debug!( + team_id, + matched_idle_count, "team idle cleanup skipped because at least one member runtime task is active" + ); + continue; + } + + cleanup_teams.push((team_id, agents, matched_idle_count)); + } + + for (team_id, agents, matched_idle_count) in cleanup_teams { + info!( + team_id, + matched_idle_count, + member_count = agents.len(), + "team idle cleanup stopping idle team session" + ); + self.stop_session_unchecked(&team_id); + for agent in agents { + self.task_manager + .kill_and_wait(&agent.conversation_id, Some(AgentKillReason::IdleTimeout)) + .await; + } + } + + idle_conversation_ids + .into_iter() + .filter(|conversation_id| !handled_conversations.contains(conversation_id)) + .collect() + } + + pub async fn send_message( + &self, + user_id: &str, + team_id: &str, + content: &str, + files: Option>, + ) -> Result { + self.load_owned_team(user_id, team_id).await?; + self.ensure_session_inner(team_id).await?; + let session = { + let entry = self + .sessions + .get(team_id) + .ok_or_else(|| TeamError::SessionNotFound(team_id.into()))?; + Arc::clone(&entry.session) + }; + session.send_message(content, files).await + } + + pub async fn send_message_to_agent( + &self, + user_id: &str, + team_id: &str, + slot_id: &str, + content: &str, + files: Option>, ) -> Result { self.load_owned_team(user_id, team_id).await?; self.ensure_session_inner(team_id).await?; @@ -888,20 +1890,34 @@ impl TeamSessionService { pub async fn set_session_mode(&self, user_id: &str, team_id: &str, mode: &str) -> Result<(), TeamError> { let team = self.load_owned_team(user_id, team_id).await?; let provisioner = self.provisioner(); + self.repo + .update_team( + team_id, + &UpdateTeamParams { + session_mode: Some(mode.to_owned()), + ..Default::default() + }, + ) + .await?; for agent in &team.agents { - if let Some(instance) = self.task_manager.get_task(&agent.conversation_id) - && let Err(e) = set_active_agent_session_mode(&instance, mode).await - { - warn!( - team_id, - slot_id = %agent.slot_id, - conversation_id = %agent.conversation_id, - error = %e, - "failed to set session mode on agent" - ); - } - if let Err(e) = provisioner.update_session_mode_seed(agent, mode).await { + let mode_applied = match self.task_manager.get_task(&agent.conversation_id) { + Some(instance) => match set_active_agent_session_mode(&instance, mode).await { + Ok(()) => true, + Err(e) => { + warn!( + team_id, + slot_id = %agent.slot_id, + conversation_id = %agent.conversation_id, + error = %e, + "failed to set session mode on agent" + ); + false + } + }, + None => true, + }; + if mode_applied && let Err(e) = provisioner.update_session_mode_seed(agent, mode).await { warn!( team_id, slot_id = %agent.slot_id, @@ -952,41 +1968,6 @@ impl TeamSessionService { session.shutdown_agent(caller_slot_id, target_slot_id, reason).await } - pub(crate) fn notify_reserved_wake_for_team_work( - &self, - team_id: &str, - slot_id: &str, - target_role: TeamRunTargetRole, - source: TeamWakeSource, - ) { - let Some(entry) = self.sessions.get(team_id) else { - warn!( - team_id, - slot_id, - target_role = ?target_role, - wake_source = %source, - "reserved wake notify skipped because session is missing" - ); - return; - }; - entry - .session - .notify_reserved_wake_for_team_work(slot_id, target_role, source); - } - - pub(crate) fn notify_mailbox_only_wake(&self, team_id: &str, slot_id: &str, source: TeamWakeSource) { - let Some(entry) = self.sessions.get(team_id) else { - warn!( - team_id, - slot_id, - wake_source = %source, - "mailbox-only wake notify skipped because session is missing" - ); - return; - }; - entry.session.notify_mailbox_only_wake(slot_id, source); - } - /// Friendly pre-check used before invoking run-scoped team tools. This is /// not a concurrency guarantee; any operation /// that writes mailbox, projection, scheduler, spawn, shutdown, or wake state @@ -996,7 +1977,7 @@ impl TeamSessionService { .sessions .get(team_id) .ok_or_else(|| TeamError::SessionNotFound(team_id.into()))?; - if entry.session.team_run_manager().active_run_id().await.is_some() { + if entry.session.team_run_manager().current_active_run_id().is_some() { return Ok(()); } Err(TeamError::InvalidRequest( @@ -1004,27 +1985,11 @@ impl TeamSessionService { )) } - pub(crate) async fn notify_leader_spawn_attach_failed( - &self, - team_id: &str, - failed_slot_id: &str, - error: &str, - ) -> Result<(), TeamError> { - let entry = self - .sessions - .get(team_id) - .ok_or_else(|| TeamError::SessionNotFound(team_id.into()))?; - entry - .session - .notify_leader_spawn_attach_failed(failed_slot_id, error) - .await - } - pub(crate) async fn wake_leader_after_recovery_message( &self, team_id: &str, source_slot_id: &str, - source: TeamWakeSource, + source: WorkSource, ) -> Result<(), TeamError> { let entry = self .sessions @@ -1046,12 +2011,326 @@ async fn set_active_agent_session_mode(instance: &AgentInstance, mode: &str) -> } } +fn is_idle_collectable_team_member(task: &AgentInstance, now: TimestampMs, idle_threshold_ms: TimestampMs) -> bool { + if !matches!( + task.status(), + None | Some(ConversationStatus::Pending | ConversationStatus::Finished) + ) { + return false; + } + now.saturating_sub(task.last_activity_at()) > idle_threshold_ms +} + #[cfg(test)] mod tests { + use std::collections::HashMap; + use std::sync::{Arc, Mutex}; + + use aionui_ai_agent::types::{BuildTaskOptions, SendMessageData}; + use aionui_ai_agent::{ + ActiveLeaseRegistry, AgentError, AgentInstance, AgentSendError, AgentStreamEvent, IAgentTask, IMockAgent, + IWorkerTaskManager, IdleCleanupCoordinator, + }; + use aionui_api_types::{AddAgentRequest, ConfigOptionConfirmation, SetConfigOptionResponse}; + use aionui_common::{AgentKillReason, AgentType, ConversationStatus, TimestampMs, now_ms}; + use aionui_db::{IConversationRepository, ITeamRepository}; + use tokio::sync::broadcast; + + use super::TeamIdleCleanupCoordinator; use crate::test_utils::workspace_harness::{ - setup_with_factory_metadata_team_repo_and_conversation_repo, single_agent_team_request, + setup_with_factory_metadata_team_repo_and_conversation_repo, + setup_with_factory_metadata_team_repo_conversation_repo_and_broadcaster, + setup_with_factory_metadata_team_repo_conversation_repo_broadcaster_and_task_manager, + single_agent_team_request, }; + struct ModeSettingAgent { + conversation_id: String, + agent_type: AgentType, + mode_result: Mutex>, + event_tx: broadcast::Sender, + status: Option, + last_activity_at: TimestampMs, + } + + impl ModeSettingAgent { + fn accepts_mode(conversation_id: &str) -> Self { + Self::new(conversation_id, Ok(())) + } + + fn rejects_mode(conversation_id: &str, message: &str) -> Self { + Self::new(conversation_id, Err(message.to_owned())) + } + + fn new(conversation_id: &str, mode_result: Result<(), String>) -> Self { + let (event_tx, _) = broadcast::channel(1); + Self { + conversation_id: conversation_id.to_owned(), + agent_type: AgentType::Acp, + mode_result: Mutex::new(mode_result), + event_tx, + status: None, + last_activity_at: now_ms(), + } + } + + fn idle_finished(conversation_id: &str) -> Self { + Self::accepts_mode(conversation_id) + .with_status(Some(ConversationStatus::Finished)) + .with_last_activity(now_ms() - 600_000) + } + + fn idle_pending_aionrs(conversation_id: &str) -> Self { + Self::accepts_mode(conversation_id) + .with_agent_type(AgentType::Aionrs) + .with_status(Some(ConversationStatus::Pending)) + .with_last_activity(now_ms() - 600_000) + } + + fn with_agent_type(mut self, agent_type: AgentType) -> Self { + self.agent_type = agent_type; + self + } + + fn with_status(mut self, status: Option) -> Self { + self.status = status; + self + } + + fn with_last_activity(mut self, last_activity_at: TimestampMs) -> Self { + self.last_activity_at = last_activity_at; + self + } + } + + #[async_trait::async_trait] + impl IAgentTask for ModeSettingAgent { + fn agent_type(&self) -> AgentType { + self.agent_type + } + + fn conversation_id(&self) -> &str { + &self.conversation_id + } + + fn workspace(&self) -> &str { + "/tmp/aioncore-team-mode-test" + } + + fn status(&self) -> Option { + self.status + } + + fn last_activity_at(&self) -> TimestampMs { + self.last_activity_at + } + + fn subscribe(&self) -> broadcast::Receiver { + self.event_tx.subscribe() + } + + async fn send_message(&self, _data: SendMessageData) -> Result<(), AgentSendError> { + Ok(()) + } + + async fn cancel(&self) -> Result<(), AgentError> { + Ok(()) + } + + fn kill(&self, _reason: Option) -> Result<(), AgentError> { + Ok(()) + } + } + + #[async_trait::async_trait] + impl IMockAgent for ModeSettingAgent { + async fn set_config_option(&self, option_id: &str, value: &str) -> Result { + assert_eq!(option_id, "mode"); + assert_eq!(value, "read-only"); + match self.mode_result.lock().unwrap().clone() { + Ok(()) => Ok(SetConfigOptionResponse { + confirmation: ConfigOptionConfirmation::Observed, + config_options: None, + }), + Err(message) => Err(AgentError::bad_request(message)), + } + } + } + + struct StaticTaskManager { + tasks: HashMap, + } + + impl StaticTaskManager { + fn new(tasks: HashMap) -> Self { + Self { tasks } + } + } + + #[async_trait::async_trait] + impl IWorkerTaskManager for StaticTaskManager { + fn get_task(&self, conversation_id: &str) -> Option { + self.tasks.get(conversation_id).cloned() + } + + async fn get_or_build_task( + &self, + _conversation_id: &str, + _options: BuildTaskOptions, + ) -> Result { + Err(AgentError::internal("static task manager does not build tasks")) + } + + fn kill(&self, _conversation_id: &str, _reason: Option) -> Result<(), AgentError> { + Ok(()) + } + + fn kill_and_wait( + &self, + _conversation_id: &str, + _reason: Option, + ) -> std::pin::Pin + Send>> { + Box::pin(std::future::ready(())) + } + + async fn clear(&self) {} + + fn active_count(&self) -> usize { + self.tasks.len() + } + + fn collect_idle(&self, _idle_threshold_ms: TimestampMs) -> Vec { + Vec::new() + } + } + + struct MutableTaskManager { + tasks: Mutex>, + kills: Mutex>, + } + + impl MutableTaskManager { + fn new() -> Self { + Self { + tasks: Mutex::new(HashMap::new()), + kills: Mutex::new(Vec::new()), + } + } + + fn insert_mode_agent(&self, conversation_id: &str) { + self.tasks.lock().unwrap().insert( + conversation_id.to_owned(), + AgentInstance::Mock(Arc::new(ModeSettingAgent::accepts_mode(conversation_id))), + ); + } + + fn insert_idle_finished_agent(&self, conversation_id: &str) { + self.tasks.lock().unwrap().insert( + conversation_id.to_owned(), + AgentInstance::Mock(Arc::new(ModeSettingAgent::idle_finished(conversation_id))), + ); + } + + fn insert_idle_pending_aionrs_agent(&self, conversation_id: &str) { + self.tasks.lock().unwrap().insert( + conversation_id.to_owned(), + AgentInstance::Mock(Arc::new(ModeSettingAgent::idle_pending_aionrs(conversation_id))), + ); + } + + fn remove(&self, conversation_id: &str) { + self.tasks.lock().unwrap().remove(conversation_id); + } + + fn reset_kills(&self) { + self.kills.lock().unwrap().clear(); + } + + fn kills(&self) -> Vec { + self.kills.lock().unwrap().clone() + } + } + + fn two_agent_team_request(name: &str) -> aionui_api_types::CreateTeamRequest { + aionui_api_types::CreateTeamRequest { + name: name.into(), + agents: vec![ + aionui_api_types::TeamAgentInput { + name: "Lead".into(), + role: "lead".into(), + backend: Some("acp".into()), + model: "claude".into(), + assistant_id: None, + conversation_id: None, + }, + aionui_api_types::TeamAgentInput { + name: "Worker".into(), + role: "teammate".into(), + backend: Some("acp".into()), + model: "claude".into(), + assistant_id: None, + conversation_id: None, + }, + ], + workspace: None, + } + } + + fn team_with_aionrs_worker_request(name: &str) -> aionui_api_types::CreateTeamRequest { + let mut request = two_agent_team_request(name); + request.agents.push(aionui_api_types::TeamAgentInput { + name: "Butler".into(), + role: "teammate".into(), + backend: Some("aionrs".into()), + model: "claude-sonnet".into(), + assistant_id: None, + conversation_id: None, + }); + request + } + + #[async_trait::async_trait] + impl IWorkerTaskManager for MutableTaskManager { + fn get_task(&self, conversation_id: &str) -> Option { + self.tasks.lock().unwrap().get(conversation_id).cloned() + } + + async fn get_or_build_task( + &self, + _conversation_id: &str, + _options: BuildTaskOptions, + ) -> Result { + Err(AgentError::internal("mutable task manager does not build tasks")) + } + + fn kill(&self, conversation_id: &str, _reason: Option) -> Result<(), AgentError> { + self.kills.lock().unwrap().push(conversation_id.to_owned()); + self.remove(conversation_id); + Ok(()) + } + + fn kill_and_wait( + &self, + conversation_id: &str, + reason: Option, + ) -> std::pin::Pin + Send>> { + let _ = self.kill(conversation_id, reason); + Box::pin(std::future::ready(())) + } + + async fn clear(&self) { + self.tasks.lock().unwrap().clear(); + } + + fn active_count(&self) -> usize { + self.tasks.lock().unwrap().len() + } + + fn collect_idle(&self, _idle_threshold_ms: TimestampMs) -> Vec { + Vec::new() + } + } + #[tokio::test] async fn session_has_slow_monitor() { let (svc, _repo, _task_manager, _conv_repo) = setup_with_factory_metadata_team_repo_and_conversation_repo(); @@ -1066,6 +2345,526 @@ mod tests { svc.stop_session("user-test", &created.id).await.unwrap(); } + #[tokio::test] + async fn ensure_session_emits_agent_runtime_ready_after_member_warmup() { + let (svc, _repo, _task_manager, _conv_repo, broadcaster) = + setup_with_factory_metadata_team_repo_conversation_repo_and_broadcaster(); + let created = svc + .create_team("user-test", single_agent_team_request("Runtime Events")) + .await + .unwrap(); + let assistant = created.assistants.first().expect("team assistant"); + + svc.ensure_session("user-test", &created.id).await.unwrap(); + + let events = broadcaster.events_by_name("team.agentRuntimeStatusChanged"); + let statuses: Vec<&str> = events + .iter() + .map(|event| event.data.get("status").and_then(serde_json::Value::as_str).unwrap()) + .collect(); + + assert_eq!(statuses, vec!["pending", "ready"]); + assert_eq!( + events[0].data.get("team_id").and_then(serde_json::Value::as_str), + Some(created.id.as_str()) + ); + assert_eq!( + events[0].data.get("slot_id").and_then(serde_json::Value::as_str), + Some(assistant.slot_id.as_str()) + ); + assert_eq!( + events[0] + .data + .get("conversation_id") + .and_then(serde_json::Value::as_str), + Some(assistant.conversation_id.as_str()) + ); + } + + #[tokio::test] + async fn ensure_session_repairs_only_missing_member_runtime_in_place() { + let task_manager = Arc::new(MutableTaskManager::new()); + let (svc, _repo, _task_manager, _conv_repo, broadcaster) = + setup_with_factory_metadata_team_repo_conversation_repo_broadcaster_and_task_manager(task_manager.clone()); + let created = svc + .create_team("user-test", two_agent_team_request("Runtime Repair")) + .await + .unwrap(); + let lead = created.assistants.iter().find(|agent| agent.role == "lead").unwrap(); + let worker = created + .assistants + .iter() + .find(|agent| agent.role == "teammate") + .unwrap(); + + svc.ensure_session("user-test", &created.id).await.unwrap(); + task_manager.insert_mode_agent(&lead.conversation_id); + task_manager.insert_mode_agent(&worker.conversation_id); + task_manager.reset_kills(); + let original_session = Arc::clone(&svc.sessions.get(&created.id).expect("session").session); + let original_generation = original_session.generation(); + task_manager.remove(&worker.conversation_id); + + svc.ensure_session("user-test", &created.id).await.unwrap(); + + let current_session = Arc::clone(&svc.sessions.get(&created.id).expect("session").session); + assert!(Arc::ptr_eq(&original_session, ¤t_session)); + assert_eq!(current_session.generation(), original_generation); + assert_eq!(task_manager.kills(), vec![worker.conversation_id.clone()]); + assert!(current_session.event_loops().has(&lead.slot_id)); + assert!(current_session.event_loops().has(&worker.slot_id)); + + let events = broadcaster.events_by_name("team.agentRuntimeStatusChanged"); + let statuses: Vec<&str> = events + .iter() + .filter(|event| { + event.data.get("slot_id").and_then(serde_json::Value::as_str) == Some(worker.slot_id.as_str()) + }) + .map(|event| event.data.get("status").and_then(serde_json::Value::as_str).unwrap()) + .collect(); + + assert_eq!(statuses, vec!["pending", "ready", "pending", "ready"]); + } + + #[tokio::test] + async fn no_work_reconciliation_rejects_replaced_session() { + let (svc, _repo, _task_manager, _conv_repo) = setup_with_factory_metadata_team_repo_and_conversation_repo(); + let created = svc + .create_team("user-test", single_agent_team_request("No work replacement")) + .await + .unwrap(); + svc.ensure_session("user-test", &created.id).await.unwrap(); + let old = Arc::clone(&svc.sessions.get(&created.id).unwrap().session); + svc.stop_session("user-test", &created.id).await.unwrap(); + svc.ensure_session("user-test", &created.id).await.unwrap(); + + let result = svc + .complete_member_runtime_reconciliation(&created.id, "user-test", old, Vec::new()) + .await; + assert!(matches!(result, Err(crate::TeamError::SessionNotFound(_)))); + } + + #[tokio::test] + async fn replaced_session_cannot_publish_dynamic_runtime_ready() { + let (svc, _repo, _task_manager, _conv_repo, broadcaster) = + setup_with_factory_metadata_team_repo_conversation_repo_and_broadcaster(); + let created = svc + .create_team("user-test", single_agent_team_request("Runtime ready generation fence")) + .await + .unwrap(); + svc.ensure_session("user-test", &created.id).await.unwrap(); + let old = Arc::clone(&svc.sessions.get(&created.id).unwrap().session); + let agent = old.scheduler().list_agents().await.remove(0); + + svc.stop_session("user-test", &created.id).await.unwrap(); + svc.ensure_session("user-test", &created.id).await.unwrap(); + let ready_before = broadcaster + .events_by_name("team.agentRuntimeStatusChanged") + .into_iter() + .filter(|event| event.data.get("status").and_then(serde_json::Value::as_str) == Some("ready")) + .count(); + + assert!(!svc.publish_member_runtime_ready_if_current(&old, &agent)); + let ready_after = broadcaster + .events_by_name("team.agentRuntimeStatusChanged") + .into_iter() + .filter(|event| event.data.get("status").and_then(serde_json::Value::as_str) == Some("ready")) + .count(); + assert_eq!(ready_after, ready_before, "old generations must not publish Ready"); + } + + #[tokio::test] + async fn stale_cleanup_does_not_kill_replacement_runtime() { + let task_manager = Arc::new(MutableTaskManager::new()); + let (svc, _repo, _task_manager, _conv_repo, _broadcaster) = + setup_with_factory_metadata_team_repo_conversation_repo_broadcaster_and_task_manager(task_manager.clone()); + let created = svc + .create_team("user-test", single_agent_team_request("Cleanup fence")) + .await + .unwrap(); + svc.ensure_session("user-test", &created.id).await.unwrap(); + let old = Arc::clone(&svc.sessions.get(&created.id).unwrap().session); + svc.stop_session("user-test", &created.id).await.unwrap(); + svc.ensure_session("user-test", &created.id).await.unwrap(); + let conversation_id = created.assistants[0].conversation_id.clone(); + task_manager.insert_mode_agent(&conversation_id); + + svc.cleanup_stale_member_runtime_task(&old, &conversation_id).await; + + assert!(task_manager.get_task(&conversation_id).is_some()); + } + + #[tokio::test] + async fn idle_cleanup_stops_team_session_and_kills_all_members_when_team_is_collectable() { + let task_manager = Arc::new(MutableTaskManager::new()); + let (svc, _repo, _task_manager, _conv_repo, _broadcaster) = + setup_with_factory_metadata_team_repo_conversation_repo_broadcaster_and_task_manager(task_manager.clone()); + let created = svc + .create_team("user-test", two_agent_team_request("Idle Cleanup")) + .await + .unwrap(); + let lead = created.assistants.iter().find(|agent| agent.role == "lead").unwrap(); + let worker = created + .assistants + .iter() + .find(|agent| agent.role == "teammate") + .unwrap(); + task_manager.insert_idle_finished_agent(&lead.conversation_id); + task_manager.insert_idle_finished_agent(&worker.conversation_id); + + svc.ensure_session("user-test", &created.id).await.unwrap(); + + let unhandled = svc + .cleanup_idle_team_runtime_tasks(vec![lead.conversation_id.clone()], &ActiveLeaseRegistry::new(), 300_000) + .await; + + assert!(unhandled.is_empty()); + assert_eq!(svc.session_count_for_test(), 0); + assert_eq!(task_manager.active_count(), 0); + } + + #[test] + fn idle_collectable_team_member_accepts_idle_pending_aionrs_runtime() { + let task = AgentInstance::Mock(Arc::new(ModeSettingAgent::idle_pending_aionrs("aionrs-idle"))); + + assert!(super::is_idle_collectable_team_member(&task, now_ms(), 300_000)); + } + + #[test] + fn idle_collectable_team_member_rejects_running_aionrs_runtime() { + let task = AgentInstance::Mock(Arc::new( + ModeSettingAgent::accepts_mode("aionrs-running") + .with_agent_type(AgentType::Aionrs) + .with_status(Some(ConversationStatus::Running)) + .with_last_activity(now_ms() - 600_000), + )); + + assert!(!super::is_idle_collectable_team_member(&task, now_ms(), 300_000)); + } + + #[tokio::test] + async fn idle_cleanup_stops_team_session_when_aionrs_member_is_idle_pending() { + let task_manager = Arc::new(MutableTaskManager::new()); + let (svc, _repo, _task_manager, _conv_repo, _broadcaster) = + setup_with_factory_metadata_team_repo_conversation_repo_broadcaster_and_task_manager(task_manager.clone()); + let created = svc + .create_team("user-test", team_with_aionrs_worker_request("Idle Aionrs Cleanup")) + .await + .unwrap(); + let lead = created.assistants.iter().find(|agent| agent.role == "lead").unwrap(); + let acp_worker = created.assistants.iter().find(|agent| agent.name == "Worker").unwrap(); + let aionrs_worker = created.assistants.iter().find(|agent| agent.name == "Butler").unwrap(); + task_manager.insert_idle_finished_agent(&lead.conversation_id); + task_manager.insert_idle_finished_agent(&acp_worker.conversation_id); + task_manager.insert_idle_pending_aionrs_agent(&aionrs_worker.conversation_id); + + svc.ensure_session("user-test", &created.id).await.unwrap(); + + let unhandled = svc + .cleanup_idle_team_runtime_tasks( + vec![lead.conversation_id.clone(), acp_worker.conversation_id.clone()], + &ActiveLeaseRegistry::new(), + 300_000, + ) + .await; + + assert!(unhandled.is_empty()); + assert_eq!(svc.session_count_for_test(), 0); + assert_eq!(task_manager.active_count(), 0); + } + + #[tokio::test] + async fn team_idle_cleanup_coordinator_delegates_to_team_service() { + let task_manager = Arc::new(MutableTaskManager::new()); + let (svc, _repo, _task_manager, _conv_repo, _broadcaster) = + setup_with_factory_metadata_team_repo_conversation_repo_broadcaster_and_task_manager(task_manager.clone()); + let created = svc + .create_team("user-test", two_agent_team_request("Idle Coordinator")) + .await + .unwrap(); + for agent in &created.assistants { + task_manager.insert_idle_finished_agent(&agent.conversation_id); + } + + svc.ensure_session("user-test", &created.id).await.unwrap(); + let coordinator = TeamIdleCleanupCoordinator::new(svc.clone(), Arc::new(ActiveLeaseRegistry::new())); + + let unhandled = coordinator + .cleanup_idle_conversations(vec![created.assistants[0].conversation_id.clone()], 300_000) + .await; + + assert!(unhandled.is_empty()); + assert_eq!(svc.session_count_for_test(), 0); + assert_eq!(task_manager.active_count(), 0); + } + + #[tokio::test] + async fn manual_add_agent_in_active_session_emits_runtime_ready_after_background_attach() { + let (svc, _repo, _task_manager, _conv_repo, broadcaster) = + setup_with_factory_metadata_team_repo_conversation_repo_and_broadcaster(); + let created = svc + .create_team("user-test", single_agent_team_request("Manual Runtime Events")) + .await + .unwrap(); + svc.ensure_session("user-test", &created.id).await.unwrap(); + + let added = svc + .add_agent( + "user-test", + &created.id, + AddAgentRequest { + name: "Worker".to_owned(), + role: "teammate".to_owned(), + backend: Some("acp".to_owned()), + model: "claude".to_owned(), + assistant_id: None, + }, + ) + .await + .unwrap(); + + let events = tokio::time::timeout(std::time::Duration::from_secs(1), async { + loop { + let events = broadcaster.events_by_name("team.agentRuntimeStatusChanged"); + let added_events: Vec<_> = events + .into_iter() + .filter(|event| { + event.data.get("slot_id").and_then(serde_json::Value::as_str) == Some(added.slot_id.as_str()) + }) + .collect(); + if added_events + .iter() + .any(|event| event.data.get("status").and_then(serde_json::Value::as_str) == Some("ready")) + { + break added_events; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("runtime ready event should be emitted"); + let statuses: Vec<&str> = events + .iter() + .map(|event| event.data.get("status").and_then(serde_json::Value::as_str).unwrap()) + .collect(); + + assert_eq!(statuses, vec!["pending", "ready"]); + let session = Arc::clone(&svc.sessions.get(&created.id).expect("session").session); + assert_eq!(session.event_loops().len(), 2, "new slot must own one event loop"); + assert!(session.event_loops().has(&added.slot_id)); + assert_eq!( + session.member_runtimes().snapshot(&added.slot_id), + crate::member_runtime::MemberRuntimeSnapshot::Ready + ); + } + + #[tokio::test] + async fn dynamic_member_is_reserved_before_membership_event() { + let (svc, _repo, _task_manager, _conv_repo, broadcaster) = + setup_with_factory_metadata_team_repo_conversation_repo_and_broadcaster(); + let created = svc + .create_team("user-test", single_agent_team_request("Reservation ordering")) + .await + .unwrap(); + svc.ensure_session("user-test", &created.id).await.unwrap(); + let session = Arc::clone(&svc.sessions.get(&created.id).unwrap().session); + let observed = Arc::new(std::sync::Mutex::new(None)); + let observed_for_event = Arc::clone(&observed); + broadcaster.set_observer(Arc::new(move |event| { + if event.name != crate::events::TEAM_AGENT_SPAWNED_EVENT { + return; + } + let slot_id = event + .data + .get("assistant") + .and_then(|assistant| assistant.get("slot_id")) + .and_then(serde_json::Value::as_str); + if let Some(slot_id) = slot_id { + *observed_for_event.lock().unwrap() = Some(session.member_runtimes().snapshot(slot_id)); + } + })); + + svc.add_agent( + "user-test", + &created.id, + AddAgentRequest { + name: "Worker".to_owned(), + role: "teammate".to_owned(), + backend: Some("acp".to_owned()), + model: "claude".to_owned(), + assistant_id: None, + }, + ) + .await + .unwrap(); + + assert!(matches!( + observed.lock().unwrap().as_ref(), + Some( + crate::member_runtime::MemberRuntimeSnapshot::Attaching { .. } + | crate::member_runtime::MemberRuntimeSnapshot::Ready + ) + )); + } + + #[tokio::test] + async fn set_session_mode_persists_team_mode_and_new_agents_inherit_it() { + let (svc, repo, _task_manager, conv_repo) = setup_with_factory_metadata_team_repo_and_conversation_repo(); + let created = svc + .create_team("user-test", single_agent_team_request("Team Mode Seed")) + .await + .unwrap(); + + svc.set_session_mode("user-test", &created.id, "full_auto") + .await + .unwrap(); + + let row = repo.get_team(&created.id).await.unwrap().expect("team row"); + assert_eq!(row.session_mode.as_deref(), Some("full_auto")); + + let added = svc + .add_agent( + "user-test", + &created.id, + AddAgentRequest { + name: "Worker".to_owned(), + role: "teammate".to_owned(), + backend: Some("acp".to_owned()), + model: "claude".to_owned(), + assistant_id: None, + }, + ) + .await + .unwrap(); + let extra = conv_repo + .get_extra(&added.conversation_id) + .expect("added conversation extra"); + + assert_eq!( + extra.get("session_mode").and_then(serde_json::Value::as_str), + Some("full_auto") + ); + } + + #[tokio::test] + async fn set_session_mode_does_not_persist_agent_seed_when_active_runtime_rejects_mode() { + let accepting_conversation_id = "conv-accepts"; + let rejecting_conversation_id = "conv-rejects"; + let task_manager = Arc::new(StaticTaskManager::new(HashMap::from([ + ( + accepting_conversation_id.to_owned(), + AgentInstance::Mock(Arc::new(ModeSettingAgent::accepts_mode(accepting_conversation_id))), + ), + ( + rejecting_conversation_id.to_owned(), + AgentInstance::Mock(Arc::new(ModeSettingAgent::rejects_mode( + rejecting_conversation_id, + "Value 'read-only' is not selectable for config option 'mode'", + ))), + ), + ]))); + let (svc, repo, _task_manager, conv_repo, _broadcaster) = + setup_with_factory_metadata_team_repo_conversation_repo_broadcaster_and_task_manager(task_manager); + let created = svc + .create_team("user-test", single_agent_team_request("Partial Mode Seed")) + .await + .unwrap(); + let mut row = repo.get_team(&created.id).await.unwrap().expect("team row"); + row.agents = serde_json::json!([ + { + "slot_id": "slot-accepts", + "name": "Codex CLI", + "role": "lead", + "conversation_id": accepting_conversation_id, + "backend": "codex", + "model": "openai.gpt-5.5", + "assistant_id": "bare:codex" + }, + { + "slot_id": "slot-rejects", + "name": "Claude Code", + "role": "teammate", + "conversation_id": rejecting_conversation_id, + "backend": "claude", + "model": "global.anthropic.claude-opus-4-8", + "assistant_id": "bare:claude" + } + ]) + .to_string(); + repo.update_team( + &created.id, + &aionui_db::UpdateTeamParams { + agents: Some(row.agents), + ..Default::default() + }, + ) + .await + .unwrap(); + conv_repo + .create(&aionui_db::models::ConversationRow { + id: accepting_conversation_id.to_owned(), + user_id: "user-test".to_owned(), + name: "Codex CLI".to_owned(), + r#type: AgentType::Acp.serde_name().to_owned(), + extra: serde_json::json!({ + "current_mode_id": "default", + "session_mode": "default" + }) + .to_string(), + model: None, + status: Some("pending".to_owned()), + source: None, + channel_chat_id: None, + pinned: false, + pinned_at: None, + created_at: now_ms(), + updated_at: now_ms(), + }) + .await + .unwrap(); + conv_repo + .create(&aionui_db::models::ConversationRow { + id: rejecting_conversation_id.to_owned(), + user_id: "user-test".to_owned(), + name: "Claude Code".to_owned(), + r#type: AgentType::Acp.serde_name().to_owned(), + extra: serde_json::json!({ + "current_mode_id": "default", + "session_mode": "default" + }) + .to_string(), + model: None, + status: Some("pending".to_owned()), + source: None, + channel_chat_id: None, + pinned: false, + pinned_at: None, + created_at: now_ms(), + updated_at: now_ms(), + }) + .await + .unwrap(); + + svc.set_session_mode("user-test", &created.id, "read-only") + .await + .unwrap(); + + let team = repo.get_team(&created.id).await.unwrap().expect("team row"); + assert_eq!(team.session_mode.as_deref(), Some("read-only")); + + let accepting_extra = conv_repo.get_extra(accepting_conversation_id).unwrap(); + assert_eq!( + accepting_extra.get("session_mode").and_then(serde_json::Value::as_str), + Some("read-only") + ); + + let rejecting_extra = conv_repo.get_extra(rejecting_conversation_id).unwrap(); + assert_eq!( + rejecting_extra.get("session_mode").and_then(serde_json::Value::as_str), + Some("default") + ); + } + #[tokio::test] async fn run_state_returns_none_without_session_and_does_not_create_session() { let (svc, _repo, _task_manager, _conv_repo) = setup_with_factory_metadata_team_repo_and_conversation_repo(); @@ -1083,6 +2882,26 @@ mod tests { assert_eq!(svc.session_count_for_test(), 0); } + #[tokio::test] + async fn config_options_returns_snapshot_without_creating_team_session() { + let (svc, _repo, _task_manager, _conv_repo) = setup_with_factory_metadata_team_repo_and_conversation_repo(); + let created = svc + .create_team("user-test", single_agent_team_request("Config Options")) + .await + .unwrap(); + let conversation_id = &created.assistants[0].conversation_id; + + assert_eq!(svc.session_count_for_test(), 0); + + let options = svc + .get_conversation_config_options("user-test", &created.id, conversation_id) + .await + .unwrap(); + + assert_eq!(options.config_options[0].id, "model"); + assert_eq!(svc.session_count_for_test(), 0); + } + #[tokio::test] async fn run_state_returns_current_active_payload() { let (svc, _repo, _task_manager, _conv_repo) = setup_with_factory_metadata_team_repo_and_conversation_repo(); @@ -1096,12 +2915,90 @@ mod tests { let active_run = state.active_run.expect("active run state"); assert_eq!(active_run.team_id, created.id); - assert_eq!(active_run.team_run_id, ack.team_run_id); - assert_eq!(active_run.status, ack.status); - assert_eq!(active_run.target_slot_id, ack.target_slot_id); - assert_eq!(active_run.target_role, ack.target_role); - assert_eq!(active_run.pending_wake_count, 1); + assert_eq!(active_run.team_run_id, ack.run.team_run_id); + assert_eq!(active_run.status, ack.run.status); + assert_eq!(active_run.target_slot_id, ack.run.target_slot_id); + assert_eq!(active_run.target_role, ack.run.target_role); + assert_eq!(active_run.queued_intent_count, 1); assert_eq!(active_run.slot_work.len(), 1); - assert_eq!(active_run.slot_work[0].slot_id, ack.accepted_slot_id); + assert_eq!(active_run.slot_work[0].slot_id, ack.run.slot_work[0].slot_id); + } + + #[tokio::test] + async fn config_options_return_member_runtime_snapshot() { + let (svc, _repo, _task_manager, _conv_repo) = setup_with_factory_metadata_team_repo_and_conversation_repo(); + let created = svc + .create_team("user-test", single_agent_team_request("Team Config")) + .await + .unwrap(); + let conversation_id = created.assistants[0].conversation_id.clone(); + + let response = svc + .get_conversation_config_options("user-test", &created.id, &conversation_id) + .await + .unwrap(); + + let model = response + .config_options + .iter() + .find(|option| option.id == "model") + .expect("model config option"); + assert_eq!(model.current_value.as_deref(), Some("claude")); + } + + #[tokio::test] + async fn config_options_reports_runtime_not_ready_for_member_conversation() { + let (svc, _repo, _task_manager, conv_repo) = setup_with_factory_metadata_team_repo_and_conversation_repo(); + let created = svc + .create_team("user-test", single_agent_team_request("Team Config Pending")) + .await + .unwrap(); + let conversation_id = created.assistants[0].conversation_id.clone(); + conv_repo.mark_runtime_not_ready(&conversation_id); + + let err = svc + .get_conversation_config_options("user-test", &created.id, &conversation_id) + .await + .expect_err("member runtime readiness should be reported distinctly"); + + assert!(matches!( + err, + crate::error::TeamError::RuntimeNotReady { + conversation_id: ref id + } if id == &conversation_id + )); + } + + #[tokio::test] + async fn config_options_reject_non_member_conversation() { + let (svc, _repo, _task_manager, _conv_repo) = setup_with_factory_metadata_team_repo_and_conversation_repo(); + let created = svc + .create_team("user-test", single_agent_team_request("Team Config Reject")) + .await + .unwrap(); + + let err = svc + .get_conversation_config_options("user-test", &created.id, "other-conversation") + .await + .expect_err("non-member conversation must be rejected"); + + assert!(matches!(err, crate::error::TeamError::AgentNotFound(_))); + } + + #[tokio::test] + async fn config_options_reject_cross_user_access() { + let (svc, _repo, _task_manager, _conv_repo) = setup_with_factory_metadata_team_repo_and_conversation_repo(); + let created = svc + .create_team("user-test", single_agent_team_request("Team Config Owner")) + .await + .unwrap(); + let conversation_id = created.assistants[0].conversation_id.clone(); + + let err = svc + .get_conversation_config_options("other-user", &created.id, &conversation_id) + .await + .expect_err("team config options must reject cross-user access"); + + assert!(matches!(err, crate::error::TeamError::Forbidden(_))); } } diff --git a/crates/aionui-team/src/service/describe_support.rs b/crates/aionui-team/src/service/describe_support.rs index 82cd5d70f..00b12c229 100644 --- a/crates/aionui-team/src/service/describe_support.rs +++ b/crates/aionui-team/src/service/describe_support.rs @@ -25,14 +25,45 @@ impl TeamSessionService { .unwrap_or(definition.agent_id.as_str()); let effective_backend = resolve_runtime_backend(&self.agent_metadata_repo, effective_agent_id).await?; - Ok(render_assistant_description( + Ok(render_assistant_description_json( &definition, &effective_backend, locale.unwrap_or("en-US"), - )) + )?) } } +fn render_assistant_description_json( + definition: &AssistantDefinitionRow, + effective_backend: &str, + locale: &str, +) -> Result { + let name_map = decode_str_map(&definition.name_i18n); + let description_map = decode_str_map(&definition.description_i18n); + let prompts_map = decode_list_map(&definition.recommended_prompts_i18n); + + let name = localized_text(&name_map, &definition.name, locale); + let description = localized_optional_text(&description_map, definition.description.as_deref(), locale) + .unwrap_or_else(|| "No description available.".to_owned()); + let example_tasks = localized_list(&prompts_map, &definition.recommended_prompts, locale).unwrap_or_default(); + let skills = decode_string_list(&definition.default_skill_ids) + .into_iter() + .chain(decode_string_list(&definition.custom_skill_names)) + .collect::>(); + let description_markdown = render_assistant_description(definition, effective_backend, locale); + + serde_json::to_string_pretty(&serde_json::json!({ + "status": "ok", + "assistant_id": definition.assistant_id, + "name": name, + "description": description, + "description_markdown": description_markdown, + "skills": skills, + "example_tasks": example_tasks, + "default_model": definition.default_model_value.clone(), + })) +} + fn render_assistant_description(definition: &AssistantDefinitionRow, effective_backend: &str, locale: &str) -> String { let name_map = decode_str_map(&definition.name_i18n); let description_map = decode_str_map(&definition.description_i18n); diff --git a/crates/aionui-team/src/service/spawn_support.rs b/crates/aionui-team/src/service/spawn_support.rs index 955477463..7175854ed 100644 --- a/crates/aionui-team/src/service/spawn_support.rs +++ b/crates/aionui-team/src/service/spawn_support.rs @@ -1,11 +1,10 @@ use super::*; -use aionui_api_types::BehaviorPolicy; use aionui_common::AgentType; -use aionui_common::constants::{TEAM_CAPABLE_BACKENDS, has_mcp_capability}; use aionui_db::models::AgentMetadataRow; use aionui_db::{IAgentMetadataRepository, resolve_agent_binding_from_rows}; use std::sync::Arc; +use crate::ports::TeamAssistantCatalogEntry; use crate::prompts::AvailableAssistant; use crate::provisioning::PersistSpawnedAgentRequest; @@ -61,6 +60,18 @@ pub(crate) async fn resolve_runtime_backend( } impl TeamSessionService { + pub(crate) async fn resolve_team_selectable_assistant( + &self, + assistant_id: &str, + ) -> Result { + self.assistant_catalog + .resolve_team_selectable_assistant(assistant_id) + .await? + .ok_or_else(|| { + TeamError::InvalidRequest(format!("Assistant is not available for team mode: {assistant_id}")) + }) + } + pub(crate) async fn resolve_spawn_backend_and_model( &self, assistant_id: Option<&str>, @@ -69,18 +80,13 @@ impl TeamSessionService { fallback_model: &str, ) -> Result<(String, String), TeamError> { if let Some(assistant_id) = assistant_id.map(str::trim).filter(|value| !value.is_empty()) { + let selectable = self.resolve_team_selectable_assistant(assistant_id).await?; let definition = self .assistant_definition_repo .get_by_assistant_id(assistant_id) .await? .ok_or_else(|| TeamError::InvalidRequest(format!("Preset assistant not found: {assistant_id}")))?; - let overlay = self.assistant_overlay_repo.get(&definition.id).await?; - let effective_agent_id = overlay - .as_ref() - .and_then(|row| row.agent_id_override.as_deref()) - .filter(|value| !value.trim().is_empty()) - .unwrap_or(definition.agent_id.as_str()); - let backend = resolve_runtime_backend(&self.agent_metadata_repo, effective_agent_id).await?; + let backend = selectable.backend; let requested_model = requested_model .map(str::trim) .filter(|value| !value.is_empty()) @@ -110,31 +116,6 @@ impl TeamSessionService { Ok((backend, model)) } - /// Check if a backend is allowed to participate in team mode. - /// Hard whitelist passes immediately; then checks behavior_policy.supports_team; - /// finally queries persisted `agent_capabilities` for MCP transport declarations. - pub(crate) async fn is_backend_team_capable(&self, backend: &str) -> bool { - if TEAM_CAPABLE_BACKENDS.contains(&backend) { - return true; - } - let Ok(Some(row)) = self.agent_metadata_repo.find_builtin_by_backend(backend).await else { - return false; - }; - let bp_supports = row - .behavior_policy - .as_deref() - .and_then(|s| serde_json::from_str::(s).ok()) - .is_some_and(|bp| bp.supports_team); - if bp_supports { - return true; - } - let caps = row - .agent_capabilities - .as_deref() - .and_then(|s| serde_json::from_str::(s).ok()); - has_mcp_capability(caps.as_ref()) - } - /// Return all enabled assistants that can currently participate in team mode. /// This consumes the same assistant projection as the Team creation UI, so /// `team_selectable` has a single source of truth. @@ -155,47 +136,6 @@ impl TeamSessionService { .collect() } - /// Return the `team_list_models` response built from DB rows. - /// Falls back to the hardcoded response if the DB query fails. - /// For internal agents (like aionrs with backend=NULL), enriches - /// with models from the providers table. - pub(crate) async fn list_models_from_db( - &self, - assistant_id_filter: Option<&str>, - ) -> Result { - let Ok(rows) = self.agent_metadata_repo.list_all().await else { - return Ok(crate::mcp::tools::handle_team_list_models(&serde_json::Value::Null)); - }; - let backend_filter = match assistant_id_filter.map(str::trim).filter(|value| !value.is_empty()) { - Some(assistant_id) => { - let definition = self - .assistant_definition_repo - .get_by_assistant_id(assistant_id) - .await? - .ok_or_else(|| TeamError::InvalidRequest(format!("Assistant not found: {assistant_id}")))?; - let overlay = self.assistant_overlay_repo.get(&definition.id).await?; - Some( - resolve_runtime_backend( - &self.agent_metadata_repo, - overlay - .as_ref() - .and_then(|row| row.agent_id_override.as_deref()) - .filter(|value| !value.trim().is_empty()) - .unwrap_or(definition.agent_id.as_str()), - ) - .await?, - ) - } - None => None, - }; - let provider_models = self.collect_provider_models().await; - Ok(crate::mcp::tools::build_list_models_from_rows( - &rows, - backend_filter.as_deref(), - &provider_models, - )) - } - /// Collect all enabled provider model IDs grouped by provider name. /// Returns a flat list of model IDs for use by internal agents (aionrs). async fn collect_provider_models(&self) -> Vec { @@ -558,6 +498,69 @@ mod tests { } } + fn assistant_definition(assistant_id: &str, agent_id: &str) -> AssistantDefinitionRow { + AssistantDefinitionRow { + id: format!("def-{assistant_id}"), + assistant_id: assistant_id.into(), + source: "builtin".into(), + owner_type: "system".into(), + source_ref: Some(assistant_id.into()), + source_version: None, + source_hash: None, + name: assistant_id.into(), + name_i18n: "{}".into(), + description: None, + description_i18n: "{}".into(), + avatar_type: "emoji".into(), + avatar_value: None, + agent_id: agent_id.into(), + rule_resource_type: "inline".into(), + rule_resource_ref: None, + rule_inline_content: None, + recommended_prompts: "[]".into(), + recommended_prompts_i18n: "{}".into(), + default_model_mode: "auto".into(), + default_model_value: None, + default_permission_mode: "auto".into(), + default_permission_value: None, + default_thought_level_mode: "auto".into(), + default_thought_level_value: None, + default_skills_mode: "auto".into(), + default_skill_ids: "[]".into(), + custom_skill_names: "[]".into(), + default_disabled_builtin_skill_ids: "[]".into(), + default_mcps_mode: "auto".into(), + default_mcp_ids: "[]".into(), + created_at: 0, + updated_at: 0, + deleted_at: None, + } + } + + fn service_with_selectable_catalog( + assistants: Vec, + definitions: Vec, + ) -> Arc { + let (base, _, _, _) = setup_with_factory_metadata_team_repo_and_conversation_repo(); + TeamSessionService::new( + base.repo.clone(), + base.agent_metadata_repo.clone(), + Arc::new(RowsTeamAssistantCatalog { rows: assistants }), + Arc::new(MultiAssistantDefinitionRepo { rows: definitions }), + Arc::new(MultiAssistantOverlayRepo { rows: vec![] }), + Arc::new(SingleProviderRepo { + rows: vec![provider_row("openai", &["gpt-5-mini"])], + }), + base.conversation_port.clone(), + base.projection_store.clone(), + base.broadcaster.clone(), + base.task_manager.clone(), + base.turn_port.clone(), + base.cancellation_port.clone(), + base.backend_binary_path.clone(), + ) + } + #[test] fn parse_agent_type_accepts_top_level_supported_runtimes() { assert_eq!(parse_agent_type("acp").unwrap(), AgentType::Acp); @@ -617,6 +620,69 @@ mod tests { assert_eq!(ids, vec!["assistant-unchecked"]); } + #[tokio::test] + async fn create_team_rejects_assistant_missing_from_team_selectable_catalog() { + let svc = service_with_selectable_catalog(vec![], vec![assistant_definition("word-creator", "aionrs")]); + let mut req = single_agent_team_request("Catalog Gate"); + req.agents[0].assistant_id = Some("word-creator".into()); + req.agents[0].backend = None; + req.agents[0].model = "gpt-5-mini".into(); + + let err = svc + .create_team("user1", req) + .await + .expect_err("assistant must be rejected when not team selectable"); + + assert!( + matches!(&err, TeamError::InvalidRequest(msg) if msg.contains("not available for team mode")), + "expected team-selectable assistant error, got {err:?}" + ); + } + + #[tokio::test] + async fn add_agent_rejects_assistant_missing_from_team_selectable_catalog() { + let svc = service_with_selectable_catalog(vec![], vec![assistant_definition("word-creator", "aionrs")]); + let created = svc + .create_team("user1", single_agent_team_request("Catalog Gate Add")) + .await + .unwrap(); + + let err = svc + .add_agent( + "user1", + &created.id, + aionui_api_types::AddAgentRequest { + name: "Worker".into(), + role: "teammate".into(), + backend: None, + model: "gpt-5-mini".into(), + assistant_id: Some("word-creator".into()), + }, + ) + .await + .expect_err("assistant must be rejected when not team selectable"); + + assert!( + matches!(&err, TeamError::InvalidRequest(msg) if msg.contains("not available for team mode")), + "expected team-selectable assistant error, got {err:?}" + ); + } + + #[tokio::test] + async fn resolve_spawn_backend_and_model_rejects_assistant_missing_from_team_selectable_catalog() { + let svc = service_with_selectable_catalog(vec![], vec![assistant_definition("word-creator", "aionrs")]); + + let err = svc + .resolve_spawn_backend_and_model(Some("word-creator"), None, "gemini", "gemini-2.5-pro") + .await + .expect_err("spawn must reject assistants outside the team-selectable catalog"); + + assert!( + matches!(&err, TeamError::InvalidRequest(msg) if msg.contains("not available for team mode")), + "expected team-selectable assistant error, got {err:?}" + ); + } + #[tokio::test] async fn persist_spawned_agent_uses_team_workspace_resolver() { let (svc, team_repo, _, conv_repo) = setup_with_factory_metadata_team_repo_and_conversation_repo(); @@ -659,7 +725,9 @@ mod tests { let svc = TeamSessionService::new( svc.repo.clone(), svc.agent_metadata_repo.clone(), - Arc::new(RowsTeamAssistantCatalog { rows: vec![] }), + Arc::new(RowsTeamAssistantCatalog { + rows: vec![team_assistant_entry("word-creator", "Word Creator", "aionrs")], + }), Arc::new(SingleAssistantDefinitionRepo { row: AssistantDefinitionRow { id: "def-1".into(), diff --git a/crates/aionui-team/src/session.rs b/crates/aionui-team/src/session.rs index 8cbad912a..7a3e0452e 100644 --- a/crates/aionui-team/src/session.rs +++ b/crates/aionui-team/src/session.rs @@ -1,11 +1,12 @@ use std::path::PathBuf; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Weak}; +use std::time::Instant; use aionui_ai_agent::IWorkerTaskManager; use aionui_api_types::{ - TeamRunAckResponse, TeamRunTargetRole, TeamSendMessageDelivery, TeamSendMessageReason, - TeamSendMessageTargetQueueState, TeamSlotRuntimeHealth, TeamSlotWorkPayload, + TeamAgentRuntimeStatus, TeamChildTurnPayload, TeamMessageEnqueueStatus, TeamRunAckResponse, TeamRunStatus, + TeamRunTargetRole, TeamSlotWorkPayload, }; use aionui_common::{AgentKillReason, generate_id}; use aionui_db::ITeamRepository; @@ -14,36 +15,38 @@ use tracing::{info, warn}; use crate::error::TeamError; use crate::event_loop::EventLoopRegistry; -use crate::events::TeamEventEmitter; +use crate::events::{TEAM_CHILD_TURN_CANCELLED_EVENT, TeamEventEmitter}; use crate::mailbox::Mailbox; use crate::mcp::{TeamMcpServer, TeamMcpStdioConfig, TeamMcpStdioServerSpec}; +use crate::member_runtime::{ + AttachLease, AttachOutcome, BeginRemove, MemberRuntimeFailure, MemberRuntimeRegistry, MemberRuntimeSnapshot, + ReserveAttach, +}; use crate::message_projection::{ TeamMessageProjection, TeamProjectionMessageStore, TeamProjectionRequest, TeamProjectionSource, teammate_dedupe_key, }; use crate::ports::{AgentTurnCancellationPort, AgentTurnExecutionPort}; +use crate::prompt_dump::{TeamPromptDumpConfig, TeamWakePromptDump, dump_team_wake_prompt}; use crate::prompts::{build_lead_prompt, build_teammate_prompt, build_wake_payload}; use crate::provisioning::PersistSpawnedAgentRequest; use crate::scheduler::{TeammateManager, normalize_name}; use crate::service::TeamSessionService; use crate::task_board::TaskBoard; -use crate::team_run::{ - ChildCancelTarget, RecoveryBacklogResult, RecoveryWakeCandidate, TeamRunManager, TeamRunWakeAcquireOutcome, - WakeRecordDecision, target_role_for, -}; +use crate::team_run::{TeamRunManager, target_role_for}; use crate::types::{MailboxMessageType, Team, TeamAgent, TeammateRole, TeammateStatus}; -use crate::wake::TeamWakeSource; +use crate::work_coordinator::{ + CausalBinding, CommitResult, EnqueueCommit, EnqueueDisposition, EnqueueLease, EnqueueRequest, ReconcileDecision, + RuntimeConstraint, SlotWorkCoordinator, WorkBatch, +}; +use crate::work_source::WorkSource; /// Input for the wake path. Produced by [`TeamSession::compute_wake_input`], /// consumed by D7b's `send_message` / `send_message_to_agent` (not implemented /// in D7a). `first_message` includes the role prompt on cold starts. #[derive(Debug, Clone)] pub struct WakeInput { - pub team_run_id: Option, pub conversation_id: String, pub first_message: String, - /// `false` when the mailbox is empty — caller should skip wake and - /// leave the agent idle. - pub should_send: bool, /// Unread mailbox rows used to build `first_message`. Returned so the /// caller can mirror non-user senders into the target agent's conversation /// as left bubbles (matches AionUi `TeammateManager.wake()`). These are @@ -52,15 +55,20 @@ pub struct WakeInput { pub unread: Vec, /// Role of the wake target. pub agent_role: TeammateRole, - pub(crate) wake_source: Option, - pub(crate) trigger_message_id: Option, } #[derive(Debug, Clone, PartialEq, Eq)] pub struct AgentMessageQueueResult { - pub team_run_id: String, - pub delivery: TeamSendMessageDelivery, - pub target: TeamSendMessageTargetQueueState, + pub team_run_id: Option, + pub target: TeamSlotWorkPayload, +} + +pub(crate) enum PrepareBatchResult { + Execute { batch: Box, input: WakeInput }, + SettleSignals { intent_ids: Vec }, + WaitingForCompletion, + Blocked, + Quiescent, } /// Input for [`TeamSession::spawn_agent`]. Populated by the lead agent when @@ -69,12 +77,6 @@ pub struct AgentMessageQueueResult { pub struct SpawnAgentRequest { pub name: String, pub assistant_id: Option, - pub model: Option, -} - -enum SpawnWakePlan { - RunScoped(TeamRunTargetRole), - MailboxOnly, } pub struct TeamSession { @@ -89,6 +91,7 @@ pub struct TeamSession { cancellation_port: Arc, projection_store: Arc, team_run_manager: Arc, + work_coordinator: Arc, /// Owner user_id for this team — needed when spawn_agent creates a /// new conversation (conversations are scoped per user). user_id: String, @@ -104,6 +107,13 @@ pub struct TeamSession { /// Per-agent event loop registry. Each agent has a dedicated tokio task /// that drains its mailbox whenever notified. event_loops: Arc, + /// Per-member runtime lifecycle for this exact session incarnation. + /// Created once with a fresh generation, seeded before publication, read + /// by dynamic attach/reconcile operations, and permanently stopped before + /// bridge/event-loop shutdown. A generation is never reused by another + /// `TeamSession`, even for the same team id. + member_runtimes: Arc, + prompt_dump: TeamPromptDumpConfig, /// Set after the session lifecycle performs its system recovery mailbox scan. /// Written by `try_start_recovery_drain` and read by later scan attempts so /// ordinary event-loop notifications cannot repeatedly create recovery runs. @@ -125,13 +135,49 @@ impl TeamSession { projection_store: Arc, user_id: String, service: Weak, + ) -> Result { + Self::start_with_prompt_dump( + team, + repo, + broadcaster, + backend_binary_path, + task_manager, + turn_port, + cancellation_port, + projection_store, + user_id, + service, + TeamPromptDumpConfig::disabled(), + ) + .await + } + + #[allow(clippy::too_many_arguments)] + pub async fn start_with_prompt_dump( + team: Team, + repo: Arc, + broadcaster: Arc, + backend_binary_path: Arc, + task_manager: Arc, + turn_port: Arc, + cancellation_port: Arc, + projection_store: Arc, + user_id: String, + service: Weak, + prompt_dump: TeamPromptDumpConfig, ) -> Result { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); + let member_runtimes = Arc::new(MemberRuntimeRegistry::new(generate_id())); let team_run_manager = Arc::new(TeamRunManager::new( team.id.clone(), Arc::new(TeamEventEmitter::new(team.id.clone(), broadcaster.clone())), )); + let work_coordinator = Arc::new(SlotWorkCoordinator::new( + team.id.clone(), + member_runtimes.generation().to_owned(), + team_run_manager.clone(), + )); let scheduler = Arc::new(TeammateManager::new( team.id.clone(), @@ -142,12 +188,13 @@ impl TeamSession { )); let auth_token = aionui_common::generate_id(); - let mcp_server = TeamMcpServer::start( + let mcp_server = TeamMcpServer::start_with_prompt_dump( auth_token, scheduler.clone(), team.id.clone(), broadcaster.clone(), service.clone(), + Some(prompt_dump.clone()), ) .await?; @@ -171,10 +218,13 @@ impl TeamSession { cancellation_port, projection_store, team_run_manager, + work_coordinator, user_id, service, broadcaster, event_loops, + member_runtimes, + prompt_dump, recovery_scan_completed: AtomicBool::new(false), }) } @@ -195,6 +245,14 @@ impl TeamSession { &self.event_loops } + pub(crate) fn member_runtimes(&self) -> &Arc { + &self.member_runtimes + } + + pub(crate) fn generation(&self) -> &str { + self.member_runtimes.generation() + } + pub fn turn_port(&self) -> &Arc { &self.turn_port } @@ -203,10 +261,18 @@ impl TeamSession { &self.cancellation_port } + pub(crate) fn team_event_emitter(&self) -> Arc { + Arc::new(TeamEventEmitter::new(self.team.id.clone(), self.broadcaster.clone())) + } + pub fn team_run_manager(&self) -> &Arc { &self.team_run_manager } + pub(crate) fn work_coordinator(&self) -> &Arc { + &self.work_coordinator + } + pub fn mcp_stdio_config(&self, slot_id: &str) -> TeamMcpStdioConfig { TeamMcpStdioConfig { team_id: self.team.id.clone(), @@ -225,100 +291,133 @@ impl TeamSession { TeamMcpStdioServerSpec::from_config(binary_path.as_ref(), &self.mcp_stdio_config(slot_id)) } - /// Assemble the payload that will drive the next wake of `slot_id`. - /// - /// - Reads status, unread messages and tasks. - /// - Cold-start agents (no prior status, or last status was `Error`) - /// receive the full role prompt prepended to the wake payload. - /// - When the mailbox is empty, returns `WakeInput { should_send: false, .. }` - /// so the caller can skip the wake and mark the agent idle. - /// - Filters out messages where `from_agent_id == slot_id` (prevent self-trigger). - /// - /// Messages are **not** marked as read here. The caller is responsible for - /// calling `mailbox.mark_read_batch` after successful delivery. - pub async fn compute_wake_input(&self, slot_id: &str) -> Result, TeamError> { + pub(crate) async fn prepare_next_batch(&self, slot_id: &str) -> Result { let agent = self.scheduler.get_agent(slot_id).await?; - let all_unread = self.mailbox.peek_unread(&self.team.id, slot_id).await?; - let next_wake = self.team_run_manager.peek_next_pending_wake(slot_id).await; - // Filter out self-messages to prevent an agent from triggering itself. - // User interventions consume only the triggering foreground message in - // this turn, leaving background backlog unread for follow-up drain. - let unread: Vec<_> = if matches!( - next_wake.as_ref().map(|wake| wake.source), - Some(TeamWakeSource::UserIntervention) - ) { - match next_wake.as_ref().and_then(|wake| wake.message_id.as_ref()) { - Some(message_id) => all_unread - .into_iter() - .filter(|m| m.id == *message_id && m.from_agent_id != slot_id) - .collect(), - None => all_unread.into_iter().filter(|m| m.from_agent_id == "user").collect(), - } - } else { - all_unread.into_iter().filter(|m| m.from_agent_id != slot_id).collect() + let runtime_constraint = match self.member_runtimes.snapshot(slot_id) { + MemberRuntimeSnapshot::Absent if self.event_loops.has(slot_id) => RuntimeConstraint::Ready, + MemberRuntimeSnapshot::Absent => RuntimeConstraint::Starting { operation_id: 0 }, + MemberRuntimeSnapshot::Attaching { operation_id } => RuntimeConstraint::Starting { operation_id }, + MemberRuntimeSnapshot::Ready => RuntimeConstraint::Ready, + MemberRuntimeSnapshot::Failed { operation_id, failure } => RuntimeConstraint::Failed { + operation_id, + classification: failure.classification, + }, + MemberRuntimeSnapshot::Removing { operation_id } => RuntimeConstraint::Removing { operation_id }, + MemberRuntimeSnapshot::SessionStopped => RuntimeConstraint::SessionStopped, }; - let active_team_run_id = self.team_run_manager.active_run_id().await; - if !unread.is_empty() && (active_team_run_id.is_none() || next_wake.is_none()) { - warn!( - team_id = %self.team.id, - slot_id, - unread_count = unread.len(), - has_active_team_run = active_team_run_id.is_some(), - has_pending_wake = next_wake.is_some(), - reason = "unowned_mailbox_backlog", - "team wake input skipped because unread mailbox is not owned by a TeamRun wake" - ); - return Ok(None); + let update = self + .work_coordinator + .set_runtime_constraint(slot_id, runtime_constraint); + if !update.terminal_message_ids.is_empty() { + self.mailbox.mark_read_batch(&update.terminal_message_ids).await?; } - let tasks = self.scheduler.list_tasks().await?; - let mut wake_body = build_wake_payload(&agent, &tasks, &unread); - if matches!( - next_wake.as_ref().map(|wake| wake.source), - Some(TeamWakeSource::UserIntervention) - ) { - wake_body = format!( - "## Turn Context\n\nThis turn was triggered by a user intervention. Prioritize the user's message in this turn. Do not infer exact queue or backlog state from this notice.\n\n{wake_body}" - ); - } - - let needs_role_prompt = self.scheduler.take_needs_role_prompt(slot_id).await; - - let first_message = if needs_role_prompt { - let role_prompt = match agent.role { - TeammateRole::Lead => { - let available_assistants = match self.service.upgrade() { - Some(svc) => svc.list_team_selectable_assistants().await, - None => Vec::new(), + let unread = self + .mailbox + .peek_unread(&self.team.id, slot_id) + .await? + .into_iter() + .filter(|message| message.from_agent_id != slot_id) + .collect::>(); + let unread_ids = unread.iter().map(|message| message.id.clone()).collect::>(); + self.work_coordinator + .reconcile_mailbox(slot_id, &unread_ids, target_role_for(agent.role)); + + match self.work_coordinator.next(slot_id) { + ReconcileDecision::Claim(batch) => { + let claimed_ids = batch + .mailbox_message_ids + .iter() + .map(String::as_str) + .collect::>(); + let claimed_unread = unread + .into_iter() + .filter(|message| claimed_ids.contains(message.id.as_str())) + .collect::>(); + let tasks = match self.scheduler.list_tasks().await { + Ok(tasks) => tasks, + Err(error) => { + self.work_coordinator.retry_start(&batch, "batch_prepare_failed"); + return Err(error); + } + }; + let current_slot_ids = self + .scheduler + .list_agents() + .await + .into_iter() + .map(|member| member.slot_id) + .collect(); + let wake_body = build_wake_payload(&agent, &tasks, &claimed_unread, ¤t_slot_ids); + let needs_role_prompt = self.scheduler.take_needs_role_prompt(slot_id).await; + let first_message = if needs_role_prompt { + let role_prompt = match agent.role { + TeammateRole::Lead => { + build_lead_prompt(&agent, &self.team.name, &self.scheduler.list_agents().await, &[]) + } + TeammateRole::Teammate => { + let members = self.scheduler.list_agents().await; + build_teammate_prompt(&agent, &self.team.name, &members) + } }; - build_lead_prompt( - &self.team.name, - &self.scheduler.list_agents().await, - &available_assistants, - ) - } - TeammateRole::Teammate => { - let members = self.scheduler.list_agents().await; - build_teammate_prompt(&agent, &self.team.name, &members) + format!("{role_prompt}\n\n{wake_body}") + } else { + wake_body + }; + + match dump_team_wake_prompt( + &self.prompt_dump, + TeamWakePromptDump { + team_id: &self.team.id, + slot_id, + conversation_id: &agent.conversation_id, + role: agent.role, + needs_role_prompt, + unread_count: claimed_unread.len(), + prompt: &first_message, + }, + ) { + Ok(Some(path)) => tracing::debug!( + team_id = %self.team.id, + slot_id, + path = %path.display(), + "team wake prompt dump written" + ), + Ok(None) => {} + Err(error) => warn!( + team_id = %self.team.id, + slot_id, + error = %error, + "team wake prompt dump failed" + ), } - }; - format!("{role_prompt}\n\n{wake_body}") - } else { - wake_body - }; - let should_send = !unread.is_empty(); + Ok(PrepareBatchResult::Execute { + batch: Box::new(batch), + input: WakeInput { + conversation_id: agent.conversation_id, + first_message, + unread: claimed_unread, + agent_role: agent.role, + }, + }) + } + ReconcileDecision::SettleSignals(intent_ids) => Ok(PrepareBatchResult::SettleSignals { intent_ids }), + ReconcileDecision::WaitingForCompletion => Ok(PrepareBatchResult::WaitingForCompletion), + ReconcileDecision::Blocked(_) => Ok(PrepareBatchResult::Blocked), + ReconcileDecision::Quiescent => Ok(PrepareBatchResult::Quiescent), + } + } - Ok(Some(WakeInput { - team_run_id: active_team_run_id, - conversation_id: agent.conversation_id, - first_message, - should_send, - unread, - agent_role: agent.role, - wake_source: next_wake.as_ref().map(|wake| wake.source), - trigger_message_id: next_wake.and_then(|wake| wake.message_id), - })) + pub(crate) async fn handle_signal_intents(&self, slot_id: &str, intent_ids: &[String]) { + if self.work_coordinator.complete_signals(slot_id, intent_ids) != CommitResult::Committed { + warn!( + team_id = %self.team.id, + slot_id, + intent_count = intent_ids.len(), + "team signal intents could not be settled" + ); + } } /// Handle agent Finish/Error events. Delegates to the scheduler's @@ -385,67 +484,10 @@ impl TeamSession { .find_lead_slot_id() .await .ok_or_else(|| TeamError::AgentNotFound("no lead agent in team".into()))?; - - let lead_conv_id = self.scheduler.get_agent(&lead_slot_id).await?.conversation_id; - let (mut ack, lease) = self - .team_run_manager - .acquire_user_message_wake(&lead_slot_id, TeamRunTargetRole::Lead) - .await?; - - let mailbox_message = match self - .mailbox - .write_with_files( - &self.team.id, - &lead_slot_id, - "user", - MailboxMessageType::Message, - content, - None, - files.as_deref(), - ) + self.enqueue_user_message(&lead_slot_id, TeamRunTargetRole::Lead, content, files) .await - { - Ok(message) => message, - Err(err) => { - let _ = self - .team_run_manager - .abort_operation_lease(&lease.lease_id, "mailbox_write_failed") - .await; - return Err(err); - } - }; - ack.message_id = Some(mailbox_message.id.clone()); - - let projection = TeamMessageProjection::new(self.projection_store.clone(), self.broadcaster.clone()); - let request = TeamProjectionRequest::user_visible( - &self.team.id, - &lead_slot_id, - &lead_conv_id, - content, - files.clone().unwrap_or_default(), - ); - if let Err(e) = projection.project(request).await { - warn!( - team_id = %self.team.id, - slot_id = %lead_slot_id, - conversation_id = %lead_conv_id, - error = %e, - "failed to project user right bubble for leader (non-fatal)" - ); - } - - let _ = files; - self.team_run_manager - .commit_operation_lease(&lease.lease_id, Some(mailbox_message.id.clone())) - .await?; - self.notify_reserved_wake_for_team_work(&lead_slot_id, lease.role.clone(), lease.wake_source); - Ok(ack) } - /// Write a user message to the specified agent's mailbox and trigger a wake. - /// - /// Same log-not-throw behaviour as [`send_message`]; see that method for - /// rationale. pub async fn send_message_to_agent( &self, slot_id: &str, @@ -453,11 +495,30 @@ impl TeamSession { files: Option>, ) -> Result { let agent = self.scheduler.get_agent(slot_id).await?; - let (mut ack, lease) = self - .team_run_manager - .acquire_user_message_wake(slot_id, target_role_for(agent.role)) - .await?; + self.enqueue_user_message(slot_id, target_role_for(agent.role), content, files) + .await + } + async fn enqueue_user_message( + &self, + slot_id: &str, + role: TeamRunTargetRole, + content: &str, + files: Option>, + ) -> Result { + self.publish_runtime_constraint(slot_id).await?; + let agent = self.scheduler.get_agent(slot_id).await?; + let source = if self.team_run_manager.current_active_run_id().is_some() { + WorkSource::UserIntervention + } else { + WorkSource::UserMessage + }; + let lease = self.work_coordinator.acquire_enqueue(EnqueueRequest { + slot_id: slot_id.to_owned(), + role, + source, + binding: CausalBinding::UserVisible, + })?; let mailbox_message = match self .mailbox .write_with_files( @@ -472,15 +533,11 @@ impl TeamSession { .await { Ok(message) => message, - Err(err) => { - let _ = self - .team_run_manager - .abort_operation_lease(&lease.lease_id, "mailbox_write_failed") - .await; - return Err(err); + Err(error) => { + self.work_coordinator.abort_enqueue(&lease, "mailbox_write_failed"); + return Err(error); } }; - ack.message_id = Some(mailbox_message.id.clone()); let projection = TeamMessageProjection::new(self.projection_store.clone(), self.broadcaster.clone()); let request = TeamProjectionRequest::user_visible( @@ -488,24 +545,56 @@ impl TeamSession { slot_id, &agent.conversation_id, content, - files.clone().unwrap_or_default(), + files.unwrap_or_default(), ); - if let Err(e) = projection.project(request).await { + if let Err(error) = projection.project(request).await { warn!( team_id = %self.team.id, slot_id, conversation_id = %agent.conversation_id, - error = %e, + error = %error, "failed to project user right bubble (non-fatal)" ); } - let _ = files; - self.team_run_manager - .commit_operation_lease(&lease.lease_id, Some(mailbox_message.id.clone())) + let commit = self + .commit_persisted_enqueue(&lease, mailbox_message.id.clone()) .await?; - self.notify_reserved_wake_for_team_work(slot_id, lease.role.clone(), lease.wake_source); - Ok(ack) + self.event_loops.notify(slot_id); + let snapshot = self.work_coordinator.snapshot(); + let run = self + .team_run_manager + .current_payload(&snapshot) + .ok_or_else(|| TeamError::InvalidRequest("user-visible enqueue lost its team run".into()))?; + Ok(TeamRunAckResponse { + enqueue_status: match commit.disposition { + EnqueueDisposition::Accepted => TeamMessageEnqueueStatus::Accepted, + EnqueueDisposition::Queued => TeamMessageEnqueueStatus::Queued, + EnqueueDisposition::BlockedRuntimeStarting => TeamMessageEnqueueStatus::BlockedRuntimeStarting, + }, + message_id: mailbox_message.id, + run, + }) + } + + async fn publish_runtime_constraint(&self, slot_id: &str) -> Result<(), TeamError> { + let constraint = match self.member_runtimes.snapshot(slot_id) { + MemberRuntimeSnapshot::Absent if self.event_loops.has(slot_id) => RuntimeConstraint::Ready, + MemberRuntimeSnapshot::Absent => RuntimeConstraint::Starting { operation_id: 0 }, + MemberRuntimeSnapshot::Attaching { operation_id } => RuntimeConstraint::Starting { operation_id }, + MemberRuntimeSnapshot::Ready => RuntimeConstraint::Ready, + MemberRuntimeSnapshot::Failed { operation_id, failure } => RuntimeConstraint::Failed { + operation_id, + classification: failure.classification, + }, + MemberRuntimeSnapshot::Removing { operation_id } => RuntimeConstraint::Removing { operation_id }, + MemberRuntimeSnapshot::SessionStopped => RuntimeConstraint::SessionStopped, + }; + let update = self.work_coordinator.set_runtime_constraint(slot_id, constraint); + if !update.terminal_message_ids.is_empty() { + self.mailbox.mark_read_batch(&update.terminal_message_ids).await?; + } + Ok(()) } pub(crate) async fn send_agent_message_from_agent( @@ -516,45 +605,15 @@ impl TeamSession { ) -> Result { let to_agent = self.scheduler.get_agent(to_slot_id).await?; let from_agent = self.scheduler.get_agent(from_slot_id).await?; - let to_role = target_role_for(to_agent.role); - let outcome = self - .team_run_manager - .acquire_run_scoped_wake(to_slot_id, to_role.clone(), TeamWakeSource::McpSendMessage) - .await?; - let TeamRunWakeAcquireOutcome::Accepted(lease) = outcome else { - let (team_run_id, work) = self - .team_run_manager - .slot_work_for_slot(to_slot_id) - .await - .ok_or_else(|| TeamError::InvalidRequest(format!("no active team run work for slot {to_slot_id}")))?; - let queue_state = classify_send_message_queue_state(&work, true); - let target = TeamSendMessageTargetQueueState { - slot_id: work.slot_id, - role: work.role, - queue_state, - pending_wake_count: work.pending_wake_count, - starting_child_count: work.starting_child_count, - active_turn_id: work.active_turn_id, - suppressed_wake_count: work.suppressed_wake_count, - }; - info!( - team_id = %self.team.id, - team_run_id = %team_run_id, - caller_slot_id = from_slot_id, - target_slot_id = to_slot_id, - target_role = ?target.role, - wake_source = %TeamWakeSource::McpSendMessage, - suppressed_wake_count = target.suppressed_wake_count, - reason = ?target.queue_state, - "team_agent_message_wake_suppressed" - ); - return Ok(AgentMessageQueueResult { - team_run_id, - delivery: TeamSendMessageDelivery::WakeSuppressed, - target, - }); - }; - + self.publish_runtime_constraint(to_slot_id).await?; + let lease = self.work_coordinator.acquire_enqueue(EnqueueRequest { + slot_id: to_slot_id.to_owned(), + role: target_role_for(to_agent.role), + source: WorkSource::McpSendMessage, + binding: CausalBinding::InheritRunningBatch { + caller_slot_id: from_slot_id.to_owned(), + }, + })?; let mailbox_message = match self .mailbox .write( @@ -568,12 +627,9 @@ impl TeamSession { .await { Ok(message) => message, - Err(err) => { - let _ = self - .team_run_manager - .abort_operation_lease(&lease.lease_id, "mailbox_write_failed") - .await; - return Err(err); + Err(error) => { + self.work_coordinator.abort_enqueue(&lease, "mailbox_write_failed"); + return Err(error); } }; @@ -584,9 +640,9 @@ impl TeamSession { conversation_id: to_agent.conversation_id.clone(), source: TeamProjectionSource::Teammate { from_slot_id: from_slot_id.to_owned(), - from_name: from_agent.name.clone(), - sender_backend: Some(from_agent.backend.clone()), - sender_conversation_id: Some(from_agent.conversation_id.clone()), + from_name: from_agent.name, + sender_backend: Some(from_agent.backend), + sender_conversation_id: Some(from_agent.conversation_id), }, content: content.to_owned(), files: Vec::new(), @@ -597,76 +653,22 @@ impl TeamSession { &to_agent.conversation_id, )), }; - - if let Err(err) = projection.project(request).await { + if let Err(error) = projection.project(request).await { warn!( team_id = %self.team.id, from_slot_id, to_slot_id, - conversation_id = %to_agent.conversation_id, mailbox_message_id = %mailbox_message.id, - error = %err, - "team agent message immediate projection failed (non-fatal)" - ); - } - - self.team_run_manager - .commit_operation_lease(&lease.lease_id, Some(mailbox_message.id.clone())) - .await?; - let target_event_loop_registered = self.event_loops.has(to_slot_id); - if !target_event_loop_registered { - warn!( - team_id = %self.team.id, - slot_id = to_slot_id, - target_role = ?to_role, - wake_source = %TeamWakeSource::McpSendMessage, - "team wake recorded but event loop is not registered; pending wake retained" + error = %error, + "team agent message projection failed (non-fatal)" ); - self.team_run_manager - .mark_slot_runtime_health(to_slot_id, TeamSlotRuntimeHealth::Unhealthy) - .await; - } - - let (team_run_id, work) = self - .team_run_manager - .slot_work_for_slot(to_slot_id) - .await - .ok_or_else(|| TeamError::InvalidRequest(format!("no active team run work for slot {to_slot_id}")))?; - let queue_state = classify_send_message_queue_state(&work, target_event_loop_registered); - let target = TeamSendMessageTargetQueueState { - slot_id: work.slot_id, - role: work.role, - queue_state, - pending_wake_count: work.pending_wake_count, - starting_child_count: work.starting_child_count, - active_turn_id: work.active_turn_id, - suppressed_wake_count: work.suppressed_wake_count, - }; - - info!( - team_id = %self.team.id, - team_run_id = %team_run_id, - caller_slot_id = from_slot_id, - target_slot_id = to_slot_id, - target_role = ?target.role, - wake_source = %TeamWakeSource::McpSendMessage, - message_id = %mailbox_message.id, - slot_pending_wake_count = target.pending_wake_count, - starting_child_count = target.starting_child_count, - active_turn_id = ?target.active_turn_id.as_deref(), - suppressed_wake_count = target.suppressed_wake_count, - reason = ?target.queue_state, - "team_agent_message_queued" - ); - - if target_event_loop_registered { - self.notify_reserved_wake_for_team_work(to_slot_id, to_role, TeamWakeSource::McpSendMessage); } + let commit = self.commit_persisted_enqueue(&lease, mailbox_message.id).await?; + self.event_loops.notify(to_slot_id); Ok(AgentMessageQueueResult { - team_run_id, - delivery: TeamSendMessageDelivery::WakeRecorded, - target, + team_run_id: commit.team_run_id, + target: TeamRunManager::slot_payload(&commit.slot), }) } @@ -684,251 +686,151 @@ impl TeamSession { if target.role == TeammateRole::Lead { return Err(TeamError::InvalidRequest("cannot shutdown the team lead".into())); } - - let outcome = self - .team_run_manager - .acquire_run_scoped_wake( - target_slot_id, - target_role_for(target.role), - TeamWakeSource::McpShutdownRequest, - ) - .await?; - let TeamRunWakeAcquireOutcome::Accepted(lease) = outcome else { - return Err(TeamError::InvalidRequest("shutdown wake was suppressed".into())); - }; - + self.publish_runtime_constraint(target_slot_id).await?; + let lease = self.work_coordinator.acquire_enqueue(EnqueueRequest { + slot_id: target_slot_id.to_owned(), + role: target_role_for(target.role), + source: WorkSource::McpShutdownRequest, + binding: CausalBinding::InheritRunningBatch { + caller_slot_id: caller_slot_id.to_owned(), + }, + })?; let shutdown_message = match self .scheduler .request_shutdown_agent(caller_slot_id, target_slot_id, reason.as_deref()) .await { Ok(message) => message, - Err(err) => { - let _ = self - .team_run_manager - .abort_operation_lease(&lease.lease_id, "shutdown_scheduler_action_failed") - .await; - return Err(err); + Err(error) => { + self.work_coordinator + .abort_enqueue(&lease, "shutdown_scheduler_action_failed"); + return Err(error); } }; - if shutdown_message.to_agent_id != target_slot_id { - let _ = self - .team_run_manager - .abort_operation_lease(&lease.lease_id, "shutdown_mailbox_target_mismatch") - .await; + self.work_coordinator + .abort_enqueue(&lease, "shutdown_mailbox_target_mismatch"); return Err(TeamError::InvalidRequest(format!( "shutdown mailbox target mismatch: expected {target_slot_id}, got {}", shutdown_message.to_agent_id ))); } - - self.team_run_manager - .commit_operation_lease(&lease.lease_id, Some(shutdown_message.id.clone())) - .await?; - self.notify_reserved_wake_for_team_work(target_slot_id, lease.role.clone(), lease.wake_source); - Ok(()) - } - - #[cfg(test)] - pub(crate) async fn wake_agent_for_team_work( - &self, - slot_id: &str, - source: TeamWakeSource, - trigger_message_id: Option, - ) -> Result<(), TeamError> { - let Some(target_role) = self - .reserve_wake_for_team_work(slot_id, source, trigger_message_id) - .await? - else { - return Ok(()); - }; - - if self.event_loops.has(slot_id) { - self.notify_reserved_wake_for_team_work(slot_id, target_role, source); - return Ok(()); - } - - warn!( - team_id = %self.team.id, - slot_id, - target_role = ?target_role, - wake_source = %source, - "team wake recorded but event loop is not registered; pending wake retained" - ); + self.commit_persisted_enqueue(&lease, shutdown_message.id).await?; + self.event_loops.notify(target_slot_id); Ok(()) } - pub(crate) fn notify_mailbox_only_wake(&self, slot_id: &str, source: TeamWakeSource) { - if self.event_loops.has(slot_id) { - self.event_loops.notify(slot_id); - info!( - team_id = %self.team.id, - slot_id, - wake_source = %source, - wake_policy = "mailbox_only", - "team mailbox-only wake notified" - ); - return; - } - - info!( - team_id = %self.team.id, - slot_id, - wake_source = %source, - wake_policy = "mailbox_only_deferred", - "team mailbox-only wake deferred because event loop is not registered" - ); - } - - pub(crate) async fn reserve_wake_for_team_work( + async fn enqueue_existing_work( &self, slot_id: &str, - source: TeamWakeSource, - trigger_message_id: Option, - ) -> Result, TeamError> { + source: WorkSource, + mailbox_message_id: Option, + binding: CausalBinding, + ) -> Result, TeamError> { let agent = self.scheduler.get_agent(slot_id).await?; - let target_role = target_role_for(agent.role); - let decision = self - .team_run_manager - .record_or_suppress_wake(slot_id, target_role.clone(), source, trigger_message_id) - .await?; - Ok(matches!(decision, WakeRecordDecision::Recorded).then_some(target_role)) + self.publish_runtime_constraint(slot_id).await?; + let lease = self.work_coordinator.acquire_enqueue(EnqueueRequest { + slot_id: slot_id.to_owned(), + role: target_role_for(agent.role), + source, + binding, + })?; + let commit = match mailbox_message_id { + Some(message_id) => self.commit_persisted_enqueue(&lease, message_id).await?, + None => match self.work_coordinator.commit_enqueue(&lease, None) { + Ok(commit) => commit, + Err(error) => { + self.work_coordinator.abort_enqueue(&lease, "enqueue_commit_failed"); + return Err(error); + } + }, + }; + self.event_loops.notify(slot_id); + Ok(commit.team_run_id) } - pub(crate) fn notify_reserved_wake_for_team_work( + async fn commit_persisted_enqueue( &self, - slot_id: &str, - target_role: TeamRunTargetRole, - source: TeamWakeSource, - ) { - if self.event_loops.has(slot_id) { - self.event_loops.notify(slot_id); - info!( - team_id = %self.team.id, - slot_id, - target_role = ?target_role, - wake_source = %source, - "team reserved wake notified" - ); - return; + lease: &EnqueueLease, + mailbox_message_id: String, + ) -> Result { + match self + .work_coordinator + .commit_enqueue(lease, Some(mailbox_message_id.clone())) + { + Ok(commit) => Ok(commit), + Err(error) => { + self.work_coordinator.abort_enqueue(lease, "enqueue_commit_failed"); + if let Err(mark_read_error) = self + .mailbox + .mark_read_batch(std::slice::from_ref(&mailbox_message_id)) + .await + { + warn!( + team_id = %self.team.id, + slot_id = %lease.slot_id, + mailbox_message_id, + error_classification = "enqueue_commit_compensation_failed", + error = %mark_read_error, + "persisted team message could not be settled after enqueue rejection" + ); + } + Err(error) + } } - - warn!( - team_id = %self.team.id, - slot_id, - target_role = ?target_role, - wake_source = %source, - "team reserved wake notification skipped because event loop is not registered" - ); } - pub(crate) async fn scheduler_wake_agent_for_team_work( + pub(crate) async fn enqueue_leader_settle_signal( &self, slot_id: &str, - source: TeamWakeSource, + source: WorkSource, ) -> Result<(), TeamError> { - let agent = self.scheduler.get_agent(slot_id).await?; - let target_role = target_role_for(agent.role); - let outcome = self - .team_run_manager - .acquire_scheduler_wake(slot_id, target_role.clone(), source) - .await?; - let TeamRunWakeAcquireOutcome::Accepted(lease) = outcome else { - return Ok(()); - }; - - self.team_run_manager - .commit_operation_lease(&lease.lease_id, None) + self.enqueue_existing_work(slot_id, source, None, CausalBinding::ActiveRunOrBackground) .await?; - self.notify_reserved_wake_for_team_work(slot_id, target_role, source); Ok(()) } - pub(crate) async fn try_start_recovery_drain( - &self, - reason: &'static str, - ) -> Result, TeamError> { + pub(crate) async fn try_start_recovery_drain(&self, reason: &'static str) -> Result, TeamError> { if self.recovery_scan_completed.swap(true, Ordering::AcqRel) { tracing::debug!( team_id = %self.team.id, reason, "team recovery scan skipped because it already ran for this session lifecycle" ); - return Ok(None); + return Ok(Vec::new()); } - let agents = self.scheduler.list_agents().await; - let mut candidates = Vec::new(); - let mut unread_total = 0usize; - let mut missing_event_loops = Vec::new(); - - for agent in agents { - if !self.event_loops.has(&agent.slot_id) { - missing_event_loops.push(agent.slot_id.clone()); + let mut recovered_slots = Vec::new(); + for agent in self.scheduler.list_agents().await { + let unread = self + .mailbox + .peek_unread(&self.team.id, &agent.slot_id) + .await? + .into_iter() + .filter(|message| message.from_agent_id != agent.slot_id) + .map(|message| message.id) + .collect::>(); + if unread.is_empty() { continue; } - - let unread = self.mailbox.peek_unread(&self.team.id, &agent.slot_id).await?; - let recoverable_count = unread - .into_iter() - .filter(|message| message.from_agent_id != agent.slot_id) - .count(); - if recoverable_count == 0 { - continue; - } - unread_total += recoverable_count; - candidates.push(RecoveryWakeCandidate { - slot_id: agent.slot_id, - role: target_role_for(agent.role), - unread_count: recoverable_count, - }); - } - - if !missing_event_loops.is_empty() { - warn!( - team_id = %self.team.id, - reason, - missing_event_loop_count = missing_event_loops.len(), - unread_count = unread_total, - "team recovery scan found slots without event loops; unread work is retained" - ); - } - - if candidates.is_empty() { - tracing::debug!( - team_id = %self.team.id, - reason, - "team recovery scan found no recoverable mailbox backlog" - ); - return Ok(None); - } - - let result = self.team_run_manager.recover_mailbox_backlog(candidates).await; - if let Some(result) = result.as_ref() { - info!( - team_id = %self.team.id, - team_run_id = %result.team_run_id, - source = "recovery_drain", - slot_count = result.recorded_wakes.len(), - pending_wake_count = result.pending_wake_count, - reason, - "team recovery scan recorded TeamRun wakes" - ); - for slot_id in &result.recorded_wakes { - self.event_loops.notify(slot_id); + for message_id in unread { + self.enqueue_existing_work( + &agent.slot_id, + WorkSource::RecoveryDrain, + Some(message_id), + CausalBinding::Background, + ) + .await?; } - } else { - warn!( - team_id = %self.team.id, - source = "recovery_drain", - unread_count = unread_total, - reason, - "team recovery scan retained unread mailbox backlog without recording wakes" - ); + recovered_slots.push(agent.slot_id); } - - Ok(result) + info!( + team_id = %self.team.id, + recovered_slot_count = recovered_slots.len(), + reason, + "team recovery scan projected background work" + ); + Ok(recovered_slots) } /// Mirror each non-user mailbox row into the target agent's conversation @@ -1001,82 +903,50 @@ impl TeamSession { pub async fn cancel_run( &self, team_run_id: &str, - target_slot_id: Option, + _target_slot_id: Option, reason: Option, ) -> Result<(), TeamError> { - let current_run_id = self - .team_run_manager - .current_run_id() - .await - .ok_or_else(|| TeamError::InvalidRequest("no active team run to cancel".into()))?; - if current_run_id != team_run_id { + if self.team_run_manager.current_active_run_id().as_deref() != Some(team_run_id) { return Err(TeamError::InvalidRequest(format!( "team run {team_run_id} is not active" ))); } - - self.team_run_manager.begin_cancel(target_slot_id, reason).await?; - - let agent_ids = self - .scheduler - .list_agents() - .await - .into_iter() - .map(|agent| agent.slot_id) - .collect::>(); - let marked = self - .mailbox - .mark_all_unread_for_agents_read(&self.team.id, &agent_ids) - .await?; - info!( - team_id = %self.team.id, - team_run_id, - agent_count = agent_ids.len(), - marked_unread = marked, - "team_run cancel drained unread mailbox rows" - ); - - let active_child_turns = self.team_run_manager.active_child_turns().await; - let active_child_count = active_child_turns.len(); - for child in active_child_turns { - if let Err(err) = self + self.team_run_manager.begin_cancel(team_run_id, reason)?; + let result = self.work_coordinator.cancel_run(team_run_id); + if !result.terminal_message_ids.is_empty() { + self.mailbox.mark_read_batch(&result.terminal_message_ids).await?; + } + for target in result.cancel_targets { + let Some(turn_id) = target.turn_id else { + continue; + }; + let agent = self.scheduler.get_agent(&target.batch.slot_id).await?; + if let Err(error) = self .cancellation_port - .cancel_agent_turn(&self.user_id, &child.conversation_id, &child.turn_id) + .cancel_agent_turn(&self.user_id, &agent.conversation_id, &turn_id) .await { warn!( team_id = %self.team.id, team_run_id, - slot_id = %child.slot_id, - turn_id = %child.turn_id, - error = %err, - "team_run cancel child turn failed (continuing)" - ); - } - self.team_run_manager.record_child_cancelled(&child).await; - } - - match self.team_run_manager.try_complete_cancelled().await { - Some(payload) => { - info!( - team_id = %self.team.id, - team_run_id, - active_child_count, - starting_child_count = payload.starting_child_count, - pending_wake_count = payload.pending_wake_count, - marked_unread = marked, - "team_run cancel completed" - ); - } - None => { - info!( - team_id = %self.team.id, - team_run_id, - active_child_count, - marked_unread = marked, - "team_run cancel completion pending" + slot_id = %target.batch.slot_id, + turn_id, + error = %error, + "team run cancellation failed for active turn" ); } + self.team_event_emitter().broadcast_child_turn( + TEAM_CHILD_TURN_CANCELLED_EVENT, + TeamChildTurnPayload { + team_id: self.team.id.clone(), + team_run_id: team_run_id.to_owned(), + slot_id: target.batch.slot_id, + role: target_role_for(agent.role), + conversation_id: agent.conversation_id, + turn_id, + status: TeamRunStatus::Cancelled, + }, + ); } Ok(()) } @@ -1087,37 +957,40 @@ impl TeamSession { slot_id: &str, reason: Option, ) -> Result<(), TeamError> { - let target = self.team_run_manager.begin_cancel_child(slot_id).await?; - match target { - ChildCancelTarget::Active(child) => { - if child.team_run_id != team_run_id { - return Err(TeamError::InvalidRequest(format!( - "agent {slot_id} is not active in team run {team_run_id}" - ))); - } - self.cancellation_port - .cancel_agent_turn(&self.user_id, &child.conversation_id, &child.turn_id) - .await - .map_err(|err| TeamError::InvalidRequest(err.to_string()))?; - self.team_run_manager.record_child_cancelled(&child).await; - - if child.role == TeamRunTargetRole::Teammate { - self.notify_leader_child_interrupted(slot_id, reason).await?; - } else { - self.team_run_manager.maybe_complete().await; - } - } - ChildCancelTarget::Starting(reservation) => { - if reservation.team_run_id != team_run_id { - return Err(TeamError::InvalidRequest(format!( - "agent {slot_id} is not starting in team run {team_run_id}" - ))); - } - if reservation.role == TeamRunTargetRole::Teammate { - self.notify_leader_child_interrupted(slot_id, reason).await?; - } - self.team_run_manager.try_complete_cancelled().await; - } + let snapshot = self + .work_coordinator + .slot_snapshot(slot_id) + .ok_or_else(|| TeamError::AgentNotFound(slot_id.to_owned()))?; + let batch = snapshot + .active_batch + .ok_or_else(|| TeamError::InvalidRequest(format!("slot has no active batch: {slot_id}")))?; + if !batch.team_run_ids.iter().any(|run_id| run_id == team_run_id) { + return Err(TeamError::InvalidRequest(format!( + "agent {slot_id} is not active in team run {team_run_id}" + ))); + } + if let Some(turn_id) = snapshot.active_turn_id { + let agent = self.scheduler.get_agent(slot_id).await?; + self.cancellation_port + .cancel_agent_turn(&self.user_id, &agent.conversation_id, &turn_id) + .await + .map_err(|error| TeamError::InvalidRequest(error.to_string()))?; + self.team_event_emitter().broadcast_child_turn( + TEAM_CHILD_TURN_CANCELLED_EVENT, + TeamChildTurnPayload { + team_id: self.team.id.clone(), + team_run_id: team_run_id.to_owned(), + slot_id: slot_id.to_owned(), + role: snapshot.role.clone(), + conversation_id: agent.conversation_id, + turn_id, + status: TeamRunStatus::Cancelled, + }, + ); + } + self.work_coordinator.cancel_batch(&batch, "child_cancelled"); + if snapshot.role == TeamRunTargetRole::Teammate { + self.notify_leader_child_interrupted(slot_id, reason).await?; } Ok(()) } @@ -1128,49 +1001,35 @@ impl TeamSession { slot_id: &str, reason: Option, ) -> Result<(), TeamError> { - let current_run_id = self - .team_run_manager - .current_run_id() - .await - .ok_or_else(|| TeamError::InvalidRequest("no active team run to pause".into()))?; - if current_run_id != team_run_id { + if self.team_run_manager.current_active_run_id().as_deref() != Some(team_run_id) { return Err(TeamError::InvalidRequest(format!( - "agent {slot_id} is not active in team run {team_run_id}" + "team run {team_run_id} is not active" ))); } - - let outcome = self.team_run_manager.pause_slot_work(slot_id, reason.clone()).await?; - + let outcome = self.work_coordinator.pause_slot(slot_id); if let Some(target) = outcome.cancel_target { - match target { - ChildCancelTarget::Active(child) => { - if let Err(err) = self - .cancellation_port - .cancel_agent_turn(&self.user_id, &child.conversation_id, &child.turn_id) - .await - { - warn!( - team_id = %self.team.id, - team_run_id = %outcome.team_run_id, - slot_id = %child.slot_id, - turn_id = %child.turn_id, - error = %err, - "team slot pause child turn cancel failed" - ); - return Err(TeamError::InvalidRequest(err.to_string())); - } - self.team_run_manager - .complete_pause_after_child_cancelled(&child, reason.clone()) - .await; - if child.role == TeamRunTargetRole::Teammate { - self.notify_leader_child_interrupted(slot_id, reason).await?; - } - } - ChildCancelTarget::Starting(reservation) => { - if reservation.role == TeamRunTargetRole::Teammate { - self.notify_leader_child_interrupted(slot_id, reason).await?; - } - } + if let Some(turn_id) = target.turn_id { + let agent = self.scheduler.get_agent(slot_id).await?; + self.cancellation_port + .cancel_agent_turn(&self.user_id, &agent.conversation_id, &turn_id) + .await + .map_err(|error| TeamError::InvalidRequest(error.to_string()))?; + self.team_event_emitter().broadcast_child_turn( + TEAM_CHILD_TURN_CANCELLED_EVENT, + TeamChildTurnPayload { + team_id: self.team.id.clone(), + team_run_id: team_run_id.to_owned(), + slot_id: slot_id.to_owned(), + role: outcome.slot.role.clone(), + conversation_id: agent.conversation_id, + turn_id, + status: TeamRunStatus::Cancelled, + }, + ); + } + self.work_coordinator.cancel_batch(&target.batch, "slot_paused"); + if outcome.slot.role == TeamRunTargetRole::Teammate { + self.notify_leader_child_interrupted(slot_id, reason).await?; } } Ok(()) @@ -1189,7 +1048,7 @@ impl TeamSession { Some("Interrupted by user"), ) .await?; - self.wake_leader_after_recovery_message(slot_id, TeamWakeSource::InterruptedNotification) + self.wake_leader_after_recovery_message(slot_id, WorkSource::InterruptedNotification) .await?; } Ok(()) @@ -1214,46 +1073,40 @@ impl TeamSession { None, ) .await?; - self.wake_leader_after_recovery_message(failed_slot_id, TeamWakeSource::SpawnAttachFailure) + self.wake_leader_after_recovery_message(failed_slot_id, WorkSource::SpawnAttachFailure) .await } pub(crate) async fn wake_leader_after_recovery_message( &self, source_slot_id: &str, - source: TeamWakeSource, + source: WorkSource, ) -> Result<(), TeamError> { - let Some(lead_slot_id) = self.scheduler.find_lead_slot_id().await else { - return Err(TeamError::AgentNotFound("lead".into())); - }; - let outcome = match self - .team_run_manager - .acquire_run_scoped_wake(&lead_slot_id, TeamRunTargetRole::Lead, source) + let lead_slot_id = self + .scheduler + .find_lead_slot_id() .await - { - Ok(outcome) => outcome, - Err(TeamError::InvalidRequest(message)) if message == "no active team run for run-scoped wake" => { - info!( - team_id = %self.team.id, - slot_id = %lead_slot_id, - source_slot_id, - wake_source = %source, - wake_policy = "deferred_mailbox_only", - "leader recovery message deferred because no active team run exists" - ); - return Ok(()); - } - Err(err) => return Err(err), - }; - - let TeamRunWakeAcquireOutcome::Accepted(lease) = outcome else { + .ok_or_else(|| TeamError::AgentNotFound("lead".into()))?; + let message_id = self + .mailbox + .peek_unread(&self.team.id, &lead_slot_id) + .await? + .into_iter() + .rev() + .find(|message| message.from_agent_id == source_slot_id) + .map(|message| message.id); + let Some(message_id) = message_id else { return Ok(()); }; - - self.team_run_manager - .commit_operation_lease(&lease.lease_id, None) - .await?; - self.notify_reserved_wake_for_team_work(&lead_slot_id, TeamRunTargetRole::Lead, source); + self.enqueue_existing_work( + &lead_slot_id, + source, + Some(message_id), + CausalBinding::InheritRunningBatch { + caller_slot_id: source_slot_id.to_owned(), + }, + ) + .await?; Ok(()) } @@ -1261,20 +1114,179 @@ impl TeamSession { self.scheduler.add_agent(agent).await; } - pub async fn remove_agent(&self, slot_id: &str) -> Result<(), TeamError> { - self.event_loops.remove(slot_id); - let conversation_id = self.scheduler.remove_agent(slot_id).await?; - if let Some(conv_id) = conversation_id - && let Err(e) = self.task_manager.kill(&conv_id, Some(AgentKillReason::TeamDeleted)) - { + /// Reserve the runtime lifecycle before a newly persisted member becomes + /// visible through the scheduler or membership events. + pub(crate) fn reserve_dynamic_member_attach(&self, agent: &TeamAgent) -> ReserveAttach { + self.member_runtimes.reserve_attach(&agent.slot_id, false) + } + + pub(crate) async fn add_manual_agent(&self, agent: &TeamAgent) -> Result<(), TeamError> { + self.scheduler.add_agent(agent).await; + let lead_slot_id = self + .scheduler + .find_lead_slot_id() + .await + .ok_or_else(|| TeamError::InvalidRequest("team lead not found".into()))?; + let welcome = self + .mailbox + .write( + &self.team.id, + &agent.slot_id, + "team_system", + MailboxMessageType::Message, + "You were manually added to this team. Read your mailbox and wait for instructions.", + None, + ) + .await?; + let leader_notice = self + .mailbox + .write( + &self.team.id, + &lead_slot_id, + "team_system", + MailboxMessageType::Message, + &format!( + "The user manually added teammate '{}' (slot_id={}, assistant_id={}, model={}). Call team_members to get the latest roster before delegating work.", + agent.name, + agent.slot_id, + agent.assistant_id.as_deref().unwrap_or(""), + agent.model + ), + None, + ) + .await?; + self.project_team_system_message(&agent.slot_id, &agent.conversation_id, &welcome.id, &welcome.content) + .await; + let lead_agent = self.scheduler.get_agent(&lead_slot_id).await?; + self.project_team_system_message( + &lead_slot_id, + &lead_agent.conversation_id, + &leader_notice.id, + &leader_notice.content, + ) + .await; + + self.enqueue_existing_work( + &agent.slot_id, + WorkSource::SpawnWelcome, + Some(welcome.id), + CausalBinding::ActiveRunOrBackground, + ) + .await?; + self.enqueue_existing_work( + &lead_slot_id, + WorkSource::TeamMembershipChanged, + Some(leader_notice.id), + CausalBinding::ActiveRunOrBackground, + ) + .await?; + Ok(()) + } + + pub async fn notify_leader_membership_removed(&self, removed: &TeamAgent) -> Result<(), TeamError> { + let lead_slot_id = self + .scheduler + .find_lead_slot_id() + .await + .ok_or_else(|| TeamError::InvalidRequest("team lead not found".into()))?; + let notice = self + .mailbox + .write( + &self.team.id, + &lead_slot_id, + "team_system", + MailboxMessageType::Message, + &format!( + "Teammate '{}' was removed from the team (slot_id={}, assistant_id={}, model={}). Call team_members to get the latest roster before delegating work.", + removed.name, + removed.slot_id, + removed.assistant_id.as_deref().unwrap_or(""), + removed.model + ), + None, + ) + .await?; + let lead_agent = self.scheduler.get_agent(&lead_slot_id).await?; + self.project_team_system_message(&lead_slot_id, &lead_agent.conversation_id, ¬ice.id, ¬ice.content) + .await; + + self.enqueue_existing_work( + &lead_slot_id, + WorkSource::TeamMembershipChanged, + Some(notice.id), + CausalBinding::ActiveRunOrBackground, + ) + .await?; + + Ok(()) + } + + async fn project_team_system_message( + &self, + slot_id: &str, + conversation_id: &str, + mailbox_message_id: &str, + content: &str, + ) { + let projection = TeamMessageProjection::new(self.projection_store.clone(), self.broadcaster.clone()); + let request = TeamProjectionRequest::team_system_visible( + &self.team.id, + slot_id, + conversation_id, + content, + mailbox_message_id, + ); + if let Err(err) = projection.project(request).await { warn!( team_id = %self.team.id, + slot_id = %slot_id, + conversation_id = %conversation_id, + mailbox_message_id = %mailbox_message_id, + error = %err, + "team system message projection failed (non-fatal)" + ); + } + } + + pub async fn remove_agent(&self, slot_id: &str) -> Result<(), TeamError> { + let removal = match self.member_runtimes.begin_remove(slot_id) { + BeginRemove::Join(waiter) => { + let _ = waiter.wait().await; + return Ok(()); + } + removal => removal, + }; + if let BeginRemove::Start(lease) = &removal { + self.work_coordinator.set_runtime_constraint( slot_id, - conversation_id = %conv_id, - error = %e, - "remove_agent: task_manager.kill failed (non-fatal)" + RuntimeConstraint::Removing { + operation_id: lease.operation_id(), + }, ); } + let work = self.work_coordinator.remove_slot(slot_id); + if !work.terminal_message_ids.is_empty() { + self.mailbox.mark_read_batch(&work.terminal_message_ids).await?; + } + if let Some(target) = work.cancel_target + && let Some(turn_id) = target.turn_id + { + let agent = self.scheduler.get_agent(slot_id).await?; + let _ = self + .cancellation_port + .cancel_agent_turn(&self.user_id, &agent.conversation_id, &turn_id) + .await; + } + self.event_loops.remove(slot_id); + let conversation_id = self.scheduler.remove_agent(slot_id).await?; + if let Some(conversation_id) = conversation_id { + self.task_manager + .kill_and_wait(&conversation_id, Some(AgentKillReason::TeamDeleted)) + .await; + } + if let BeginRemove::Start(lease) = removal { + self.member_runtimes.finish_remove(&lease); + } Ok(()) } @@ -1335,24 +1347,12 @@ impl TeamSession { // capability checks. Assistant spawns derive backend from the preset // identity rather than inheriting the caller backend. let (backend, model) = service - .resolve_spawn_backend_and_model( - Some(assistant_id), - req.model.as_deref(), - caller.backend.as_str(), - caller.model.as_str(), - ) + .resolve_spawn_backend_and_model(Some(assistant_id), None, caller.backend.as_str(), caller.model.as_str()) .await?; - // Step 4: backend capability check. Hard whitelist passes immediately; - // otherwise query persisted agent_capabilities for MCP support. - if !crate::capability::TEAM_CAPABLE_BACKENDS.contains(&backend.as_str()) { - let capable = service.is_backend_team_capable(&backend).await; - if !capable { - return Err(TeamError::BackendNotAllowed(backend)); - } - } - - // Step 5: DB side-effects (new conversation + persisted agent slot). + // Step 4: DB side-effects (new conversation + persisted agent slot). + // Assistant-first spawn has already passed the team-selectable catalog + // gate while resolving the backend/model above. let new_slot_id = generate_id(); let new_agent = service .persist_spawned_agent(PersistSpawnedAgentRequest { @@ -1366,6 +1366,8 @@ impl TeamSession { }) .await?; + let reservation = self.reserve_dynamic_member_attach(&new_agent); + // Step 5: attach to the in-memory scheduler so wake-from-lead finds // the new slot immediately. self.scheduler.add_agent(&new_agent).await; @@ -1398,33 +1400,15 @@ impl TeamSession { } }; - let spawn_wake_plan = if self.team_run_manager.active_run_id().await.is_some() { - let spawn_welcome_role = self - .reserve_wake_for_team_work( - &new_agent.slot_id, - TeamWakeSource::SpawnWelcome, - Some(welcome_message.id), - ) - .await? - .ok_or_else(|| TeamError::InvalidRequest("spawn welcome wake was suppressed".into()))?; - info!( - team_id = %self.team.id, - slot_id = %new_agent.slot_id, - target_role = ?spawn_welcome_role, - wake_source = %TeamWakeSource::SpawnWelcome, - "spawn welcome wake reserved before runtime attach" - ); - SpawnWakePlan::RunScoped(spawn_welcome_role) - } else { - info!( - team_id = %self.team.id, - slot_id = %new_agent.slot_id, - wake_source = %TeamWakeSource::SpawnWelcome, - wake_policy = "mailbox_only", - "spawn welcome wake will use mailbox-only delivery because no active team run exists" - ); - SpawnWakePlan::MailboxOnly - }; + self.enqueue_existing_work( + &new_agent.slot_id, + WorkSource::SpawnWelcome, + Some(welcome_message.id), + CausalBinding::InheritRunningBatch { + caller_slot_id: caller_slot_id.to_owned(), + }, + ) + .await?; // Step 7: attach the CLI process and register the finish subscriber // in a background task. This involves spawning the CLI process and @@ -1432,64 +1416,17 @@ impl TeamSession { // time (10-30s). Running it asynchronously ensures `spawn_agent` // returns promptly so the MCP tool call completes without blocking // the leader's connection loop. - { - let team_id = self.team.id.clone(); - let user_id = self.user_id.clone(); - let agent_clone = new_agent.clone(); - let mcp_stdio_cfg = self.mcp_stdio_config(&new_agent.slot_id); - let task_manager = self.task_manager.clone(); - tokio::spawn(async move { - // Push the team MCP stdio config into the new conversation's - // extras, then kill + rebuild the agent task so the freshly - // spawned process boots with the MCP handshake pointing at - // our session. - if let Err(err) = Self::attach_spawned_agent_process_bg( - &service, - &agent_clone, - mcp_stdio_cfg, - &user_id, - &task_manager, - ) - .await - { - warn!( - team_id = %team_id, - slot_id = %agent_clone.slot_id, - error = %err, - "failed to attach spawned agent process; agent is persisted but not yet running" - ); - if let Err(notify_err) = service - .notify_leader_spawn_attach_failed(&team_id, &agent_clone.slot_id, &err.to_string()) - .await - { - warn!( - team_id = %team_id, - slot_id = %agent_clone.slot_id, - error = %notify_err, - "failed to notify leader about spawned agent attach failure" - ); - } - return; - } - - // Register the event loop for the newly spawned agent. - service.register_event_loop(&team_id, &agent_clone.slot_id); - - match spawn_wake_plan { - SpawnWakePlan::RunScoped(spawn_welcome_role) => { - service.notify_reserved_wake_for_team_work( - &team_id, - &agent_clone.slot_id, - spawn_welcome_role, - TeamWakeSource::SpawnWelcome, - ); - } - SpawnWakePlan::MailboxOnly => { - service.notify_mailbox_only_wake(&team_id, &agent_clone.slot_id, TeamWakeSource::SpawnWelcome); - } - } - }); - } + let captured_session = service + .capture_published_session(self) + .ok_or_else(|| TeamError::SessionNotFound(self.team.id.clone()))?; + spawn_attach_agent_process_bg( + service, + captured_session, + self.user_id.clone(), + new_agent.clone(), + self.task_manager.clone(), + reservation, + ); Ok(new_agent) } @@ -1499,7 +1436,7 @@ impl TeamSession { /// /// This is a static helper suitable for use inside `tokio::spawn` (no /// `&self` borrow). The caller passes all necessary context by value. - async fn attach_spawned_agent_process_bg( + async fn attach_spawned_agent_process( service: &TeamSessionService, agent: &TeamAgent, mcp_stdio_cfg: crate::mcp::TeamMcpStdioConfig, @@ -1513,7 +1450,10 @@ impl TeamSession { } pub fn stop(&self) { - info!(team_id = %self.team.id, "TeamSession stopping"); + info!(team_id = %self.team.id, generation = %self.generation(), "TeamSession stopping"); + self.work_coordinator.stop(); + self.member_runtimes.stop(); + self.event_loops.shutdown(); self.mcp_server.stop(); } @@ -1526,43 +1466,292 @@ impl TeamSession { } } -fn classify_send_message_queue_state( - work: &TeamSlotWorkPayload, - target_event_loop_registered: bool, -) -> TeamSendMessageReason { - if matches!(work.runtime_health, Some(TeamSlotRuntimeHealth::Disconnected)) { - return TeamSendMessageReason::TargetDisconnected; - } - if !target_event_loop_registered || matches!(work.runtime_health, Some(TeamSlotRuntimeHealth::Unhealthy)) { - return TeamSendMessageReason::TargetUnhealthy; - } - if work.paused && work.suppressed_wake_count > 0 { - return TeamSendMessageReason::SuppressedByPause; - } - if work.active_turn_id.is_some() { - return TeamSendMessageReason::BehindActiveTurn; - } - if work.starting_child_count > 0 { - return TeamSendMessageReason::BehindStartingTurn; - } - TeamSendMessageReason::QueuedForIdle -} - -#[cfg(test)] -mod tests { - use super::*; +pub(crate) fn spawn_attach_agent_process_bg( + service: Arc, + session: Arc, + user_id: String, + agent: TeamAgent, + task_manager: Arc, + reservation: ReserveAttach, +) { + tokio::spawn(async move { + let outcome = match reservation { + ReserveAttach::Start(lease) => { + attach_member_runtime( + Arc::clone(&service), + Arc::clone(&session), + user_id, + agent.clone(), + task_manager, + lease, + ) + .await + } + ReserveAttach::Join(waiter) | ReserveAttach::Removing(waiter) => { + info!( + team_id = session.team_id(), + slot_id = agent.slot_id, + conversation_id = agent.conversation_id, + operation_id = waiter.operation_id(), + generation = session.generation(), + duration_ms = 0, + error_classification = "none", + "team member runtime attach waiting" + ); + waiter.wait().await + } + ReserveAttach::AlreadyReady => AttachOutcome::Ready, + ReserveAttach::SessionStopped => AttachOutcome::SessionStopped, + }; + + if outcome != AttachOutcome::Ready { + return; + } + + let _ = service.with_published_session(&session, |current| { + current + .work_coordinator + .set_runtime_constraint(&agent.slot_id, RuntimeConstraint::Ready); + current.event_loops.notify(&agent.slot_id); + }); + }); +} + +pub(crate) async fn attach_member_runtime( + service: Arc, + session: Arc, + user_id: String, + agent: TeamAgent, + task_manager: Arc, + lease: AttachLease, +) -> AttachOutcome { + let started_at = Instant::now(); + let operation_id = lease.operation_id(); + let generation = session.generation().to_owned(); + session + .work_coordinator + .set_runtime_constraint(&agent.slot_id, RuntimeConstraint::Starting { operation_id }); + service.publish_member_runtime_starting_if_current(&session); + info!( + team_id = session.team_id(), + slot_id = agent.slot_id, + conversation_id = agent.conversation_id, + operation_id, + generation, + duration_ms = 0, + error_classification = "none", + "team member runtime attach started" + ); + + let attach_result = TeamSession::attach_spawned_agent_process( + &service, + &agent, + session.mcp_stdio_config(&agent.slot_id), + &user_id, + &task_manager, + ) + .await; + if let Err(error) = attach_result { + service + .cleanup_stale_member_runtime_task(&session, &agent.conversation_id) + .await; + let failure = sanitize_member_runtime_failure(&error); + if session.member_runtimes.commit_failed(&lease, failure.clone()) { + service.broadcast_agent_runtime_status( + session.team_id(), + &agent, + TeamAgentRuntimeStatus::Failed, + Some(failure.public_reason.clone()), + ); + let _ = session + .scheduler + .set_status(&agent.slot_id, TeammateStatus::Error) + .await; + let update = session.work_coordinator.set_runtime_constraint( + &agent.slot_id, + RuntimeConstraint::Failed { + operation_id, + classification: failure.classification, + }, + ); + if !update.terminal_message_ids.is_empty() { + let _ = session.mailbox.mark_read_batch(&update.terminal_message_ids).await; + } + service.refresh_member_runtime_status(&session).await; + if let Err(notify_error) = session + .notify_leader_spawn_attach_failed(&agent.slot_id, &failure.public_reason) + .await + { + warn!( + team_id = session.team_id(), + slot_id = agent.slot_id, + conversation_id = agent.conversation_id, + operation_id, + generation, + duration_ms = started_at.elapsed().as_millis(), + error_classification = "leader_notification_failed", + error = %notify_error, + "team member runtime attach failure notification failed" + ); + } + } else { + warn!( + team_id = session.team_id(), + slot_id = agent.slot_id, + conversation_id = agent.conversation_id, + operation_id, + generation, + duration_ms = started_at.elapsed().as_millis(), + error_classification = "stale_attach_failure", + "stale team member runtime attach failure ignored" + ); + } + info!( + team_id = session.team_id(), + slot_id = agent.slot_id, + conversation_id = agent.conversation_id, + operation_id, + generation, + duration_ms = started_at.elapsed().as_millis(), + error_classification = failure.classification, + "team member runtime attach completed" + ); + return lease.waiter().wait().await; + } + + if !session.member_runtimes.owns_attach(&lease) { + cleanup_stale_attach(&service, &session, &agent, operation_id, &generation, started_at).await; + return lease.waiter().wait().await; + } + + let registered_event_loop = if session.event_loops.has(&agent.slot_id) { + false + } else { + match service.register_event_loop(&session, &agent.slot_id) { + Ok(registered) => registered, + Err(error) => { + service + .cleanup_stale_member_runtime_task(&session, &agent.conversation_id) + .await; + let failure = MemberRuntimeFailure { + classification: "event_loop_registration_failed", + public_reason: "Agent runtime event loop could not be registered".to_owned(), + }; + session.member_runtimes.commit_failed(&lease, failure.clone()); + session.work_coordinator.set_runtime_constraint( + &agent.slot_id, + RuntimeConstraint::Failed { + operation_id, + classification: failure.classification, + }, + ); + service.refresh_member_runtime_status(&session).await; + warn!( + team_id = session.team_id(), + slot_id = agent.slot_id, + conversation_id = agent.conversation_id, + operation_id, + generation, + duration_ms = started_at.elapsed().as_millis(), + error_classification = failure.classification, + error = ?error, + "team member runtime attach event loop registration failed" + ); + return lease.waiter().wait().await; + } + } + }; + + if !session.member_runtimes.commit_ready(&lease) { + if registered_event_loop { + session.event_loops.remove(&agent.slot_id); + } + cleanup_stale_attach(&service, &session, &agent, operation_id, &generation, started_at).await; + return lease.waiter().wait().await; + } + session + .work_coordinator + .set_runtime_constraint(&agent.slot_id, RuntimeConstraint::Ready); + session.event_loops.notify(&agent.slot_id); + + if !service.publish_member_runtime_ready_if_current(&session, &agent) { + if registered_event_loop { + session.event_loops.remove(&agent.slot_id); + } + cleanup_stale_attach(&service, &session, &agent, operation_id, &generation, started_at).await; + return AttachOutcome::SessionStopped; + } + service.refresh_member_runtime_status(&session).await; + info!( + team_id = session.team_id(), + slot_id = agent.slot_id, + conversation_id = agent.conversation_id, + operation_id, + generation, + duration_ms = started_at.elapsed().as_millis(), + error_classification = "none", + "team member runtime attach completed" + ); + AttachOutcome::Ready +} + +fn sanitize_member_runtime_failure(error: &TeamError) -> MemberRuntimeFailure { + let classification = match error { + TeamError::Database(_) => "runtime_configuration_failed", + TeamError::SessionNotFound(_) => "session_stopped", + _ => "runtime_start_failed", + }; + MemberRuntimeFailure { + classification, + public_reason: "Agent runtime failed to start".to_owned(), + } +} + +async fn cleanup_stale_attach( + service: &TeamSessionService, + session: &TeamSession, + agent: &TeamAgent, + operation_id: u64, + generation: &str, + started_at: Instant, +) { + service + .cleanup_stale_member_runtime_task(session, &agent.conversation_id) + .await; + info!( + team_id = session.team_id(), + slot_id = agent.slot_id, + conversation_id = agent.conversation_id, + operation_id, + generation, + duration_ms = started_at.elapsed().as_millis(), + error_classification = "attach_cancelled", + "team member runtime attach cancelled" + ); + warn!( + team_id = session.team_id(), + slot_id = agent.slot_id, + conversation_id = agent.conversation_id, + operation_id, + generation, + duration_ms = started_at.elapsed().as_millis(), + error_classification = "stale_attach_completion", + "stale team member runtime attach completion rejected" + ); +} + +#[cfg(test)] +mod tests { + use super::*; use crate::event_loop::AgentLoopContext; - use crate::team_run::{ActiveChildTurn, RecoveryWakeCandidate}; use crate::test_utils::MockTeamRepo; use crate::types::{Team, TeamAgent, TeammateRole}; + use crate::work_coordinator::SlotPhase; use aionui_ai_agent::AgentError; use aionui_ai_agent::agent_task::AgentInstance; use aionui_ai_agent::types::BuildTaskOptions; - use aionui_api_types::{ - TeamRunSource, TeamRunStatus, TeamRunTargetRole, TeamSendMessageDelivery, TeamSendMessageReason, - WebSocketMessage, - }; - use aionui_common::{AgentKillReason, TimestampMs, now_ms}; + use aionui_api_types::WebSocketMessage; + use aionui_common::{AgentKillReason, TimestampMs}; use std::sync::{Arc, Mutex}; struct NullBroadcaster; @@ -1572,6 +1761,52 @@ mod tests { struct NoopTurnPort; + struct BlockingRunningTurnPort { + started_tx: Mutex>>, + release_rx: Mutex>>, + } + + impl BlockingRunningTurnPort { + fn new(started_tx: tokio::sync::oneshot::Sender<()>, release_rx: tokio::sync::oneshot::Receiver<()>) -> Self { + Self { + started_tx: Mutex::new(Some(started_tx)), + release_rx: Mutex::new(Some(release_rx)), + } + } + } + + #[async_trait::async_trait] + impl crate::ports::AgentTurnExecutionPort for BlockingRunningTurnPort { + async fn run_agent_turn( + &self, + request: crate::ports::AgentTurnRequest, + ) -> Result { + if let Some(on_started) = request.on_started.as_ref() { + on_started(crate::ports::AgentTurnStarted { + team_run_id: request.team_run_id.clone(), + slot_id: request.slot_id.clone(), + role: request.role.clone(), + conversation_id: request.conversation_id.clone(), + turn_id: "turn-background".into(), + }) + .await; + } + if let Some(tx) = self.started_tx.lock().unwrap().take() { + let _ = tx.send(()); + } + let release_rx = self.release_rx.lock().unwrap().take(); + if let Some(rx) = release_rx { + let _ = rx.await; + } + Ok(crate::ports::AgentTurnOutcome { + conversation_id: request.conversation_id, + turn_id: "turn-background".into(), + status: crate::ports::AgentTurnStatus::Completed, + runtime: None, + }) + } + } + #[async_trait::async_trait] impl crate::ports::AgentTurnExecutionPort for NoopTurnPort { async fn run_agent_turn( @@ -1580,7 +1815,7 @@ mod tests { ) -> Result { if let Some(on_started) = request.on_started.as_ref() { on_started(crate::ports::AgentTurnStarted { - team_run_id: request.team_run_id.clone().expect("team run id"), + team_run_id: request.team_run_id.clone(), slot_id: request.slot_id.clone(), role: request.role.clone(), conversation_id: request.conversation_id.clone(), @@ -1619,26 +1854,6 @@ mod tests { Arc::new(NoopCancellationPort) } - struct FailingCancellationPort; - - #[async_trait::async_trait] - impl crate::ports::AgentTurnCancellationPort for FailingCancellationPort { - async fn cancel_agent_turn( - &self, - _user_id: &str, - _conversation_id: &str, - _turn_id: &str, - ) -> Result<(), crate::ports::AgentTurnExecutionError> { - Err(crate::ports::AgentTurnExecutionError::Failed { - reason: "cancel unavailable".into(), - }) - } - } - - fn failing_cancellation_port() -> Arc { - Arc::new(FailingCancellationPort) - } - #[derive(Default)] struct NoopProjectionStore; @@ -1666,90 +1881,6 @@ mod tests { Arc::new(NoopProjectionStore) } - #[derive(Default)] - struct RecordingProjectionStore { - inserted: std::sync::Mutex>, - } - - #[async_trait::async_trait] - impl TeamProjectionMessageStore for RecordingProjectionStore { - fn mint_message_id(&self) -> String { - "minted-message".into() - } - - async fn find_projected_message( - &self, - _conversation_id: &str, - _msg_id: &str, - _msg_type: &str, - ) -> Result, TeamError> { - Ok(None) - } - - async fn insert_projected_message(&self, row: &aionui_db::models::MessageRow) -> Result<(), TeamError> { - self.inserted.lock().unwrap().push(row.clone()); - Ok(()) - } - } - - #[derive(Default)] - struct CompletingProjectionStore { - manager: Mutex>>, - } - - impl CompletingProjectionStore { - fn set_manager(&self, manager: Arc) { - *self.manager.lock().unwrap() = Some(manager); - } - } - - #[async_trait::async_trait] - impl TeamProjectionMessageStore for CompletingProjectionStore { - fn mint_message_id(&self) -> String { - "projection-completes-run".into() - } - - async fn find_projected_message( - &self, - _conversation_id: &str, - _msg_id: &str, - _msg_type: &str, - ) -> Result, TeamError> { - Ok(None) - } - - async fn insert_projected_message(&self, _row: &aionui_db::models::MessageRow) -> Result<(), TeamError> { - let manager = self.manager.lock().unwrap().clone(); - if let Some(manager) = manager { - manager.maybe_complete().await; - } - Ok(()) - } - } - - #[derive(Default)] - struct FailingProjectionStore; - - #[async_trait::async_trait] - impl TeamProjectionMessageStore for FailingProjectionStore { - fn mint_message_id(&self) -> String { - "projection-fails".into() - } - - async fn find_projected_message( - &self, - _conversation_id: &str, - _msg_id: &str, - _msg_type: &str, - ) -> Result, TeamError> { - Ok(None) - } - - async fn insert_projected_message(&self, _row: &aionui_db::models::MessageRow) -> Result<(), TeamError> { - Err(TeamError::InvalidRequest("projection failed for test".into())) - } - } - /// RecordingBroadcaster used by the D29d-1 ratification test below to /// assert that `team.agentSpawned` is *not* emitted on failed spawns. #[derive(Default)] @@ -1892,9 +2023,13 @@ mod tests { } async fn start_session() -> TeamSession { + start_session_with_prompt_dump(crate::prompt_dump::TeamPromptDumpConfig::disabled()).await + } + + async fn start_session_with_prompt_dump(prompt_dump: crate::prompt_dump::TeamPromptDumpConfig) -> TeamSession { let repo: Arc = Arc::new(MockTeamRepo::new()); let broadcaster: Arc = Arc::new(NullBroadcaster); - TeamSession::start( + TeamSession::start_with_prompt_dump( make_team(), repo, broadcaster, @@ -1905,6 +2040,7 @@ mod tests { noop_projection_store(), "user-test".into(), Weak::::new(), + prompt_dump, ) .await .unwrap() @@ -1914,829 +2050,128 @@ mod tests { Arc::new(start_session().await) } - async fn record_recovery_wake(session: &TeamSession, slot_id: &str, role: TeamRunTargetRole, unread_count: usize) { + fn register_test_event_loop(session: &Arc, slot_id: &str) { session - .team_run_manager() - .recover_mailbox_backlog(vec![RecoveryWakeCandidate { - slot_id: slot_id.to_owned(), - role, - unread_count, - }]) - .await - .expect("recovery wake"); + .event_loops() + .spawn(slot_id, test_loop_context(session, slot_id)) + .expect("test event loop registration"); } - fn register_test_event_loop(session: &Arc, slot_id: &str) { - session.event_loops().spawn( - slot_id, - AgentLoopContext { - team_id: session.team_id().to_owned(), - slot_id: slot_id.to_owned(), - user_id: session.user_id().to_owned(), - session: session.clone(), - scheduler: session.scheduler().clone(), - mailbox: session.mailbox().clone(), - turn_port: session.turn_port().clone(), - registry: session.event_loops().clone(), - }, + fn test_loop_context(session: &Arc, slot_id: &str) -> AgentLoopContext { + AgentLoopContext { + team_id: session.team_id().to_owned(), + slot_id: slot_id.to_owned(), + user_id: session.user_id().to_owned(), + session: session.clone(), + scheduler: session.scheduler().clone(), + mailbox: session.mailbox().clone(), + turn_port: session.turn_port().clone(), + registry: session.event_loops().clone(), + } + } + + #[tokio::test] + async fn event_loop_registry_rejects_duplicate_and_stopped_registration() { + let session = Arc::new(start_session().await); + register_test_event_loop(&session, "worker-1"); + + assert_eq!( + session + .event_loops() + .spawn("worker-1", test_loop_context(&session, "worker-1")), + Err(crate::event_loop::EventLoopRegistrationError::Duplicate) + ); + + session.stop(); + assert_eq!( + session + .event_loops() + .spawn("lead-1", test_loop_context(&session, "lead-1")), + Err(crate::event_loop::EventLoopRegistrationError::Stopped) ); } - async fn start_session_with_projection_store(store: Arc) -> TeamSession { - let repo: Arc = Arc::new(MockTeamRepo::new()); + #[tokio::test] + async fn start_and_stop() { + let session = start_session().await; + assert_eq!(session.team_id(), "t1"); + assert!(session.mcp_server.port() > 0); + session.stop(); + } + + #[tokio::test] + async fn mcp_stdio_config_for_agent() { + let session = start_session().await; + let config = session.mcp_stdio_config("lead-1"); + assert_eq!(config.team_id, "t1"); + assert_eq!(config.slot_id, "lead-1"); + assert_eq!(config.port, session.mcp_server.port()); + session.stop(); + } + + #[tokio::test] + async fn stdio_spec_uses_fixed_name_and_binary_path() { + let session = start_session().await; + let spec = session.stdio_spec("lead-1"); + assert_eq!(spec.name, crate::mcp::TEAM_MCP_SERVER_NAME); + assert_eq!(spec.command, "/tmp/aioncore-test"); + assert_eq!(spec.args, vec!["mcp-bridge".to_string()]); + assert!(spec.env.iter().any(|(k, v)| k == "TEAM_AGENT_SLOT_ID" && v == "lead-1")); + session.stop(); + } + + #[tokio::test] + async fn send_message_writes_to_lead_mailbox() { + let repo = Arc::new(MockTeamRepo::new()); + let repo_dyn: Arc = repo.clone(); let broadcaster: Arc = Arc::new(NullBroadcaster); - TeamSession::start( + let session = TeamSession::start( make_team(), - repo, + repo_dyn, broadcaster, backend_path(), empty_task_manager(), noop_turn_port(), noop_cancellation_port(), - store, + noop_projection_store(), "user-test".into(), Weak::::new(), ) .await - .unwrap() + .unwrap(); + session.send_message("Hello team", None).await.unwrap(); + + let state = repo.state.lock().unwrap(); + assert_eq!(state.messages.len(), 1); + assert_eq!(state.messages[0].to_agent_id, "lead-1"); + assert_eq!(state.messages[0].from_agent_id, "user"); + assert_eq!(state.messages[0].content, "Hello team"); + session.stop(); } - async fn start_session_with_repo_and_projection_store( - repo: Arc, - store: Arc, - ) -> TeamSession { + #[tokio::test] + async fn send_message_to_agent_writes_to_mailbox() { + let repo = Arc::new(MockTeamRepo::new()); + let repo_dyn: Arc = repo.clone(); let broadcaster: Arc = Arc::new(NullBroadcaster); - TeamSession::start( + let session = TeamSession::start( make_team(), - repo, + repo_dyn, broadcaster, backend_path(), empty_task_manager(), noop_turn_port(), noop_cancellation_port(), - store, + noop_projection_store(), "user-test".into(), Weak::::new(), ) .await - .unwrap() - } - - async fn start_session_with_cancellation_port( - cancellation_port: Arc, - ) -> TeamSession { - let repo: Arc = Arc::new(MockTeamRepo::new()); - let broadcaster: Arc = Arc::new(NullBroadcaster); - TeamSession::start( - make_team(), - repo, - broadcaster, - backend_path(), - empty_task_manager(), - noop_turn_port(), - cancellation_port, - noop_projection_store(), - "user-test".into(), - Weak::::new(), - ) - .await - .unwrap() - } - - #[tokio::test] - async fn agent_send_message_projects_recipient_visible_bubble_immediately() { - let store = Arc::new(RecordingProjectionStore::default()); - let session = start_session_with_projection_store(store.clone()).await; - session - .team_run_manager() - .accept_user_message("lead-1", TeamRunTargetRole::Lead, false, None) - .await - .expect("active run"); - - session - .send_agent_message_from_agent("lead-1", "worker-1", "Do the implementation") - .await - .expect("agent send succeeds"); - - let inserted = store.inserted.lock().unwrap(); - assert_eq!(inserted.len(), 1); - assert_eq!(inserted[0].conversation_id, "c2"); - assert_eq!(inserted[0].position.as_deref(), Some("left")); - assert!( - inserted[0] - .msg_id - .as_deref() - .unwrap_or_default() - .contains("team:t1:mailbox:"), - "projection must use mailbox-message dedupe key" - ); - session.stop(); - } - - #[tokio::test] - async fn agent_send_message_survives_completion_between_projection_and_commit() { - let store = Arc::new(CompletingProjectionStore::default()); - let session = start_session_with_projection_store(store.clone()).await; - store.set_manager(session.team_run_manager().clone()); - session - .team_run_manager() - .accept_user_message("lead-1", TeamRunTargetRole::Lead, false, None) - .await - .expect("active run before agent message"); - - session - .send_agent_message_from_agent("lead-1", "worker-1", "lease protected") - .await - .expect("agent message lease should retain run while projection completes it"); - - let unread = session - .mailbox() - .peek_unread(session.team_id(), "worker-1") - .await - .unwrap(); - assert!(unread.iter().any(|msg| msg.content == "lease protected")); - session.stop(); - } - - #[tokio::test] - async fn agent_message_to_idle_target_returns_queued_for_idle() { - let session = start_session_arc().await; - register_test_event_loop(&session, "worker-1"); - let ack = session.send_message("start", None).await.unwrap(); - - let result = session - .send_agent_message_from_agent("lead-1", "worker-1", "Do the implementation") - .await - .unwrap(); - - assert_eq!(result.team_run_id, ack.team_run_id); - assert_eq!(result.delivery, TeamSendMessageDelivery::WakeRecorded); - assert_eq!(result.target.queue_state, TeamSendMessageReason::QueuedForIdle); - assert_eq!(result.target.pending_wake_count, 1); - session.event_loops().shutdown(); - session.stop(); - } - - #[tokio::test] - async fn agent_message_to_active_target_returns_behind_active_turn() { - let session = start_session_arc().await; - register_test_event_loop(&session, "worker-1"); - let ack = session - .team_run_manager() - .accept_user_message("worker-1", TeamRunTargetRole::Teammate, false, None) - .await - .unwrap(); - session - .team_run_manager() - .record_pending_wake("worker-1", TeamRunTargetRole::Teammate, TeamWakeSource::UserMessage) - .await - .unwrap(); - let reservation = session - .team_run_manager() - .claim_wake_for_turn("worker-1", TeamRunTargetRole::Teammate, "conv-worker") - .await - .unwrap(); - session - .team_run_manager() - .record_child_started( - &reservation.reservation_id, - ActiveChildTurn { - team_run_id: ack.team_run_id.clone(), - slot_id: "worker-1".into(), - role: TeamRunTargetRole::Teammate, - conversation_id: "conv-worker".into(), - turn_id: "turn-worker".into(), - started_at_ms: now_ms(), - last_slow_notified_at_ms: None, - }, - ) - .await; - - let result = session - .send_agent_message_from_agent("lead-1", "worker-1", "Second item") - .await - .unwrap(); - - assert_eq!(result.target.queue_state, TeamSendMessageReason::BehindActiveTurn); - assert_eq!(result.target.pending_wake_count, 1); - assert_eq!(result.target.active_turn_id.as_deref(), Some("turn-worker")); - session.event_loops().shutdown(); - session.stop(); - } - - #[tokio::test] - async fn leader_message_reuses_active_run_when_leader_slot_is_free() { - let session = start_session().await; - let first = session - .send_message("First leader message", None) - .await - .expect("first leader send creates run"); - - session - .wake_agent_for_team_work("worker-1", TeamWakeSource::McpSendMessage, None) - .await - .expect("worker wake keeps run active"); - session.team_run_manager().record_empty_wake_observed("lead-1").await; - - let second = session - .send_message("Second leader message", None) - .await - .expect("leader slot free should accept active-run intervention"); - - assert_eq!(second.team_run_id, first.team_run_id); - assert_eq!(second.target_slot_id, "lead-1"); - assert_eq!(second.accepted_slot_id, "lead-1"); - assert_eq!(second.accepted_role, TeamRunTargetRole::Lead); - session.stop(); - } - - #[tokio::test] - async fn leader_message_queues_when_leader_slot_has_pending_wake() { - let session = start_session().await; - let first = session - .send_message("First leader message", None) - .await - .expect("first leader send creates run"); - - let second = session - .send_message("Second leader message", None) - .await - .expect("leader pending wake should accept additional foreground message"); - - assert_eq!(second.team_run_id, first.team_run_id); - let payload = session.team_run_manager().current_payload().await.unwrap(); - assert_eq!(payload.pending_wake_count, 2); - session.stop(); - } - - #[tokio::test] - async fn cancelling_leader_child_does_not_cancel_active_teammate_child() { - let session = start_session().await; - let ack = session - .team_run_manager() - .accept_user_message("lead-1", TeamRunTargetRole::Lead, false, None) - .await - .expect("accept run"); - - let leader = ActiveChildTurn { - team_run_id: ack.team_run_id.clone(), - slot_id: "lead-1".into(), - role: TeamRunTargetRole::Lead, - conversation_id: "c1".into(), - turn_id: "turn-lead".into(), - started_at_ms: now_ms(), - last_slow_notified_at_ms: None, - }; - let worker = ActiveChildTurn { - team_run_id: ack.team_run_id.clone(), - slot_id: "worker-1".into(), - role: TeamRunTargetRole::Teammate, - conversation_id: "c2".into(), - turn_id: "turn-worker".into(), - started_at_ms: now_ms(), - last_slow_notified_at_ms: None, - }; - - session - .team_run_manager() - .record_pending_wake("lead-1", TeamRunTargetRole::Lead, TeamWakeSource::UserMessage) - .await - .unwrap(); - let leader_reservation = session - .team_run_manager() - .claim_wake_for_turn("lead-1", TeamRunTargetRole::Lead, "c1") - .await - .unwrap(); - session - .team_run_manager() - .record_child_started(&leader_reservation.reservation_id, leader) - .await; - - session - .team_run_manager() - .record_pending_wake("worker-1", TeamRunTargetRole::Teammate, TeamWakeSource::McpSendMessage) - .await - .unwrap(); - let worker_reservation = session - .team_run_manager() - .claim_wake_for_turn("worker-1", TeamRunTargetRole::Teammate, "c2") - .await - .unwrap(); - session - .team_run_manager() - .record_child_started(&worker_reservation.reservation_id, worker) - .await; - - session - .cancel_child_turn(&ack.team_run_id, "lead-1", Some("user stopped leader".into())) - .await - .expect("leader child cancel succeeds"); - - let active = session.team_run_manager().active_child_turns().await; - assert_eq!(active.len(), 1); - assert_eq!(active[0].slot_id, "worker-1"); - assert_eq!(active[0].turn_id, "turn-worker"); - session.stop(); - } - - #[tokio::test] - async fn cancel_run_still_cancels_leader_and_teammate_children() { - let session = start_session().await; - let ack = session - .team_run_manager() - .accept_user_message("lead-1", TeamRunTargetRole::Lead, false, None) - .await - .expect("accept run"); - - session - .team_run_manager() - .record_pending_wake("lead-1", TeamRunTargetRole::Lead, TeamWakeSource::UserMessage) - .await - .unwrap(); - let leader_reservation = session - .team_run_manager() - .claim_wake_for_turn("lead-1", TeamRunTargetRole::Lead, "c1") - .await - .unwrap(); - session - .team_run_manager() - .record_child_started( - &leader_reservation.reservation_id, - ActiveChildTurn { - team_run_id: ack.team_run_id.clone(), - slot_id: "lead-1".into(), - role: TeamRunTargetRole::Lead, - conversation_id: "c1".into(), - turn_id: "turn-lead".into(), - started_at_ms: now_ms(), - last_slow_notified_at_ms: None, - }, - ) - .await; - - session - .team_run_manager() - .record_pending_wake("worker-1", TeamRunTargetRole::Teammate, TeamWakeSource::McpSendMessage) - .await - .unwrap(); - let worker_reservation = session - .team_run_manager() - .claim_wake_for_turn("worker-1", TeamRunTargetRole::Teammate, "c2") - .await - .unwrap(); - session - .team_run_manager() - .record_child_started( - &worker_reservation.reservation_id, - ActiveChildTurn { - team_run_id: ack.team_run_id.clone(), - slot_id: "worker-1".into(), - role: TeamRunTargetRole::Teammate, - conversation_id: "c2".into(), - turn_id: "turn-worker".into(), - started_at_ms: now_ms(), - last_slow_notified_at_ms: None, - }, - ) - .await; - - assert_eq!(session.team_run_manager().active_child_turns().await.len(), 2); - - session - .cancel_run(&ack.team_run_id, None, Some("stop all".into())) - .await - .expect("stop-all run cancel succeeds"); - - assert!( - session.team_run_manager().active_child_turns().await.is_empty(), - "cancel_run remains the explicit stop-all capability" - ); - session.stop(); - } - - #[tokio::test] - async fn wake_agent_for_team_work_records_pending_wake_without_registered_loop() { - let session = start_session().await; - session - .team_run_manager() - .accept_user_message("lead-1", TeamRunTargetRole::Lead, false, None) - .await - .expect("accept run"); - - session - .wake_agent_for_team_work("worker-1", TeamWakeSource::McpSendMessage, None) - .await - .expect("pending wake is recorded even before loop registration"); - - let reservation = session - .team_run_manager() - .claim_wake_for_turn("worker-1", TeamRunTargetRole::Teammate, "c2") - .await - .expect("wake should be claimable"); - - assert_eq!(reservation.slot_id, "worker-1"); - session.stop(); - } - - #[tokio::test] - async fn pause_active_child_cancel_success_marks_slot_paused_and_removes_child() { - let session = start_session().await; - let ack = session - .team_run_manager() - .accept_user_message("lead-1", TeamRunTargetRole::Lead, false, None) - .await - .unwrap(); - session - .team_run_manager() - .record_pending_wake("lead-1", TeamRunTargetRole::Lead, TeamWakeSource::UserMessage) - .await - .unwrap(); - let reservation = session - .team_run_manager() - .claim_wake_for_turn("lead-1", TeamRunTargetRole::Lead, "c1") - .await - .unwrap(); - session - .team_run_manager() - .record_child_started( - &reservation.reservation_id, - ActiveChildTurn { - team_run_id: ack.team_run_id.clone(), - slot_id: "lead-1".into(), - role: TeamRunTargetRole::Lead, - conversation_id: "c1".into(), - turn_id: "turn-lead".into(), - started_at_ms: now_ms(), - last_slow_notified_at_ms: None, - }, - ) - .await; - session - .team_run_manager() - .record_pending_wake("lead-1", TeamRunTargetRole::Lead, TeamWakeSource::McpSendMessage) - .await - .unwrap(); - - session - .pause_slot_work(&ack.team_run_id, "lead-1", Some("user stopped".into())) - .await - .unwrap(); - - let payload = session.team_run_manager().current_payload().await.unwrap(); - let lead = payload - .slot_work - .iter() - .find(|work| work.slot_id == "lead-1") - .expect("lead slot work"); - assert!(lead.paused); - assert_eq!(lead.pending_wake_count, 0); - assert_eq!(lead.suppressed_wake_count, 1); - assert_eq!(lead.active_turn_id, None); - session.stop(); - } - - #[tokio::test] - async fn pause_active_child_cancel_error_keeps_child_and_does_not_pause_slot() { - let session = start_session_with_cancellation_port(failing_cancellation_port()).await; - let ack = session - .team_run_manager() - .accept_user_message("lead-1", TeamRunTargetRole::Lead, false, None) - .await - .unwrap(); - session - .team_run_manager() - .record_pending_wake("lead-1", TeamRunTargetRole::Lead, TeamWakeSource::UserMessage) - .await - .unwrap(); - let reservation = session - .team_run_manager() - .claim_wake_for_turn("lead-1", TeamRunTargetRole::Lead, "c1") - .await - .unwrap(); - session - .team_run_manager() - .record_child_started( - &reservation.reservation_id, - ActiveChildTurn { - team_run_id: ack.team_run_id.clone(), - slot_id: "lead-1".into(), - role: TeamRunTargetRole::Lead, - conversation_id: "c1".into(), - turn_id: "turn-lead".into(), - started_at_ms: now_ms(), - last_slow_notified_at_ms: None, - }, - ) - .await; - session - .team_run_manager() - .record_pending_wake("lead-1", TeamRunTargetRole::Lead, TeamWakeSource::McpSendMessage) - .await - .unwrap(); - - let err = session - .pause_slot_work(&ack.team_run_id, "lead-1", Some("user stopped".into())) - .await - .unwrap_err(); - assert!(matches!(err, TeamError::InvalidRequest(message) if message.contains("cancel unavailable"))); - - let payload = session.team_run_manager().current_payload().await.unwrap(); - let lead = payload - .slot_work - .iter() - .find(|work| work.slot_id == "lead-1") - .expect("lead slot work"); - assert!(!lead.paused); - assert_eq!(lead.pending_wake_count, 1); - assert_eq!(lead.suppressed_wake_count, 0); - assert_eq!(lead.active_turn_id.as_deref(), Some("turn-lead")); - session.stop(); - } - - #[tokio::test] - async fn paused_slot_suppresses_mcp_send_without_mailbox_write() { - let session = start_session().await; - let ack = session.send_message("start", None).await.unwrap(); - session - .pause_slot_work(&ack.team_run_id, "lead-1", Some("user stopped".into())) - .await - .unwrap(); - - session - .send_agent_message_from_agent("worker-1", "lead-1", "background update") - .await - .unwrap(); - - let payload = session.team_run_manager().current_payload().await.expect("payload"); - let lead = payload.slot_work.iter().find(|work| work.slot_id == "lead-1").unwrap(); - assert!(lead.paused); - assert_eq!(lead.pending_wake_count, 0); - assert_eq!(lead.suppressed_wake_count, 2); - - let unread = session - .mailbox() - .peek_unread(session.team_id(), "lead-1") - .await - .unwrap(); - assert!( - unread.iter().all(|msg| msg.content != "background update"), - "suppressed acquire must not write mailbox" - ); - session.stop(); - } - - #[tokio::test] - async fn user_intervention_wake_input_prioritizes_only_user_message() { - let session = start_session().await; - let ack = session.send_message("start", None).await.unwrap(); - session - .pause_slot_work(&ack.team_run_id, "lead-1", Some("user stopped".into())) - .await - .unwrap(); - session - .send_agent_message_from_agent("worker-1", "lead-1", "old background") - .await - .unwrap(); - - let user_ack = session - .send_message_to_agent("lead-1", "please answer this first", None) - .await - .unwrap(); - - let input = session.compute_wake_input("lead-1").await.unwrap().expect("wake input"); - assert_eq!(input.wake_source, Some(TeamWakeSource::UserIntervention)); - assert_eq!(input.trigger_message_id, user_ack.message_id); - assert!(input.first_message.contains("## Turn Context")); - assert!(input.first_message.contains("Prioritize the user's message")); - assert!(input.first_message.contains("please answer this first")); - assert!(!input.first_message.contains("old background")); - assert_eq!(input.unread.len(), 1); - session.stop(); - } - - #[tokio::test] - async fn reserved_spawn_welcome_survives_leader_empty_wake_until_teammate_registers() { - let session = start_session().await; - let ack = session - .team_run_manager() - .accept_user_message("lead-1", TeamRunTargetRole::Lead, false, None) - .await - .expect("accept run"); - - let role = session - .reserve_wake_for_team_work("worker-1", TeamWakeSource::SpawnWelcome, None) - .await - .expect("reserve spawn welcome") - .expect("spawn welcome should be recorded"); - assert_eq!(role, TeamRunTargetRole::Teammate); - - assert!( - session - .team_run_manager() - .record_empty_wake_observed("lead-1") - .await - .is_none(), - "lead empty wake must not complete while worker spawn welcome is pending" - ); - assert_eq!( - session.team_run_manager().active_run_id().await.as_deref(), - Some(ack.team_run_id.as_str()) - ); - - let reservation = session - .team_run_manager() - .claim_wake_for_turn("worker-1", TeamRunTargetRole::Teammate, "c2") - .await - .expect("worker should claim reserved spawn welcome"); - - assert_eq!(reservation.slot_id, "worker-1"); - session.stop(); - } - - #[tokio::test] - async fn reserve_wake_for_team_work_rejects_without_active_run() { - let session = start_session().await; - - let err = session - .reserve_wake_for_team_work("worker-1", TeamWakeSource::SpawnWelcome, None) - .await - .expect_err("run-scoped reserve without active run must fail"); - - assert!(matches!( - err, - TeamError::InvalidRequest(message) - if message == "no active team run for run-scoped wake" - )); - session.stop(); - } - - #[tokio::test] - async fn wake_agent_for_team_work_rejects_orphan_mcp_work() { - let session = start_session().await; - - let err = session - .wake_agent_for_team_work("worker-1", TeamWakeSource::McpSendMessage, None) - .await - .expect_err("MCP work without active run must fail"); - - assert!(matches!( - err, - TeamError::InvalidRequest(message) - if message == "no active team run for run-scoped wake" - )); - session.stop(); - } - - #[tokio::test] - async fn crash_notification_in_active_run_wakes_leader() { - let session = start_session().await; - session - .team_run_manager() - .accept_user_message("lead-1", TeamRunTargetRole::Lead, false, None) - .await - .expect("accept run"); - - session - .wake_leader_after_recovery_message("worker-1", TeamWakeSource::CrashNotification) - .await - .expect("active run recovery should wake leader"); - - let reservation = session - .team_run_manager() - .claim_wake_for_turn("lead-1", TeamRunTargetRole::Lead, "c1") - .await - .expect("leader recovery wake should be claimable"); - - assert_eq!(reservation.slot_id, "lead-1"); - session.stop(); - } - - #[tokio::test] - async fn recovery_wake_without_active_run_is_deferred_without_active_run_snapshot_branch() { - let session = start_session().await; - - session - .wake_leader_after_recovery_message("worker-1", TeamWakeSource::CrashNotification) - .await - .expect("deferred recovery notification should not fail"); - - let reservation = session - .team_run_manager() - .claim_wake_for_turn("lead-1", TeamRunTargetRole::Lead, "c1") - .await; - - assert!( - reservation.is_none(), - "deferred recovery must not create Team Run reservation" - ); - session.stop(); - } - - #[tokio::test] - async fn inactivity_notification_in_active_run_wakes_leader() { - let session = start_session().await; - session - .team_run_manager() - .accept_user_message("lead-1", TeamRunTargetRole::Lead, false, None) - .await - .expect("accept run"); - - session - .wake_leader_after_recovery_message("worker-1", TeamWakeSource::InactivityTimeout) - .await - .expect("active run inactivity recovery should wake leader"); - - let reservation = session - .team_run_manager() - .claim_wake_for_turn("lead-1", TeamRunTargetRole::Lead, "c1") - .await - .expect("leader inactivity wake should be claimable"); - - assert_eq!(reservation.slot_id, "lead-1"); - session.stop(); - } - - #[tokio::test] - async fn start_and_stop() { - let session = start_session().await; - assert_eq!(session.team_id(), "t1"); - assert!(session.mcp_server.port() > 0); - session.stop(); - } - - #[tokio::test] - async fn mcp_stdio_config_for_agent() { - let session = start_session().await; - let config = session.mcp_stdio_config("lead-1"); - assert_eq!(config.team_id, "t1"); - assert_eq!(config.slot_id, "lead-1"); - assert_eq!(config.port, session.mcp_server.port()); - session.stop(); - } - - #[tokio::test] - async fn stdio_spec_uses_fixed_name_and_binary_path() { - let session = start_session().await; - let spec = session.stdio_spec("lead-1"); - assert_eq!(spec.name, crate::mcp::TEAM_MCP_SERVER_NAME); - assert_eq!(spec.command, "/tmp/aioncore-test"); - assert_eq!(spec.args, vec!["mcp-bridge".to_string()]); - assert!(spec.env.iter().any(|(k, v)| k == "TEAM_AGENT_SLOT_ID" && v == "lead-1")); - session.stop(); - } - - #[tokio::test] - async fn send_message_writes_to_lead_mailbox() { - let repo = Arc::new(MockTeamRepo::new()); - let repo_dyn: Arc = repo.clone(); - let broadcaster: Arc = Arc::new(NullBroadcaster); - let session = TeamSession::start( - make_team(), - repo_dyn, - broadcaster, - backend_path(), - empty_task_manager(), - noop_turn_port(), - noop_cancellation_port(), - noop_projection_store(), - "user-test".into(), - Weak::::new(), - ) - .await - .unwrap(); - session.send_message("Hello team", None).await.unwrap(); - - let state = repo.state.lock().unwrap(); - assert_eq!(state.messages.len(), 1); - assert_eq!(state.messages[0].to_agent_id, "lead-1"); - assert_eq!(state.messages[0].from_agent_id, "user"); - assert_eq!(state.messages[0].content, "Hello team"); - session.stop(); - } - - #[tokio::test] - async fn send_message_to_agent_writes_to_mailbox() { - let repo = Arc::new(MockTeamRepo::new()); - let repo_dyn: Arc = repo.clone(); - let broadcaster: Arc = Arc::new(NullBroadcaster); - let session = TeamSession::start( - make_team(), - repo_dyn, - broadcaster, - backend_path(), - empty_task_manager(), - noop_turn_port(), - noop_cancellation_port(), - noop_projection_store(), - "user-test".into(), - Weak::::new(), - ) - .await - .unwrap(); - session - .send_message_to_agent("worker-1", "Do this task", None) - .await - .unwrap(); + .unwrap(); + session + .send_message_to_agent("worker-1", "Do this task", None) + .await + .unwrap(); let state = repo.state.lock().unwrap(); assert_eq!(state.messages.len(), 1); @@ -2754,63 +2189,17 @@ mod tests { } #[tokio::test] - async fn shutdown_without_active_run_does_not_write_mailbox() { - let session = start_session().await; - - let err = session - .shutdown_agent("lead-1", "worker-1", Some("done".into())) - .await - .expect_err("shutdown is run-scoped"); - - assert!(matches!( - err, - TeamError::InvalidRequest(message) - if message == "no active team run for run-scoped wake" - )); - let unread = session.mailbox().peek_unread("t1", "worker-1").await.unwrap(); - assert!(unread.is_empty()); - session.stop(); - } - - #[tokio::test] - async fn shutdown_busy_target_is_accepted_as_lifecycle_wake() { + async fn shutdown_without_active_run_queues_background_work() { let session = start_session().await; - let ack = session.send_message("start", None).await.unwrap(); - let reservation = session - .team_run_manager() - .claim_wake_for_turn("lead-1", TeamRunTargetRole::Lead, "c1") - .await - .unwrap(); - session - .team_run_manager() - .record_child_started( - &reservation.reservation_id, - ActiveChildTurn { - team_run_id: ack.team_run_id, - slot_id: "lead-1".into(), - role: TeamRunTargetRole::Lead, - conversation_id: "c1".into(), - turn_id: "turn-lead".into(), - started_at_ms: now_ms(), - last_slow_notified_at_ms: None, - }, - ) - .await; session .shutdown_agent("lead-1", "worker-1", Some("done".into())) .await - .expect("shutdown lifecycle wake should be accepted"); + .expect("background caller may queue shutdown work"); let unread = session.mailbox().peek_unread("t1", "worker-1").await.unwrap(); assert_eq!(unread.len(), 1); - assert_eq!(unread[0].msg_type, MailboxMessageType::ShutdownRequest); - let reservation = session - .team_run_manager() - .claim_wake_for_turn("worker-1", TeamRunTargetRole::Teammate, "c2") - .await - .expect("shutdown wake should be pending"); - assert_eq!(reservation.wake_source, TeamWakeSource::McpShutdownRequest); + assert!(session.team_run_manager().current_active_run_id().is_none()); session.stop(); } @@ -2937,219 +2326,31 @@ mod tests { let result = session.rename_agent("worker-1", &lead_name).await; assert!(result.is_err()); - session.stop(); - } - - // -- spawn_agent helpers + guard tests ----------------------------------- - - fn sample_spawn_req() -> SpawnAgentRequest { - SpawnAgentRequest { - name: "Helper".into(), - assistant_id: Some("word-creator".into()), - model: None, - } - } - - #[tokio::test] - async fn spawn_agent_rejects_unknown_caller() { - let session = start_session().await; - let result = session.spawn_agent("nonexistent", sample_spawn_req()).await; - assert!( - matches!(&result, Err(TeamError::AgentNotFound(_))), - "unknown caller must surface AgentNotFound, got {result:?}" - ); - session.stop(); - } - - // -- D7a new method tests ------------------------------------------------ - - #[tokio::test] - async fn compute_wake_input_cold_start_injects_lead_role_prompt() { - let session = start_session().await; - // Seed one unread message. `send_message` flips status to Working — - // that is the post-send path; here we want to exercise cold-start - // detection, so write directly to the mailbox instead. - session - .mailbox - .write("t1", "lead-1", "user", MailboxMessageType::Message, "kick off", None) - .await - .unwrap(); - record_recovery_wake(&session, "lead-1", TeamRunTargetRole::Lead, 1).await; - - let input = session.compute_wake_input("lead-1").await.unwrap().expect("WakeInput"); - - assert_eq!(input.conversation_id, "c1"); - assert!(input.should_send); - assert!( - input.first_message.contains("You are the Team Leader"), - "expected lead role prompt, got: {}", - input.first_message - ); - assert!(input.first_message.contains("kick off")); - session.stop(); - } - - #[tokio::test] - async fn compute_wake_input_cold_start_injects_teammate_role_prompt() { - let session = start_session().await; - session - .mailbox - .write("t1", "worker-1", "user", MailboxMessageType::Message, "do X", None) - .await - .unwrap(); - record_recovery_wake(&session, "worker-1", TeamRunTargetRole::Teammate, 1).await; - - let input = session - .compute_wake_input("worker-1") - .await - .unwrap() - .expect("WakeInput"); - - assert!( - input.first_message.contains("## Team Governance"), - "expected Team Governance in teammate role prompt, got: {}", - input.first_message - ); - assert!( - input.first_message.contains("# You are a Team Member"), - "expected teammate role prompt, got: {}", - input.first_message - ); - assert!(input.first_message.contains("do X")); - session.stop(); - } - - #[tokio::test] - async fn teammate_first_wake_uses_canonical_prompt() { - let session = start_session().await; - session - .mailbox - .write("t1", "worker-1", "user", MailboxMessageType::Message, "do X", None) - .await - .unwrap(); - record_recovery_wake(&session, "worker-1", TeamRunTargetRole::Teammate, 1).await; - - let input = session - .compute_wake_input("worker-1") - .await - .unwrap() - .expect("WakeInput"); - let first_message = input.first_message; - - assert!(first_message.contains("## Team Governance")); - assert!(first_message.contains("You MUST use the `team_*` MCP tools for ALL team coordination.")); - assert!(first_message.contains("Use team_send_message to report results to the leader")); - assert!(first_message.contains("STOP GENERATING")); - assert!(!first_message.contains( - "You execute tasks assigned by the Lead Agent. Focus on completing your assigned work thoroughly and reporting back." - )); - assert!(first_message.contains("do X")); - session.stop(); - } - - #[tokio::test] - async fn compute_wake_input_warm_agent_skips_role_prompt() { - let session = start_session().await; - // Exit cold-start by setting a status once; any non-Error status - // means the scheduler has seen this agent before. - session - .scheduler - .set_status("lead-1", TeammateStatus::Idle) - .await - .unwrap(); - session - .mailbox - .write("t1", "lead-1", "user", MailboxMessageType::Message, "follow-up", None) - .await - .unwrap(); - record_recovery_wake(&session, "lead-1", TeamRunTargetRole::Lead, 1).await; - - let input = session.compute_wake_input("lead-1").await.unwrap().expect("WakeInput"); - - assert!(input.should_send); - assert!( - !input.first_message.contains("Lead Agent of team"), - "should not re-inject role prompt, got: {}", - input.first_message - ); - assert!(input.first_message.contains("follow-up")); - session.stop(); - } - - #[tokio::test] - async fn compute_wake_input_empty_mailbox_should_not_send() { - let session = start_session().await; - - let input = session.compute_wake_input("lead-1").await.unwrap().expect("WakeInput"); - - assert!(!input.should_send); - session.stop(); - } - - #[tokio::test] - async fn compute_wake_input_refuses_unowned_mailbox_backlog() { - let session = start_session().await; - let lead = session.scheduler().find_lead_slot_id().await.expect("lead"); - session - .mailbox() - .write( - session.team_id(), - &lead, - "worker-1", - MailboxMessageType::Message, - "recover this", - None, - ) - .await - .expect("mailbox write"); - - let input = session - .compute_wake_input(&lead) - .await - .expect("compute should not fail"); - - assert!(input.is_none(), "unowned unread mailbox must not start a turn"); - } - - #[tokio::test] - async fn compute_wake_input_accepts_recovery_pending_wake() { - let session = start_session().await; - let lead = session.scheduler().find_lead_slot_id().await.expect("lead"); - session - .mailbox() - .write( - session.team_id(), - &lead, - "worker-1", - MailboxMessageType::Message, - "recover this", - None, - ) - .await - .expect("mailbox write"); - session - .team_run_manager() - .recover_mailbox_backlog(vec![RecoveryWakeCandidate { - slot_id: lead.clone(), - role: TeamRunTargetRole::Lead, - unread_count: 1, - }]) - .await - .expect("recovery run"); + session.stop(); + } - let input = session - .compute_wake_input(&lead) - .await - .expect("compute") - .expect("owned wake input"); + // -- spawn_agent helpers + guard tests ----------------------------------- + + fn sample_spawn_req() -> SpawnAgentRequest { + SpawnAgentRequest { + name: "Helper".into(), + assistant_id: Some("word-creator".into()), + } + } - assert!(input.team_run_id.is_some()); - assert!(input.should_send); - assert_eq!(input.wake_source, Some(TeamWakeSource::RecoveryDrain)); - assert!(input.trigger_message_id.is_none()); - assert_eq!(input.unread.len(), 1); + #[tokio::test] + async fn spawn_agent_rejects_unknown_caller() { + let session = start_session().await; + let result = session.spawn_agent("nonexistent", sample_spawn_req()).await; + assert!( + matches!(&result, Err(TeamError::AgentNotFound(_))), + "unknown caller must surface AgentNotFound, got {result:?}" + ); + session.stop(); } + // -- D7a new method tests ------------------------------------------------ + #[tokio::test] async fn recovery_scan_creates_one_wake_per_slot_with_non_self_unread() { let session = start_session_arc().await; @@ -3205,12 +2406,10 @@ mod tests { let result = session .try_start_recovery_drain("test_restore") .await - .expect("scan should not fail") - .expect("recovery result"); + .expect("scan should not fail"); - assert_eq!(result.recorded_wakes, vec![lead.clone()]); - assert_eq!(result.source, TeamRunSource::RecoveryDrain); - assert_eq!(result.pending_wake_count, 1); + assert_eq!(result, vec![lead.clone()]); + assert!(session.team_run_manager().current_active_run_id().is_none()); } #[tokio::test] @@ -3223,7 +2422,7 @@ mod tests { .try_start_recovery_drain("test_restore_no_work") .await .expect("first scan"); - assert!(first.is_none(), "empty first scan must not create recovery work"); + assert!(first.is_empty(), "empty first scan must not create recovery work"); session .mailbox() @@ -3244,286 +2443,201 @@ mod tests { .expect("second scan should not fail"); assert!( - second.is_none(), + second.is_empty(), "ordinary new work must not create a second recovery scan" ); } #[tokio::test] - async fn compute_wake_input_returns_unread_rows_and_role_for_teammate() { + async fn on_agent_finish_marks_idle_and_returns_lead_when_all_settled() { let session = start_session().await; + + // Worker is Working; on finish → mark idle → since the lead is the + // only remaining non-idle member (actually also idle), all-idle + // check returns the lead slot_id. session - .mailbox - .write( - "t1", - "worker-1", - "lead-1", - MailboxMessageType::Message, - "from lead", - None, - ) - .await - .unwrap(); - session - .mailbox - .write("t1", "worker-1", "user", MailboxMessageType::Message, "from user", None) + .scheduler + .set_status("worker-1", TeammateStatus::Working) .await .unwrap(); - record_recovery_wake(&session, "worker-1", TeamRunTargetRole::Teammate, 2).await; - let input = session - .compute_wake_input("worker-1") - .await - .unwrap() - .expect("WakeInput"); + let result = session.on_agent_finish("c2", false).await.unwrap(); + assert_eq!(result.as_deref(), Some("lead-1")); - assert_eq!(input.unread.len(), 2); - assert!(matches!(input.agent_role, TeammateRole::Teammate)); - assert!(input.unread.iter().any(|m| m.from_agent_id == "lead-1")); - assert!(input.unread.iter().any(|m| m.from_agent_id == "user")); + let status = session.scheduler.get_status("worker-1").await.unwrap(); + assert_eq!(status, TeammateStatus::Idle); session.stop(); } #[tokio::test] - async fn compute_wake_input_returns_lead_role() { + async fn on_agent_finish_lead_returns_none() { let session = start_session().await; session - .mailbox - .write("t1", "lead-1", "user", MailboxMessageType::Message, "hi lead", None) + .scheduler + .set_status("lead-1", TeammateStatus::Working) .await .unwrap(); - record_recovery_wake(&session, "lead-1", TeamRunTargetRole::Lead, 1).await; - let input = session.compute_wake_input("lead-1").await.unwrap().expect("WakeInput"); - - assert!(matches!(input.agent_role, TeammateRole::Lead)); - assert_eq!(input.unread.len(), 1); + let result = session.on_agent_finish("c1", false).await.unwrap(); + assert!(result.is_none()); session.stop(); } #[tokio::test] - async fn mirror_unread_to_conversation_skips_when_service_weak_is_dangling_for_leader() { + async fn on_agent_finish_unknown_conversation_returns_error() { let session = start_session().await; - session - .mailbox - .write( - "t1", - "lead-1", - "worker-1", - MailboxMessageType::Message, - "lead-gets-this", - None, - ) - .await - .unwrap(); - record_recovery_wake(&session, "lead-1", TeamRunTargetRole::Lead, 1).await; + let result = session.on_agent_finish("nope", false).await; + assert!(result.is_err()); + session.stop(); + } - let input = session.compute_wake_input("lead-1").await.unwrap().expect("WakeInput"); + // -- D7b wake-path tests ------------------------------------------------- - // In unit tests, `service` is a dangling Weak — the mirror helper must - // skip gracefully even for leader targets. - session.mirror_unread_to_conversation(&input).await; - session.stop(); + async fn start_session_with(task_manager: Arc) -> (TeamSession, Arc) { + let repo = Arc::new(MockTeamRepo::new()); + let repo_dyn: Arc = repo.clone(); + let broadcaster: Arc = Arc::new(NullBroadcaster); + let session = TeamSession::start( + make_team(), + repo_dyn, + broadcaster, + backend_path(), + task_manager, + noop_turn_port(), + noop_cancellation_port(), + noop_projection_store(), + "user-test".into(), + Weak::::new(), + ) + .await + .unwrap(); + (session, repo) } #[tokio::test] - async fn mirror_unread_to_conversation_skips_idle_notification_bubbles() { - let store = Arc::new(RecordingProjectionStore::default()); - let session = start_session_with_projection_store(store.clone()).await; + async fn send_message_persists_files_in_mailbox_without_inline_wake() { + let (session, _repo) = start_session_with(empty_task_manager()).await; + session - .mailbox - .write( - "t1", - "lead-1", - "worker-1", - MailboxMessageType::IdleNotification, - "idle", - Some("idle"), - ) + .send_message("Hello", Some(vec!["/tmp/a.txt".into(), "/tmp/b.txt".into()])) .await .unwrap(); - record_recovery_wake(&session, "lead-1", TeamRunTargetRole::Lead, 1).await; - let input = session.compute_wake_input("lead-1").await.unwrap().expect("WakeInput"); + let unread = session.mailbox.peek_unread("t1", "lead-1").await.unwrap(); + assert_eq!(unread.len(), 1); assert_eq!( - input.unread.len(), - 1, - "idle notification must still be delivered to wake payload" - ); - - session.mirror_unread_to_conversation(&input).await; - - assert!( - store.inserted.lock().unwrap().is_empty(), - "idle notification must not become a visible chat bubble" + unread[0].files.as_deref(), + Some(&["/tmp/a.txt".into(), "/tmp/b.txt".into()][..]) ); + assert_eq!(unread[0].content, "Hello"); session.stop(); } #[tokio::test] - async fn mirror_unread_to_conversation_still_projects_teammate_messages() { - let store = Arc::new(RecordingProjectionStore::default()); - let session = start_session_with_projection_store(store.clone()).await; + async fn rejected_commit_releases_lease_and_settles_the_persisted_message() { + let (session, _repo) = start_session_with(empty_task_manager()).await; session + .work_coordinator + .set_runtime_constraint("lead-1", RuntimeConstraint::Ready); + let lease = session + .work_coordinator + .acquire_enqueue(EnqueueRequest { + slot_id: "lead-1".into(), + role: TeamRunTargetRole::Lead, + source: WorkSource::UserMessage, + binding: CausalBinding::UserVisible, + }) + .unwrap(); + let message = session .mailbox .write( "t1", "lead-1", - "worker-1", + "user", MailboxMessageType::Message, - "work is done", + "persisted before removal", None, ) .await .unwrap(); - record_recovery_wake(&session, "lead-1", TeamRunTargetRole::Lead, 1).await; - - let input = session.compute_wake_input("lead-1").await.unwrap().expect("WakeInput"); - session.mirror_unread_to_conversation(&input).await; - - let inserted = store.inserted.lock().unwrap(); - assert_eq!(inserted.len(), 1, "normal teammate message must remain visible"); - let content: serde_json::Value = serde_json::from_str(&inserted[0].content).unwrap(); - assert_eq!(content["content"], "work is done"); - session.stop(); - } - - #[tokio::test] - async fn send_message_to_agent_survives_completion_between_projection_and_commit() { - let store = Arc::new(CompletingProjectionStore::default()); - let session = start_session_with_projection_store(store.clone()).await; - store.set_manager(session.team_run_manager().clone()); - session - .team_run_manager() - .accept_user_message("lead-1", TeamRunTargetRole::Lead, false, None) - .await - .expect("active run before user intervention"); + session.work_coordinator.remove_slot("lead-1"); - let ack = session - .send_message_to_agent("worker-1", "race window message", None) + let error = session + .commit_persisted_enqueue(&lease, message.id.clone()) .await - .expect("lease should retain run while projection completes it"); + .expect_err("removed slot must reject the commit"); - assert_eq!(ack.accepted_slot_id, "worker-1"); - let reservation = session - .team_run_manager() - .claim_wake_for_turn("worker-1", TeamRunTargetRole::Teammate, "c2") - .await - .expect("committed user intervention should be pending"); - assert_eq!(reservation.message_id, ack.message_id); - session.stop(); + assert!(matches!(error, TeamError::InvalidRequest(_))); + assert!(session.team_run_manager.current_active_run_id().is_none()); + assert!(session.mailbox.peek_unread("t1", "lead-1").await.unwrap().is_empty()); } #[tokio::test] - async fn send_message_projection_failure_still_commits_wake() { - let session = start_session_with_projection_store(Arc::new(FailingProjectionStore)).await; - - let ack = session - .send_message_to_agent("worker-1", "projection can fail", None) - .await - .expect("projection failure is non-fatal"); + async fn prepare_next_batch_requeues_claim_when_task_snapshot_fails() { + let (session, repo) = start_session_with(empty_task_manager()).await; + let session = Arc::new(session); + session.send_message("retry after task read", None).await.unwrap(); + repo.state.lock().unwrap().fail_task_lists = true; + register_test_event_loop(&session, "lead-1"); + + let error = match session.prepare_next_batch("lead-1").await { + Err(error) => error, + Ok(_) => panic!("task snapshot failure must be returned"), + }; - let reservation = session - .team_run_manager() - .claim_wake_for_turn("worker-1", TeamRunTargetRole::Teammate, "c2") - .await - .expect("wake should still be committed"); - assert_eq!(reservation.message_id, ack.message_id); + assert!(matches!(error, TeamError::Database(_))); + let slot = session.work_coordinator.slot_snapshot("lead-1").unwrap(); + assert_eq!(slot.state, SlotPhase::Queued); + assert!(slot.active_batch.is_none()); + assert_eq!(slot.queued_foreground_count, 1); session.stop(); } #[tokio::test] - async fn send_message_mailbox_failure_aborts_lease_and_allows_completion() { - let repo: Arc = Arc::new(MockTeamRepo::with_message_write_failure()); - let session = start_session_with_repo_and_projection_store(repo, noop_projection_store()).await; - - let err = session - .send_message("mailbox fails", None) - .await - .expect_err("mailbox write failure must be returned"); - assert!( - err.to_string().contains("forced mailbox write failure"), - "unexpected error: {err}" - ); - - let completed = session - .team_run_manager() - .maybe_complete() - .await - .expect("aborted lease should allow empty run to complete"); - assert_eq!(completed.status, TeamRunStatus::Completed); - session.stop(); - } + async fn background_turn_is_reported_running_without_a_team_run() { + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + let repo: Arc = Arc::new(MockTeamRepo::new()); + let broadcaster: Arc = Arc::new(NullBroadcaster); + let session = TeamSession::start( + make_team(), + repo, + broadcaster, + backend_path(), + empty_task_manager(), + Arc::new(BlockingRunningTurnPort::new(started_tx, release_rx)), + noop_cancellation_port(), + noop_projection_store(), + "user-test".into(), + Weak::::new(), + ) + .await + .unwrap(); + let session = Arc::new(session); + register_test_event_loop(&session, "worker-1"); - #[tokio::test] - async fn mirror_unread_to_conversation_skips_when_service_weak_is_dangling() { - let session = start_session().await; session - .mailbox - .write("t1", "worker-1", "lead-1", MailboxMessageType::Message, "do it", None) + .shutdown_agent("lead-1", "worker-1", Some("background control".into())) .await .unwrap(); - record_recovery_wake(&session, "worker-1", TeamRunTargetRole::Teammate, 1).await; - - let input = session - .compute_wake_input("worker-1") - .await - .unwrap() - .expect("WakeInput"); - - // In unit tests, `service` is a dangling Weak — the mirror helper must - // skip gracefully (no panic, no broadcast), leaving the wake path to - // still forward `first_message` to the agent. - session.mirror_unread_to_conversation(&input).await; - session.stop(); - } - - #[tokio::test] - async fn on_agent_finish_marks_idle_and_returns_lead_when_all_settled() { - let session = start_session().await; - - // Worker is Working; on finish → mark idle → since the lead is the - // only remaining non-idle member (actually also idle), all-idle - // check returns the lead slot_id. - session - .scheduler - .set_status("worker-1", TeammateStatus::Working) + tokio::time::timeout(std::time::Duration::from_secs(2), started_rx) .await - .unwrap(); - - let result = session.on_agent_finish("c2", false).await.unwrap(); - assert_eq!(result.as_deref(), Some("lead-1")); - - let status = session.scheduler.get_status("worker-1").await.unwrap(); - assert_eq!(status, TeammateStatus::Idle); - session.stop(); - } + .expect("background turn should start") + .expect("start signal should be sent"); - #[tokio::test] - async fn on_agent_finish_lead_returns_none() { - let session = start_session().await; - session - .scheduler - .set_status("lead-1", TeammateStatus::Working) - .await - .unwrap(); + let slot = session.work_coordinator.slot_snapshot("worker-1").unwrap(); + assert_eq!(slot.state, SlotPhase::Running); + assert_eq!(slot.active_turn_id.as_deref(), Some("turn-background")); + assert!(session.team_run_manager.current_active_run_id().is_none()); - let result = session.on_agent_finish("c1", false).await.unwrap(); - assert!(result.is_none()); + release_tx.send(()).unwrap(); session.stop(); } #[tokio::test] - async fn on_agent_finish_unknown_conversation_returns_error() { - let session = start_session().await; - let result = session.on_agent_finish("nope", false).await; - assert!(result.is_err()); - session.stop(); - } - - // -- D7b wake-path tests ------------------------------------------------- - - async fn start_session_with(task_manager: Arc) -> (TeamSession, Arc) { + async fn busy_lead_user_message_is_persisted_and_acknowledged_as_queued() { + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); let repo = Arc::new(MockTeamRepo::new()); let repo_dyn: Arc = repo.clone(); let broadcaster: Arc = Arc::new(NullBroadcaster); @@ -3532,8 +2646,8 @@ mod tests { repo_dyn, broadcaster, backend_path(), - task_manager, - noop_turn_port(), + empty_task_manager(), + Arc::new(BlockingRunningTurnPort::new(started_tx, release_rx)), noop_cancellation_port(), noop_projection_store(), "user-test".into(), @@ -3541,28 +2655,142 @@ mod tests { ) .await .unwrap(); - (session, repo) + let session = Arc::new(session); + register_test_event_loop(&session, "lead-1"); + + let first = session.send_message("first", None).await.unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(2), started_rx) + .await + .expect("first turn should start") + .expect("start signal should be sent"); + let queued = session.send_message("queued", None).await.unwrap(); + + assert_eq!(queued.enqueue_status, TeamMessageEnqueueStatus::Queued); + assert_eq!(queued.run.team_run_id, first.run.team_run_id); + assert!( + session + .mailbox + .peek_unread("t1", "lead-1") + .await + .unwrap() + .iter() + .any(|message| message.id == queued.message_id) + ); + release_tx.send(()).unwrap(); + session.stop(); } #[tokio::test] - async fn send_message_persists_files_in_mailbox_without_inline_wake() { - let (session, _repo) = start_session_with(empty_task_manager()).await; + async fn mailbox_failure_aborts_enqueue_lease_and_allows_run_completion() { + let (session, repo) = start_session_with(empty_task_manager()).await; + repo.state.lock().unwrap().fail_message_writes = true; - session - .send_message("Hello", Some(vec!["/tmp/a.txt".into(), "/tmp/b.txt".into()])) + let error = session.send_message("will fail", None).await.unwrap_err(); + + assert!(error.to_string().contains("forced mailbox write failure")); + assert!(session.team_run_manager.current_active_run_id().is_none()); + assert!(session.work_coordinator.snapshot().active_run_summary.is_none()); + session.stop(); + } + + #[tokio::test] + async fn message_to_runtime_starting_member_returns_blocked_runtime_starting() { + let session = start_session().await; + + let ack = session + .send_message_to_agent("worker-1", "wait for runtime", None) .await .unwrap(); - let unread = session.mailbox.peek_unread("t1", "lead-1").await.unwrap(); - assert_eq!(unread.len(), 1); + assert_eq!(ack.enqueue_status, TeamMessageEnqueueStatus::BlockedRuntimeStarting); assert_eq!( - unread[0].files.as_deref(), - Some(&["/tmp/a.txt".into(), "/tmp/b.txt".into()][..]) + ack.run.slot_work[0].blocked_reason, + Some(aionui_api_types::TeamSlotBlockedReason::RuntimeStarting) ); - assert_eq!(unread[0].content, "Hello"); session.stop(); } + #[tokio::test] + async fn runtime_failed_removing_and_stopped_members_reject_before_mailbox_write() { + let (failed_session, failed_repo) = start_session_with(empty_task_manager()).await; + let ReserveAttach::Start(failed_lease) = failed_session.member_runtimes.reserve_attach("lead-1", false) else { + panic!("failed runtime setup must reserve attach"); + }; + assert!(failed_session.member_runtimes.commit_failed( + &failed_lease, + MemberRuntimeFailure { + classification: "test_runtime_failed", + public_reason: "test failure".into(), + }, + )); + let failed_error = failed_session.send_message("rejected", None).await.unwrap_err(); + assert!(failed_error.to_string().contains("runtime failed")); + assert!(failed_repo.state.lock().unwrap().messages.is_empty()); + failed_session.stop(); + + let (removing_session, removing_repo) = start_session_with(empty_task_manager()).await; + let ReserveAttach::Start(removing_lease) = removing_session.member_runtimes.reserve_attach("lead-1", false) + else { + panic!("removing runtime setup must reserve attach"); + }; + assert!(removing_session.member_runtimes.commit_ready(&removing_lease)); + assert!(matches!( + removing_session.member_runtimes.begin_remove("lead-1"), + BeginRemove::Start(_) + )); + let removing_error = removing_session.send_message("rejected", None).await.unwrap_err(); + assert!(removing_error.to_string().contains("being removed")); + assert!(removing_repo.state.lock().unwrap().messages.is_empty()); + removing_session.stop(); + + let (stopped_session, stopped_repo) = start_session_with(empty_task_manager()).await; + stopped_session.member_runtimes.stop(); + let stopped_error = stopped_session.send_message("rejected", None).await.unwrap_err(); + assert!(stopped_error.to_string().contains("session stopped")); + assert!(stopped_repo.state.lock().unwrap().messages.is_empty()); + stopped_session.stop(); + } + + #[test] + fn session_replacement_rejects_old_generation_batch_and_attach_completion() { + let old_broadcaster: Arc = Arc::new(NullBroadcaster); + let old_emitter = Arc::new(TeamEventEmitter::new("t1".into(), old_broadcaster)); + let old_runs = Arc::new(TeamRunManager::new("t1".into(), old_emitter)); + let old_coordinator = SlotWorkCoordinator::new("t1".into(), "old-generation".into(), old_runs); + old_coordinator.set_runtime_constraint("lead-1", RuntimeConstraint::Ready); + let lease = old_coordinator + .acquire_enqueue(EnqueueRequest { + slot_id: "lead-1".into(), + role: TeamRunTargetRole::Lead, + source: WorkSource::RecoveryDrain, + binding: CausalBinding::Background, + }) + .unwrap(); + old_coordinator.commit_enqueue(&lease, Some("m-old".into())).unwrap(); + let ReconcileDecision::Claim(old_batch) = old_coordinator.next("lead-1") else { + panic!("old generation batch must be claimed"); + }; + + let new_broadcaster: Arc = Arc::new(NullBroadcaster); + let new_emitter = Arc::new(TeamEventEmitter::new("t1".into(), new_broadcaster)); + let new_runs = Arc::new(TeamRunManager::new("t1".into(), new_emitter)); + let new_coordinator = SlotWorkCoordinator::new("t1".into(), "new-generation".into(), new_runs); + new_coordinator.set_runtime_constraint("lead-1", RuntimeConstraint::Ready); + assert_eq!( + new_coordinator.mark_started(&old_batch, "turn-old"), + crate::work_coordinator::StartCommitResult::StaleOwner + ); + assert!(new_coordinator.slot_snapshot("lead-1").unwrap().active_batch.is_none()); + + let old_registry = MemberRuntimeRegistry::new("old-generation"); + let ReserveAttach::Start(old_attach) = old_registry.reserve_attach("lead-1", false) else { + panic!("old attach must be reserved"); + }; + let new_registry = MemberRuntimeRegistry::new("new-generation"); + assert!(!new_registry.commit_ready(&old_attach)); + assert_eq!(new_registry.snapshot("lead-1"), MemberRuntimeSnapshot::Absent); + } + #[tokio::test] async fn send_message_without_active_task_does_not_error() { // Empty task_manager → get_task returns None → log-not-throw: the @@ -3671,7 +2899,6 @@ mod tests { SpawnAgentRequest { name: "Helper".into(), assistant_id: assistant_id.map(str::to_owned), - model: None, } } diff --git a/crates/aionui-team/src/slot_wake_gate.rs b/crates/aionui-team/src/slot_wake_gate.rs deleted file mode 100644 index f21d5043e..000000000 --- a/crates/aionui-team/src/slot_wake_gate.rs +++ /dev/null @@ -1,255 +0,0 @@ -use std::collections::HashMap; - -use aionui_api_types::TeamRunTargetRole; -use aionui_common::{TimestampMs, now_ms}; - -use crate::wake::{TeamWakeClass, TeamWakeSource}; - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum WakeGateDecision { - Record { resumed_from_pause: bool }, - Suppress, -} - -#[derive(Debug, Clone, Default, PartialEq, Eq)] -pub(crate) struct SlotWakeSnapshot { - pub paused: bool, - pub suppressed_wake_count: usize, -} - -#[derive(Debug, Clone)] -struct PausedSlotState { - paused: bool, - #[allow(dead_code)] - paused_at: TimestampMs, - role_hint: TeamRunTargetRole, - #[allow(dead_code)] - reason: String, - suppressed_wake_count: usize, - last_suppressed_source: Option, -} - -#[derive(Debug, Clone, Default)] -pub(crate) struct SlotWakeGate { - slots: HashMap, -} - -impl SlotWakeGate { - pub(crate) fn clear(&mut self) { - self.slots.clear(); - } - - pub(crate) fn pause(&mut self, slot_id: &str, role: TeamRunTargetRole, reason: impl Into) { - let reason = reason.into(); - let now = now_ms(); - let entry = self.slots.entry(slot_id.to_owned()).or_insert_with(|| PausedSlotState { - paused: true, - paused_at: now, - role_hint: role.clone(), - reason: reason.clone(), - suppressed_wake_count: 0, - last_suppressed_source: None, - }); - entry.paused = true; - entry.role_hint = role; - entry.paused_at = now; - entry.reason = reason; - } - - pub(crate) fn before_wake( - &mut self, - slot_id: &str, - source: TeamWakeSource, - _trigger_message_id: Option, - ) -> WakeGateDecision { - if source.resumes_paused_slot() { - let resumed_from_pause = if let Some(entry) = self.slots.get_mut(slot_id) { - let was_paused = entry.paused; - entry.paused = false; - entry.reason = "resumed_by_user".into(); - was_paused - } else { - false - }; - return WakeGateDecision::Record { resumed_from_pause }; - } - - let Some(entry) = self.slots.get_mut(slot_id) else { - return WakeGateDecision::Record { - resumed_from_pause: false, - }; - }; - - if source.bypasses_pause() { - return WakeGateDecision::Record { - resumed_from_pause: false, - }; - } - - match source.class() { - TeamWakeClass::Background | TeamWakeClass::SystemRecovery => { - entry.suppressed_wake_count += 1; - entry.last_suppressed_source = Some(source); - WakeGateDecision::Suppress - } - TeamWakeClass::Foreground | TeamWakeClass::Lifecycle => WakeGateDecision::Record { - resumed_from_pause: false, - }, - } - } - - pub(crate) fn snapshot_for_slot(&self, slot_id: &str) -> SlotWakeSnapshot { - self.slots - .get(slot_id) - .map(|entry| SlotWakeSnapshot { - paused: entry.paused, - suppressed_wake_count: entry.suppressed_wake_count, - }) - .unwrap_or_default() - } - - #[allow(dead_code)] - pub(crate) fn role_hint(&self, slot_id: &str) -> Option { - self.slots.get(slot_id).map(|entry| entry.role_hint.clone()) - } - - #[allow(dead_code)] - pub(crate) fn slot_ids(&self) -> impl Iterator { - self.slots.keys() - } - - pub(crate) fn has_retained_work(&self) -> bool { - self.slots - .values() - .any(|entry| entry.paused || entry.suppressed_wake_count > 0) - } - - #[allow(dead_code)] - pub(crate) fn add_suppressed(&mut self, slot_id: &str, count: usize) { - if count == 0 { - return; - } - if let Some(entry) = self.slots.get_mut(slot_id) { - entry.suppressed_wake_count += count; - entry - .last_suppressed_source - .get_or_insert(TeamWakeSource::McpSendMessage); - } - } - - pub(crate) fn release_suppressed_if_resumed(&mut self, slot_id: &str) -> Option { - let source = { - let entry = self.slots.get_mut(slot_id)?; - if entry.paused || entry.suppressed_wake_count == 0 { - return None; - } - let source = entry.last_suppressed_source.unwrap_or(TeamWakeSource::McpSendMessage); - entry.suppressed_wake_count = 0; - source - }; - if self.slots.get(slot_id).is_some_and(should_prune) { - self.slots.remove(slot_id); - } - Some(source) - } -} - -fn should_prune(entry: &PausedSlotState) -> bool { - !entry.paused && entry.suppressed_wake_count == 0 -} - -#[cfg(test)] -mod tests { - use super::*; - - use aionui_api_types::TeamRunTargetRole; - - use crate::wake::TeamWakeSource; - - #[test] - fn paused_slot_suppresses_background_but_not_lifecycle() { - let mut gate = SlotWakeGate::default(); - gate.pause("lead-1", TeamRunTargetRole::Lead, "user_stop"); - - let background = gate.before_wake("lead-1", TeamWakeSource::McpSendMessage, None); - assert_eq!(background, WakeGateDecision::Suppress); - assert_eq!(gate.snapshot_for_slot("lead-1").suppressed_wake_count, 1); - - let lifecycle = gate.before_wake("lead-1", TeamWakeSource::McpShutdownRequest, None); - assert_eq!( - lifecycle, - WakeGateDecision::Record { - resumed_from_pause: false - } - ); - assert!(gate.snapshot_for_slot("lead-1").paused); - } - - #[test] - fn user_intervention_resumes_and_records_one_time_decision() { - let mut gate = SlotWakeGate::default(); - gate.pause("worker-1", TeamRunTargetRole::Teammate, "user_stop"); - assert!(gate.snapshot_for_slot("worker-1").paused); - - let decision = gate.before_wake( - "worker-1", - TeamWakeSource::UserIntervention, - Some("msg-user".to_owned()), - ); - - assert_eq!( - decision, - WakeGateDecision::Record { - resumed_from_pause: true - } - ); - let snapshot = gate.snapshot_for_slot("worker-1"); - assert!(!snapshot.paused); - } - - #[test] - fn release_suppressed_returns_one_background_drain_when_resumed() { - let mut gate = SlotWakeGate::default(); - gate.pause("lead-1", TeamRunTargetRole::Lead, "user_stop"); - assert_eq!( - gate.before_wake("lead-1", TeamWakeSource::InterruptedNotification, None), - WakeGateDecision::Suppress - ); - assert_eq!( - gate.before_wake("lead-1", TeamWakeSource::IdleNotification, None), - WakeGateDecision::Suppress - ); - assert_eq!( - gate.before_wake("lead-1", TeamWakeSource::UserIntervention, Some("msg-user".into())), - WakeGateDecision::Record { - resumed_from_pause: true - } - ); - - let released = gate.release_suppressed_if_resumed("lead-1"); - assert_eq!(released, Some(TeamWakeSource::IdleNotification)); - let snapshot = gate.snapshot_for_slot("lead-1"); - assert_eq!(snapshot.suppressed_wake_count, 0); - } - - #[test] - fn release_suppressed_removes_slot_when_no_retained_work_remains() { - let mut gate = SlotWakeGate::default(); - gate.pause("lead-1", TeamRunTargetRole::Lead, "user stopped"); - assert_eq!( - gate.before_wake("lead-1", TeamWakeSource::IdleNotification, None), - WakeGateDecision::Suppress - ); - assert_eq!( - gate.before_wake("lead-1", TeamWakeSource::UserMessage, Some("msg-user".into())), - WakeGateDecision::Record { - resumed_from_pause: true - } - ); - - let released = gate.release_suppressed_if_resumed("lead-1"); - assert_eq!(released, Some(TeamWakeSource::IdleNotification)); - assert!(!gate.has_retained_work()); - assert_eq!(gate.slot_ids().count(), 0, "empty resumed gate entry should be pruned"); - } -} diff --git a/crates/aionui-team/src/team_run.rs b/crates/aionui-team/src/team_run.rs deleted file mode 100644 index 3297572ad..000000000 --- a/crates/aionui-team/src/team_run.rs +++ /dev/null @@ -1,3411 +0,0 @@ -use std::collections::{HashMap, VecDeque}; -use std::sync::Arc; - -use aionui_api_types::{ - TeamChildTurnPayload, TeamRunAckResponse, TeamRunPayload, TeamRunSource, TeamRunStatus, TeamRunTargetRole, - TeamSlotRuntimeHealth, TeamSlotWorkPayload, -}; -use aionui_common::{TimestampMs, generate_id, now_ms}; -use tokio::sync::Mutex; -use tracing::{debug, error, info, warn}; - -use crate::error::TeamError; -use crate::events::{ - TEAM_CHILD_TURN_CANCELLED_EVENT, TEAM_CHILD_TURN_COMPLETED_EVENT, TEAM_CHILD_TURN_STARTED_EVENT, - TEAM_RUN_ACCEPTED_EVENT, TEAM_RUN_CANCELLED_EVENT, TEAM_RUN_COMPLETED_EVENT, TEAM_RUN_FAILED_EVENT, - TEAM_RUN_STARTED_EVENT, TEAM_RUN_UPDATED_EVENT, TeamEventEmitter, -}; -use crate::slot_wake_gate::{SlotWakeGate, WakeGateDecision}; -use crate::types::TeammateRole; -use crate::wake::TeamWakeSource; - -const ACTIVE_CHILD_SLOW_THRESHOLD_MS: u64 = 10 * 60 * 1000; -const ACTIVE_CHILD_SLOW_REPEAT_MS: u64 = 10 * 60 * 1000; - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ActiveChildTurn { - pub team_run_id: String, - pub slot_id: String, - pub role: TeamRunTargetRole, - pub conversation_id: String, - pub turn_id: String, - pub started_at_ms: TimestampMs, - pub last_slow_notified_at_ms: Option, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum StartingReservationState { - Starting, - Cancelling, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct StartingChildReservation { - pub reservation_id: String, - pub team_run_id: String, - pub slot_id: String, - pub role: TeamRunTargetRole, - pub conversation_id: String, - pub(crate) wake_source: TeamWakeSource, - pub(crate) message_id: Option, - pub state: StartingReservationState, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -struct PendingWake { - slot_id: String, - role: TeamRunTargetRole, - source: TeamWakeSource, - message_id: Option, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum ChildStartDecision { - Accepted, - CancelImmediately(ActiveChildTurn), - Ignored, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum ChildCancelTarget { - Active(ActiveChildTurn), - Starting(StartingChildReservation), -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum WakeRecordDecision { - Recorded, - Suppressed, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) struct TeamRunOperationLease { - pub lease_id: String, - pub team_run_id: String, - pub slot_id: String, - pub role: TeamRunTargetRole, - pub wake_source: TeamWakeSource, - pub accepted_as_new_run: bool, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -struct ActiveOperationLease { - lease: TeamRunOperationLease, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) enum TeamRunWakeAcquireOutcome { - Accepted(TeamRunOperationLease), - Suppressed, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum TeamRunSlotState { - Busy, - Pending, - Paused, - Idle, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum AcquirePolicyDecision { - Accept, - Suppress(&'static str), - RejectSlotBusy, - RejectInvalid(&'static str), -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum TeamRunWakeIntent { - ExternalRequest, - SchedulerWakeTarget, -} - -#[derive(Debug, Clone)] -pub struct PauseSlotOutcome { - pub team_run_id: String, - pub cancel_target: Option, - pub payload: TeamRunPayload, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) struct PendingWakeView { - pub source: TeamWakeSource, - pub message_id: Option, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) struct RecoveryWakeCandidate { - pub slot_id: String, - pub role: TeamRunTargetRole, - pub unread_count: usize, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) struct RecoveryBacklogResult { - pub team_run_id: String, - pub source: TeamRunSource, - pub recorded_wakes: Vec, - pub pending_wake_count: usize, -} - -#[derive(Debug, Clone)] -struct TeamRunRecord { - team_run_id: String, - team_id: String, - source: TeamRunSource, - has_user_intervention: bool, - target_slot_id: String, - target_role: TeamRunTargetRole, - status: TeamRunStatus, - started_at: Option, - completed_at: Option, - cancelled_at: Option, - cancel_reason: Option, - active_child_turns: HashMap, - starting_reservations: HashMap, - pending_wakes: HashMap>, - slot_runtime_health: HashMap, - slot_wake_gate: SlotWakeGate, - active_operation_leases: HashMap, -} - -impl TeamRunRecord { - fn pending_wake_count(&self) -> usize { - self.pending_wakes.values().map(VecDeque::len).sum() - } - - fn pending_wake_count_for_slot(&self, slot_id: &str) -> usize { - self.pending_wakes.get(slot_id).map(VecDeque::len).unwrap_or(0) - } - - fn starting_child_count_for_slot(&self, slot_id: &str) -> usize { - self.starting_reservations - .values() - .filter(|reservation| reservation.slot_id == slot_id) - .count() - } - - fn active_operation_lease_count(&self) -> usize { - self.active_operation_leases.len() - } - - fn slot_run_state(&self, slot_id: &str) -> TeamRunSlotState { - if self.active_child_turns.contains_key(slot_id) || self.starting_child_count_for_slot(slot_id) > 0 { - return TeamRunSlotState::Busy; - } - if self.pending_wake_count_for_slot(slot_id) > 0 { - return TeamRunSlotState::Pending; - } - if self.slot_wake_gate.snapshot_for_slot(slot_id).paused { - return TeamRunSlotState::Paused; - } - TeamRunSlotState::Idle - } - - fn has_spawn_welcome_for_slot(&self, slot_id: &str) -> bool { - self.pending_wakes - .get(slot_id) - .is_some_and(|wakes| wakes.iter().any(|wake| wake.source == TeamWakeSource::SpawnWelcome)) - || self.active_operation_leases.values().any(|active| { - active.lease.slot_id == slot_id && active.lease.wake_source == TeamWakeSource::SpawnWelcome - }) - } - - fn role_for_slot(&self, slot_id: &str) -> Option { - self.active_child_turns - .get(slot_id) - .map(|child| child.role.clone()) - .or_else(|| { - self.starting_reservations - .values() - .find(|reservation| reservation.slot_id == slot_id) - .map(|reservation| reservation.role.clone()) - }) - .or_else(|| { - self.pending_wakes - .get(slot_id) - .and_then(|wakes| wakes.front()) - .map(|wake| wake.role.clone()) - }) - .or_else(|| self.slot_wake_gate.role_hint(slot_id)) - .or_else(|| (self.target_slot_id == slot_id).then(|| self.target_role.clone())) - } - - fn slot_work(&self) -> Vec { - let now = now_ms(); - let mut slot_ids = self - .pending_wakes - .keys() - .cloned() - .chain( - self.starting_reservations - .values() - .map(|reservation| reservation.slot_id.clone()), - ) - .chain(self.active_child_turns.keys().cloned()) - .chain(self.slot_runtime_health.keys().cloned()) - .chain(self.slot_wake_gate.slot_ids().cloned()) - .collect::>(); - slot_ids.sort(); - slot_ids.dedup(); - - slot_ids - .into_iter() - .filter_map(|slot_id| { - let role = self.role_for_slot(&slot_id)?; - let gate = self.slot_wake_gate.snapshot_for_slot(&slot_id); - let active_child = self.active_child_turns.get(&slot_id); - let active_elapsed_ms = active_child.map(|child| { - child - .last_slow_notified_at_ms - .unwrap_or(now) - .saturating_sub(child.started_at_ms) - .max(0) as u64 - }); - Some(TeamSlotWorkPayload { - pending_wake_count: self.pending_wake_count_for_slot(&slot_id), - starting_child_count: self.starting_child_count_for_slot(&slot_id), - paused: gate.paused, - suppressed_wake_count: gate.suppressed_wake_count, - active_turn_id: active_child.map(|child| child.turn_id.clone()), - active_turn_started_at_ms: active_child.map(|child| child.started_at_ms), - active_turn_elapsed_ms: active_elapsed_ms, - active_turn_slow: active_elapsed_ms.map(|elapsed| elapsed >= ACTIVE_CHILD_SLOW_THRESHOLD_MS), - active_turn_slow_threshold_ms: active_child.map(|_| ACTIVE_CHILD_SLOW_THRESHOLD_MS), - runtime_health: self.slot_runtime_health.get(&slot_id).cloned(), - slot_id, - role, - }) - }) - .collect() - } - - fn slot_is_busy(&self, slot_id: &str) -> bool { - self.pending_wake_count_for_slot(slot_id) > 0 - || self.starting_child_count_for_slot(slot_id) > 0 - || self.active_child_turns.contains_key(slot_id) - } - - fn has_retained_wake_gate_work(&self) -> bool { - self.slot_wake_gate.has_retained_work() - } - - fn payload(&self) -> TeamRunPayload { - TeamRunPayload { - team_id: self.team_id.clone(), - team_run_id: self.team_run_id.clone(), - source: self.source.clone(), - has_user_intervention: self.has_user_intervention, - target_slot_id: self.target_slot_id.clone(), - target_role: self.target_role.clone(), - status: self.status.clone(), - active_child_count: self.active_child_turns.len(), - pending_wake_count: self.pending_wake_count(), - starting_child_count: self.starting_reservations.len(), - slot_work: self.slot_work(), - } - } - - fn ack( - &self, - accepted_slot_id: &str, - accepted_role: TeamRunTargetRole, - message_id: Option, - ) -> TeamRunAckResponse { - TeamRunAckResponse { - team_run_id: self.team_run_id.clone(), - team_id: self.team_id.clone(), - source: self.source.clone(), - has_user_intervention: self.has_user_intervention, - target_slot_id: self.target_slot_id.clone(), - target_role: self.target_role.clone(), - accepted_slot_id: accepted_slot_id.to_owned(), - accepted_role, - status: self.status.clone(), - message_id, - } - } - - fn is_active(&self) -> bool { - matches!(self.status, TeamRunStatus::Accepted | TeamRunStatus::Running) - } -} - -fn new_operation_lease( - run: &mut TeamRunRecord, - slot_id: &str, - role: TeamRunTargetRole, - wake_source: TeamWakeSource, - accepted_as_new_run: bool, -) -> TeamRunOperationLease { - let lease = TeamRunOperationLease { - lease_id: generate_id(), - team_run_id: run.team_run_id.clone(), - slot_id: slot_id.to_owned(), - role, - wake_source, - accepted_as_new_run, - }; - run.active_operation_leases - .insert(lease.lease_id.clone(), ActiveOperationLease { lease: lease.clone() }); - lease -} - -fn new_team_run_record( - team_id: String, - target_slot_id: &str, - target_role: TeamRunTargetRole, - source: TeamRunSource, - has_user_intervention: bool, -) -> TeamRunRecord { - TeamRunRecord { - team_run_id: generate_id(), - team_id, - source, - has_user_intervention, - target_slot_id: target_slot_id.to_owned(), - target_role, - status: TeamRunStatus::Accepted, - started_at: None, - completed_at: None, - cancelled_at: None, - cancel_reason: None, - active_child_turns: HashMap::new(), - starting_reservations: HashMap::new(), - pending_wakes: HashMap::new(), - slot_runtime_health: HashMap::new(), - slot_wake_gate: SlotWakeGate::default(), - active_operation_leases: HashMap::new(), - } -} - -fn push_pending_wake_locked( - run: &mut TeamRunRecord, - slot_id: String, - role: TeamRunTargetRole, - source: TeamWakeSource, - message_id: Option, -) { - let wake = PendingWake { - slot_id: slot_id.clone(), - role, - source, - message_id, - }; - let queue = run.pending_wakes.entry(slot_id).or_default(); - if is_foreground_wake(source) { - let insert_at = queue - .iter() - .position(|pending| !is_foreground_wake(pending.source)) - .unwrap_or(queue.len()); - queue.insert(insert_at, wake); - } else { - queue.push_back(wake); - } -} - -fn is_foreground_wake(source: TeamWakeSource) -> bool { - matches!(source, TeamWakeSource::UserMessage | TeamWakeSource::UserIntervention) -} - -fn acquire_policy( - source: TeamWakeSource, - slot_state: TeamRunSlotState, - has_spawn_welcome: bool, - intent: TeamRunWakeIntent, -) -> AcquirePolicyDecision { - match source { - TeamWakeSource::UserMessage | TeamWakeSource::UserIntervention => match slot_state { - TeamRunSlotState::Busy => AcquirePolicyDecision::RejectSlotBusy, - TeamRunSlotState::Pending | TeamRunSlotState::Paused | TeamRunSlotState::Idle => { - AcquirePolicyDecision::Accept - } - }, - TeamWakeSource::McpSendMessage => match slot_state { - TeamRunSlotState::Paused => AcquirePolicyDecision::Suppress("paused_slot_background_wake"), - TeamRunSlotState::Busy | TeamRunSlotState::Pending | TeamRunSlotState::Idle => { - AcquirePolicyDecision::Accept - } - }, - TeamWakeSource::SpawnWelcome => { - if has_spawn_welcome { - return AcquirePolicyDecision::Suppress("duplicate_spawn_welcome"); - } - match slot_state { - TeamRunSlotState::Busy => AcquirePolicyDecision::RejectInvalid("spawn welcome target is already busy"), - TeamRunSlotState::Pending => AcquirePolicyDecision::Suppress("spawn_welcome_already_pending"), - TeamRunSlotState::Paused | TeamRunSlotState::Idle => AcquirePolicyDecision::Accept, - } - } - TeamWakeSource::McpShutdownRequest => AcquirePolicyDecision::Accept, - TeamWakeSource::IdleNotification | TeamWakeSource::InterruptedNotification => match slot_state { - TeamRunSlotState::Idle if intent == TeamRunWakeIntent::SchedulerWakeTarget => AcquirePolicyDecision::Accept, - TeamRunSlotState::Idle => AcquirePolicyDecision::Suppress("background_notification_without_wake_target"), - TeamRunSlotState::Busy | TeamRunSlotState::Pending | TeamRunSlotState::Paused => { - AcquirePolicyDecision::Suppress("background_notification_deduped") - } - }, - TeamWakeSource::RecoveryDrain - | TeamWakeSource::CrashNotification - | TeamWakeSource::InactivityTimeout - | TeamWakeSource::SpawnAttachFailure - | TeamWakeSource::ShutdownRejected => match slot_state { - TeamRunSlotState::Paused => AcquirePolicyDecision::Suppress("paused_slot_recovery_wake"), - TeamRunSlotState::Busy | TeamRunSlotState::Pending | TeamRunSlotState::Idle => { - AcquirePolicyDecision::Accept - } - }, - } -} - -#[derive(Clone)] -pub struct TeamRunManager { - team_id: String, - emitter: Arc, - state: Arc>>, -} - -impl TeamRunManager { - pub fn new(team_id: String, emitter: Arc) -> Self { - Self { - team_id, - emitter, - state: Arc::new(Mutex::new(None)), - } - } - - pub async fn accept_user_message( - &self, - target_slot_id: &str, - target_role: TeamRunTargetRole, - allow_active_intervention: bool, - message_id: Option, - ) -> Result { - let mut guard = self.state.lock().await; - if let Some(active) = guard.as_mut().filter(|r| r.is_active()) { - if allow_active_intervention { - if active.slot_is_busy(target_slot_id) { - info!( - team_id = %self.team_id, - team_run_id = %active.team_run_id, - target_slot_id, - target_role = ?target_role, - "team_run active intervention rejected because target slot is busy" - ); - return Err(TeamError::SlotBusy(target_slot_id.to_owned())); - } - debug!( - team_id = %self.team_id, - team_run_id = %active.team_run_id, - target_slot_id = %target_slot_id, - target_role = ?target_role, - active_target_slot_id = %active.target_slot_id, - active_target_role = ?active.target_role, - "team_run active intervention accepted" - ); - active.has_user_intervention = true; - return Ok(active.ack(target_slot_id, target_role, message_id)); - } - return Err(TeamError::InvalidRequest("team run is already active".into())); - } - if let Some(cancelling) = guard.as_ref().filter(|r| matches!(r.status, TeamRunStatus::Cancelling)) { - return Err(TeamError::InvalidRequest(format!( - "team run {} is cancelling", - cancelling.team_run_id - ))); - } - - let record = new_team_run_record( - self.team_id.clone(), - target_slot_id, - target_role.clone(), - TeamRunSource::UserMessage, - true, - ); - let ack = record.ack(target_slot_id, target_role, message_id); - let payload = record.payload(); - *guard = Some(record); - drop(guard); - - info!( - team_id = %self.team_id, - team_run_id = %ack.team_run_id, - target_slot_id = %ack.target_slot_id, - target_role = ?ack.target_role, - "team_run accepted" - ); - self.emitter.broadcast_team_run(TEAM_RUN_ACCEPTED_EVENT, payload); - Ok(ack) - } - - pub(crate) async fn acquire_user_message_wake( - &self, - slot_id: &str, - role: TeamRunTargetRole, - ) -> Result<(TeamRunAckResponse, TeamRunOperationLease), TeamError> { - let mut guard = self.state.lock().await; - - if let Some(run) = guard.as_mut().filter(|r| r.is_active()) { - if matches!(run.slot_run_state(slot_id), TeamRunSlotState::Busy) { - info!( - team_id = %self.team_id, - team_run_id = %run.team_run_id, - slot_id, - role = ?role, - "team_run user wake acquire rejected because slot is busy" - ); - return Err(TeamError::SlotBusy(slot_id.to_owned())); - } - - let source = TeamWakeSource::UserIntervention; - let _ = run.slot_wake_gate.before_wake(slot_id, source, None); - run.has_user_intervention = true; - let lease = new_operation_lease(run, slot_id, role.clone(), source, false); - let ack = run.ack(slot_id, role, None); - info!( - team_id = %self.team_id, - team_run_id = %lease.team_run_id, - lease_id = %lease.lease_id, - slot_id = %lease.slot_id, - wake_source = %lease.wake_source, - accepted_as_new_run = lease.accepted_as_new_run, - "team_run operation lease acquired" - ); - return Ok((ack, lease)); - } - - if let Some(cancelling) = guard.as_ref().filter(|r| matches!(r.status, TeamRunStatus::Cancelling)) { - return Err(TeamError::InvalidRequest(format!( - "team run {} is cancelling", - cancelling.team_run_id - ))); - } - - let mut record = new_team_run_record( - self.team_id.clone(), - slot_id, - role.clone(), - TeamRunSource::UserMessage, - true, - ); - let lease = new_operation_lease(&mut record, slot_id, role.clone(), TeamWakeSource::UserMessage, true); - let ack = record.ack(slot_id, role, None); - let payload = record.payload(); - *guard = Some(record); - drop(guard); - - info!( - team_id = %self.team_id, - team_run_id = %ack.team_run_id, - target_slot_id = %ack.target_slot_id, - target_role = ?ack.target_role, - "team_run accepted" - ); - info!( - team_id = %self.team_id, - team_run_id = %lease.team_run_id, - lease_id = %lease.lease_id, - slot_id = %lease.slot_id, - wake_source = %lease.wake_source, - accepted_as_new_run = lease.accepted_as_new_run, - "team_run operation lease acquired" - ); - self.emitter.broadcast_team_run(TEAM_RUN_ACCEPTED_EVENT, payload); - Ok((ack, lease)) - } - - pub(crate) async fn recover_mailbox_backlog( - &self, - candidates: Vec, - ) -> Option { - if candidates.is_empty() { - return None; - } - - let mut guard = self.state.lock().await; - if let Some(run) = guard.as_ref().filter(|r| matches!(r.status, TeamRunStatus::Cancelling)) { - warn!( - team_id = %self.team_id, - team_run_id = %run.team_run_id, - source = %TeamWakeSource::RecoveryDrain, - candidate_count = candidates.len(), - reason = "active_run_cancelling", - "team recovery backlog skipped" - ); - return None; - } - - if guard.as_ref().is_none_or(|run| !run.is_active()) { - let first = candidates.first().expect("checked non-empty"); - let mut run = new_team_run_record( - self.team_id.clone(), - &first.slot_id, - first.role.clone(), - TeamRunSource::RecoveryDrain, - false, - ); - for candidate in &candidates { - push_pending_wake_locked( - &mut run, - candidate.slot_id.clone(), - candidate.role.clone(), - TeamWakeSource::RecoveryDrain, - None, - ); - } - let payload = run.payload(); - let result = RecoveryBacklogResult { - team_run_id: run.team_run_id.clone(), - source: run.source.clone(), - recorded_wakes: candidates.iter().map(|candidate| candidate.slot_id.clone()).collect(), - pending_wake_count: payload.pending_wake_count, - }; - *guard = Some(run); - drop(guard); - - info!( - team_id = %self.team_id, - team_run_id = %result.team_run_id, - source = "recovery_drain", - slot_count = result.recorded_wakes.len(), - pending_wake_count = result.pending_wake_count, - reason = "orphan_mailbox_backlog", - "team recovery drain accepted" - ); - self.emitter.broadcast_team_run(TEAM_RUN_ACCEPTED_EVENT, payload); - return Some(result); - } - - let run = guard.as_mut().expect("active checked"); - let mut recorded_wakes = Vec::new(); - for candidate in candidates { - let slot_state = run.slot_run_state(&candidate.slot_id); - if !matches!(slot_state, TeamRunSlotState::Idle) { - debug!( - team_id = %self.team_id, - team_run_id = %run.team_run_id, - slot_id = %candidate.slot_id, - source = "recovery_drain", - unread_count = candidate.unread_count, - slot_state = ?slot_state, - reason = "slot_already_has_run_work", - "team recovery backlog wake skipped" - ); - continue; - } - push_pending_wake_locked( - run, - candidate.slot_id.clone(), - candidate.role, - TeamWakeSource::RecoveryDrain, - None, - ); - recorded_wakes.push(candidate.slot_id); - } - - if recorded_wakes.is_empty() { - return None; - } - - let payload = run.payload(); - let result = RecoveryBacklogResult { - team_run_id: run.team_run_id.clone(), - source: run.source.clone(), - recorded_wakes, - pending_wake_count: payload.pending_wake_count, - }; - info!( - team_id = %self.team_id, - team_run_id = %result.team_run_id, - source = "recovery_drain", - slot_count = result.recorded_wakes.len(), - pending_wake_count = result.pending_wake_count, - reason = "attached_to_active_run", - "team recovery backlog attached to active run" - ); - drop(guard); - self.emitter.broadcast_team_run(TEAM_RUN_UPDATED_EVENT, payload); - Some(result) - } - - pub(crate) async fn commit_operation_lease( - &self, - lease_id: &str, - trigger_message_id: Option, - ) -> Result<(), TeamError> { - let mut guard = self.state.lock().await; - let Some(run) = guard.as_mut().filter(|r| r.is_active()) else { - error!( - team_id = %self.team_id, - lease_id, - "team_run operation lease commit failed because no active run exists" - ); - return Err(TeamError::InvalidRequest(format!( - "team run operation lease missing: {lease_id}" - ))); - }; - let Some(active) = run.active_operation_leases.remove(lease_id) else { - error!( - team_id = %self.team_id, - team_run_id = %run.team_run_id, - lease_id, - "team_run operation lease commit failed because lease is missing" - ); - return Err(TeamError::InvalidRequest(format!( - "team run operation lease missing: {lease_id}" - ))); - }; - - let lease = active.lease; - push_pending_wake_locked( - run, - lease.slot_id.clone(), - lease.role.clone(), - lease.wake_source, - trigger_message_id, - ); - let payload = run.payload(); - info!( - team_id = %self.team_id, - team_run_id = %run.team_run_id, - lease_id = %lease.lease_id, - slot_id = %lease.slot_id, - wake_source = %lease.wake_source, - pending_wake_count = payload.pending_wake_count, - active_operation_lease_count = run.active_operation_lease_count(), - "team_run operation lease committed" - ); - drop(guard); - self.emitter.broadcast_team_run(TEAM_RUN_UPDATED_EVENT, payload); - Ok(()) - } - - pub(crate) async fn abort_operation_lease(&self, lease_id: &str, reason: &str) -> Result<(), TeamError> { - let mut guard = self.state.lock().await; - let Some(run) = guard.as_mut().filter(|r| r.is_active()) else { - error!( - team_id = %self.team_id, - lease_id, - reason, - "team_run operation lease abort failed because no active run exists" - ); - return Err(TeamError::InvalidRequest(format!( - "team run operation lease missing: {lease_id}" - ))); - }; - let Some(active) = run.active_operation_leases.remove(lease_id) else { - error!( - team_id = %self.team_id, - team_run_id = %run.team_run_id, - lease_id, - reason, - "team_run operation lease abort failed because lease is missing" - ); - return Err(TeamError::InvalidRequest(format!( - "team run operation lease missing: {lease_id}" - ))); - }; - let payload = run.payload(); - warn!( - team_id = %self.team_id, - team_run_id = %run.team_run_id, - lease_id = %active.lease.lease_id, - slot_id = %active.lease.slot_id, - wake_source = %active.lease.wake_source, - reason, - active_operation_lease_count = run.active_operation_lease_count(), - "team_run operation lease aborted" - ); - drop(guard); - self.emitter.broadcast_team_run(TEAM_RUN_UPDATED_EVENT, payload); - Ok(()) - } - - async fn acquire_run_scoped_wake_with_intent( - &self, - slot_id: &str, - role: TeamRunTargetRole, - source: TeamWakeSource, - intent: TeamRunWakeIntent, - ) -> Result { - let mut guard = self.state.lock().await; - let Some(run) = guard.as_mut().filter(|r| r.is_active()) else { - warn!( - team_id = %self.team_id, - slot_id, - role = ?role, - wake_source = %source, - "team_run wake acquire rejected because no active run exists" - ); - return Err(TeamError::InvalidRequest( - "no active team run for run-scoped wake".into(), - )); - }; - - let slot_state = run.slot_run_state(slot_id); - let decision = acquire_policy(source, slot_state, run.has_spawn_welcome_for_slot(slot_id), intent); - match decision { - AcquirePolicyDecision::RejectSlotBusy => Err(TeamError::SlotBusy(slot_id.to_owned())), - AcquirePolicyDecision::RejectInvalid(message) => Err(TeamError::InvalidRequest(message.into())), - AcquirePolicyDecision::Suppress(reason) => { - let _ = run.slot_wake_gate.before_wake(slot_id, source, None); - let payload = run.payload(); - debug!( - team_id = %self.team_id, - team_run_id = %run.team_run_id, - slot_id, - wake_source = %source, - slot_state = ?slot_state, - reason, - "team_run wake acquire suppressed" - ); - drop(guard); - self.emitter.broadcast_team_run(TEAM_RUN_UPDATED_EVENT, payload); - Ok(TeamRunWakeAcquireOutcome::Suppressed) - } - AcquirePolicyDecision::Accept => { - let _ = run.slot_wake_gate.before_wake(slot_id, source, None); - let lease = new_operation_lease(run, slot_id, role, source, false); - info!( - team_id = %self.team_id, - team_run_id = %lease.team_run_id, - lease_id = %lease.lease_id, - slot_id = %lease.slot_id, - wake_source = %lease.wake_source, - accepted_as_new_run = lease.accepted_as_new_run, - "team_run operation lease acquired" - ); - Ok(TeamRunWakeAcquireOutcome::Accepted(lease)) - } - } - } - - pub(crate) async fn acquire_run_scoped_wake( - &self, - slot_id: &str, - role: TeamRunTargetRole, - source: TeamWakeSource, - ) -> Result { - self.acquire_run_scoped_wake_with_intent(slot_id, role, source, TeamRunWakeIntent::ExternalRequest) - .await - } - - pub(crate) async fn acquire_scheduler_wake( - &self, - slot_id: &str, - role: TeamRunTargetRole, - source: TeamWakeSource, - ) -> Result { - self.acquire_run_scoped_wake_with_intent(slot_id, role, source, TeamRunWakeIntent::SchedulerWakeTarget) - .await - } - - pub async fn active_run_id(&self) -> Option { - let guard = self.state.lock().await; - guard.as_ref().filter(|r| r.is_active()).map(|r| r.team_run_id.clone()) - } - - pub async fn current_run_id(&self) -> Option { - let guard = self.state.lock().await; - guard.as_ref().map(|r| r.team_run_id.clone()) - } - - pub async fn active_child_turns(&self) -> Vec { - let guard = self.state.lock().await; - guard - .as_ref() - .map(|run| run.active_child_turns.values().cloned().collect()) - .unwrap_or_default() - } - - pub async fn current_payload(&self) -> Option { - let guard = self.state.lock().await; - guard.as_ref().filter(|run| run.is_active()).map(TeamRunRecord::payload) - } - - pub(crate) async fn slot_work_for_slot(&self, slot_id: &str) -> Option<(String, TeamSlotWorkPayload)> { - let guard = self.state.lock().await; - let run = guard.as_ref().filter(|run| run.is_active())?; - let payload = run.payload(); - let work = payload.slot_work.iter().find(|work| work.slot_id == slot_id).cloned()?; - Some((payload.team_run_id, work)) - } - - pub async fn mark_slot_runtime_health( - &self, - slot_id: &str, - health: TeamSlotRuntimeHealth, - ) -> Option { - let mut guard = self.state.lock().await; - let run = guard.as_mut().filter(|run| run.is_active())?; - run.slot_runtime_health.insert(slot_id.to_owned(), health); - let payload = run.payload(); - drop(guard); - self.emitter.broadcast_team_run(TEAM_RUN_UPDATED_EVENT, payload.clone()); - Some(payload) - } - - pub async fn observe_slow_child_turns(&self, now: TimestampMs) -> Option { - let mut guard = self.state.lock().await; - let run = guard.as_mut().filter(|run| run.is_active())?; - let mut observed = false; - - for child in run.active_child_turns.values_mut() { - let elapsed_ms = now.saturating_sub(child.started_at_ms).max(0) as u64; - if elapsed_ms < ACTIVE_CHILD_SLOW_THRESHOLD_MS { - continue; - } - let due = child - .last_slow_notified_at_ms - .map(|last| now.saturating_sub(last).max(0) as u64 >= ACTIVE_CHILD_SLOW_REPEAT_MS) - .unwrap_or(true); - if !due { - continue; - } - child.last_slow_notified_at_ms = Some(now); - observed = true; - info!( - team_id = %self.team_id, - team_run_id = %child.team_run_id, - slot_id = %child.slot_id, - role = ?child.role, - conversation_id = %child.conversation_id, - turn_id = %child.turn_id, - elapsed_ms, - slow_threshold_ms = ACTIVE_CHILD_SLOW_THRESHOLD_MS, - "team_child_turn slow" - ); - } - - if !observed { - return None; - } - - let payload = run.payload(); - drop(guard); - self.emitter.broadcast_team_run(TEAM_RUN_UPDATED_EVENT, payload.clone()); - Some(payload) - } - - pub async fn pause_slot_work(&self, slot_id: &str, reason: Option) -> Result { - let mut guard = self.state.lock().await; - let Some(run) = guard.as_mut().filter(|r| r.is_active()) else { - return Err(TeamError::InvalidRequest("no active team run to pause".into())); - }; - - let role = run - .role_for_slot(slot_id) - .ok_or_else(|| TeamError::AgentNotFound(slot_id.to_owned()))?; - let reason = reason.unwrap_or_else(|| "user_stop".into()); - - if let Some(child) = run.active_child_turns.get(slot_id).cloned() { - info!( - team_id = %self.team_id, - team_run_id = %run.team_run_id, - slot_id, - active_turn_id = %child.turn_id, - reason = %reason, - "team slot pause requested for active child" - ); - return Ok(PauseSlotOutcome { - team_run_id: run.team_run_id.clone(), - cancel_target: Some(ChildCancelTarget::Active(child)), - payload: run.payload(), - }); - } - - let pending_count = run.pending_wakes.remove(slot_id).map(|wakes| wakes.len()).unwrap_or(0); - run.slot_wake_gate.pause(slot_id, role, reason.clone()); - run.slot_wake_gate.add_suppressed(slot_id, pending_count); - - let cancel_target = take_starting_cancel_target_locked(run, slot_id); - let payload = run.payload(); - info!( - team_id = %self.team_id, - team_run_id = %run.team_run_id, - slot_id, - active_turn_id = Option::<&str>::None, - pending_wake_count = pending_count, - reason = %reason, - "team slot paused" - ); - let outcome = PauseSlotOutcome { - team_run_id: run.team_run_id.clone(), - cancel_target, - payload: payload.clone(), - }; - drop(guard); - self.emitter.broadcast_team_run(TEAM_RUN_UPDATED_EVENT, payload); - Ok(outcome) - } - - #[cfg(test)] - pub(crate) async fn record_pending_wake( - &self, - slot_id: &str, - target_role: TeamRunTargetRole, - wake_source: TeamWakeSource, - ) -> Result<(), TeamError> { - self.record_or_suppress_wake(slot_id, target_role, wake_source, None) - .await - .map(|_| ()) - } - - pub(crate) async fn record_or_suppress_wake( - &self, - slot_id: &str, - target_role: TeamRunTargetRole, - wake_source: TeamWakeSource, - trigger_message_id: Option, - ) -> Result { - let mut guard = self.state.lock().await; - let Some(run) = guard.as_mut().filter(|r| r.is_active()) else { - warn!( - team_id = %self.team_id, - slot_id, - target_role = ?target_role, - wake_source = %wake_source, - "team_run pending wake rejected because no active run exists" - ); - return Err(TeamError::InvalidRequest( - "no active team run for run-scoped wake".into(), - )); - }; - - match run - .slot_wake_gate - .before_wake(slot_id, wake_source, trigger_message_id.clone()) - { - WakeGateDecision::Suppress => { - let suppressed_wake_count = run.slot_wake_gate.snapshot_for_slot(slot_id).suppressed_wake_count; - let payload = run.payload(); - info!( - team_id = %self.team_id, - team_run_id = %run.team_run_id, - slot_id, - wake_source = %wake_source, - suppressed_wake_count, - "team wake suppressed" - ); - drop(guard); - self.emitter.broadcast_team_run(TEAM_RUN_UPDATED_EVENT, payload); - Ok(WakeRecordDecision::Suppressed) - } - WakeGateDecision::Record { resumed_from_pause } => { - push_pending_wake_locked( - run, - slot_id.to_owned(), - target_role.clone(), - wake_source, - trigger_message_id.clone(), - ); - let slot_pending_wake_count = run.pending_wake_count_for_slot(slot_id); - let payload = run.payload(); - if resumed_from_pause { - info!( - team_id = %self.team_id, - team_run_id = %run.team_run_id, - slot_id, - resume_source = %wake_source, - foreground_message_id = ?trigger_message_id.as_deref(), - "team slot resumed" - ); - } - info!( - team_id = %self.team_id, - team_run_id = %run.team_run_id, - slot_id, - target_role = ?target_role, - wake_source = %wake_source, - slot_pending_wake_count, - pending_wake_count = payload.pending_wake_count, - starting_child_count = payload.starting_child_count, - active_child_count = payload.active_child_count, - "team_run pending wake recorded" - ); - drop(guard); - self.emitter.broadcast_team_run(TEAM_RUN_UPDATED_EVENT, payload); - Ok(WakeRecordDecision::Recorded) - } - } - } - - pub async fn claim_wake_for_turn( - &self, - slot_id: &str, - role: TeamRunTargetRole, - conversation_id: &str, - ) -> Option { - let mut guard = self.state.lock().await; - let run = guard.as_mut().filter(|r| r.is_active())?; - let pending = match run.pending_wakes.get_mut(slot_id).and_then(VecDeque::pop_front) { - Some(pending) => pending, - None => { - warn!( - team_id = %self.team_id, - team_run_id = %run.team_run_id, - slot_id, - role = ?role, - pending_wake_count = run.pending_wake_count(), - "team_run reservation claim ignored because no pending wake exists for slot" - ); - return None; - } - }; - if run.pending_wakes.get(slot_id).is_some_and(VecDeque::is_empty) { - run.pending_wakes.remove(slot_id); - } - run.slot_runtime_health.remove(slot_id); - if pending.slot_id != slot_id { - warn!( - team_id = %self.team_id, - team_run_id = %run.team_run_id, - expected_slot_id = %pending.slot_id, - actual_slot_id = %slot_id, - wake_source = %pending.source, - "team_run reservation claimed with slot mismatch" - ); - } - if pending.role != role { - warn!( - team_id = %self.team_id, - team_run_id = %run.team_run_id, - slot_id, - expected_role = ?pending.role, - actual_role = ?role, - wake_source = %pending.source, - "team_run reservation claimed with role mismatch" - ); - } - let reservation = StartingChildReservation { - reservation_id: generate_id(), - team_run_id: run.team_run_id.clone(), - slot_id: pending.slot_id.clone(), - role, - conversation_id: conversation_id.to_owned(), - wake_source: pending.source, - message_id: pending.message_id.clone(), - state: StartingReservationState::Starting, - }; - run.starting_reservations - .insert(reservation.reservation_id.clone(), reservation.clone()); - let payload = run.payload(); - info!( - team_id = %self.team_id, - team_run_id = %reservation.team_run_id, - reservation_id = %reservation.reservation_id, - slot_id = %reservation.slot_id, - role = ?reservation.role, - wake_source = %pending.source, - slot_pending_wake_count = run.pending_wake_count_for_slot(slot_id), - pending_wake_count = payload.pending_wake_count, - starting_child_count = payload.starting_child_count, - active_child_count = payload.active_child_count, - "team_run reservation claimed" - ); - drop(guard); - self.emitter.broadcast_team_run(TEAM_RUN_UPDATED_EVENT, payload); - Some(reservation) - } - - pub async fn retry_child_start_later(&self, reservation_id: &str, reason: &str) -> Option { - let mut guard = self.state.lock().await; - let run = guard.as_mut().filter(|run| run.is_active())?; - let reservation = match run.starting_reservations.remove(reservation_id) { - Some(reservation) => reservation, - None => { - warn!( - team_id = %self.team_id, - reservation_id, - error = %reason, - "team_run reservation retry ignored because reservation is missing" - ); - return None; - } - }; - - let pending = PendingWake { - slot_id: reservation.slot_id.clone(), - role: reservation.role, - source: reservation.wake_source, - message_id: reservation.message_id.clone(), - }; - run.pending_wakes - .entry(reservation.slot_id.clone()) - .or_default() - .push_front(pending); - let payload = run.payload(); - info!( - team_id = %self.team_id, - team_run_id = %run.team_run_id, - reservation_id = %reservation.reservation_id, - slot_id = %reservation.slot_id, - error = %reason, - slot_pending_wake_count = run.pending_wake_count_for_slot(&reservation.slot_id), - pending_wake_count = payload.pending_wake_count, - starting_child_count = payload.starting_child_count, - active_child_count = payload.active_child_count, - "team_run reservation deferred for retry" - ); - drop(guard); - self.emitter.broadcast_team_run(TEAM_RUN_UPDATED_EVENT, payload.clone()); - Some(payload) - } - - pub(crate) async fn peek_next_pending_wake(&self, slot_id: &str) -> Option { - let guard = self.state.lock().await; - guard - .as_ref() - .filter(|r| r.is_active()) - .and_then(|run| run.pending_wakes.get(slot_id)) - .and_then(|wakes| wakes.front()) - .map(|wake| PendingWakeView { - source: wake.source, - message_id: wake.message_id.clone(), - }) - } - - pub async fn record_empty_wake_observed(&self, slot_id: &str) -> Option { - let mut guard = self.state.lock().await; - let run = guard.as_mut()?; - let consumed = match run.pending_wakes.get_mut(slot_id).and_then(VecDeque::pop_front) { - Some(pending) => { - if run.pending_wakes.get(slot_id).is_some_and(VecDeque::is_empty) { - run.pending_wakes.remove(slot_id); - } - Some(pending) - } - None => None, - }; - let payload_before_completion = run.payload(); - debug!( - team_id = %self.team_id, - team_run_id = %run.team_run_id, - slot_id, - consumed_wake_source = ?consumed.as_ref().map(|wake| wake.source.as_str()), - slot_pending_wake_count = run.pending_wake_count_for_slot(slot_id), - pending_wake_count = payload_before_completion.pending_wake_count, - starting_child_count = payload_before_completion.starting_child_count, - active_child_count = payload_before_completion.active_child_count, - "team_run empty mailbox wake observed" - ); - let payload = maybe_complete_locked(run, &self.emitter); - if payload.is_some() { - *guard = None; - } - payload - } - - pub async fn record_child_started(&self, reservation_id: &str, child: ActiveChildTurn) -> ChildStartDecision { - let mut guard = self.state.lock().await; - let Some(run) = guard.as_mut() else { - warn!( - team_id = %self.team_id, - team_run_id = %child.team_run_id, - reservation_id, - slot_id = %child.slot_id, - turn_id = %child.turn_id, - "team_run child start ignored because no run is active" - ); - return ChildStartDecision::Ignored; - }; - let Some(reservation) = run.starting_reservations.remove(reservation_id) else { - warn!( - team_id = %self.team_id, - team_run_id = %child.team_run_id, - reservation_id, - slot_id = %child.slot_id, - turn_id = %child.turn_id, - "team_run child start ignored because reservation is missing" - ); - return ChildStartDecision::Ignored; - }; - if run.team_run_id != child.team_run_id { - warn!( - team_id = %self.team_id, - expected_team_run_id = %run.team_run_id, - actual_team_run_id = %child.team_run_id, - reservation_id, - slot_id = %child.slot_id, - turn_id = %child.turn_id, - "team_run child start ignored because run id mismatched" - ); - return ChildStartDecision::Ignored; - } - - let should_cancel = matches!(run.status, TeamRunStatus::Cancelling) - || matches!(reservation.state, StartingReservationState::Cancelling); - let first_child_for_run = run.started_at.is_none(); - if first_child_for_run { - run.started_at = Some(now_ms()); - } - if !should_cancel { - run.status = TeamRunStatus::Running; - } - run.active_child_turns.insert(child.slot_id.clone(), child.clone()); - let run_payload = run.payload(); - let child_payload = child_payload(&run.team_id, &child, TeamRunStatus::Running); - drop(guard); - - if first_child_for_run && !should_cancel { - info!( - team_id = %self.team_id, - team_run_id = %child.team_run_id, - target_slot_id = %run_payload.target_slot_id, - target_role = ?run_payload.target_role, - active_child_count = run_payload.active_child_count, - pending_wake_count = run_payload.pending_wake_count, - starting_child_count = run_payload.starting_child_count, - "team_run started" - ); - self.emitter - .broadcast_team_run(TEAM_RUN_STARTED_EVENT, run_payload.clone()); - } else { - self.emitter - .broadcast_team_run(TEAM_RUN_UPDATED_EVENT, run_payload.clone()); - } - info!( - team_id = %self.team_id, - team_run_id = %child.team_run_id, - reservation_id, - slot_id = %child.slot_id, - role = ?child.role, - conversation_id = %child.conversation_id, - turn_id = %child.turn_id, - cancelling = should_cancel, - active_child_count = run_payload.active_child_count, - pending_wake_count = run_payload.pending_wake_count, - starting_child_count = run_payload.starting_child_count, - "team_child_turn started" - ); - self.emitter - .broadcast_child_turn(TEAM_CHILD_TURN_STARTED_EVENT, child_payload); - if should_cancel { - ChildStartDecision::CancelImmediately(child) - } else { - ChildStartDecision::Accepted - } - } - - pub async fn record_child_start_failed(&self, reservation_id: &str, reason: &str) -> Option { - let mut guard = self.state.lock().await; - let run = guard.as_mut()?; - let Some(reservation) = run.starting_reservations.remove(reservation_id) else { - warn!( - team_id = %self.team_id, - reservation_id, - error = %reason, - "team_run reservation start failure ignored because reservation is missing" - ); - return None; - }; - - if matches!(run.status, TeamRunStatus::Cancelling) - || matches!(reservation.state, StartingReservationState::Cancelling) - { - let payload = maybe_cancelled_locked(run, &self.emitter); - if payload.is_some() { - *guard = None; - } - return payload; - } - - warn!( - team_id = %self.team_id, - team_run_id = %run.team_run_id, - reservation_id = %reservation.reservation_id, - slot_id = %reservation.slot_id, - error = %reason, - active_child_count = run.active_child_turns.len(), - pending_wake_count = run.pending_wake_count(), - starting_child_count = run.starting_reservations.len(), - "team_run reservation failed before start" - ); - run.status = TeamRunStatus::Failed; - run.completed_at = Some(now_ms()); - let payload = run.payload(); - *guard = None; - drop(guard); - self.emitter.broadcast_team_run(TEAM_RUN_FAILED_EVENT, payload.clone()); - Some(payload) - } - - pub async fn record_child_completed( - &self, - slot_id: &str, - turn_id: &str, - status: TeamRunStatus, - ) -> Option { - let mut guard = self.state.lock().await; - let run = guard.as_mut()?; - let child = run.active_child_turns.remove(slot_id)?; - if child.turn_id != turn_id { - run.active_child_turns.insert(slot_id.to_owned(), child); - return None; - } - - let child_payload = child_payload(&run.team_id, &child, status.clone()); - match status { - TeamRunStatus::Failed => { - warn!( - team_id = %self.team_id, - team_run_id = %run.team_run_id, - slot_id = %child.slot_id, - role = ?child.role, - conversation_id = %child.conversation_id, - turn_id = %child.turn_id, - active_child_count = run.active_child_turns.len(), - pending_wake_count = run.pending_wake_count(), - starting_child_count = run.starting_reservations.len(), - "team_child_turn failed" - ); - run.status = TeamRunStatus::Failed; - run.completed_at = Some(now_ms()); - let payload = run.payload(); - *guard = None; - drop(guard); - self.emitter - .broadcast_child_turn(TEAM_CHILD_TURN_COMPLETED_EVENT, child_payload); - warn!( - team_id = %payload.team_id, - team_run_id = %payload.team_run_id, - target_slot_id = %payload.target_slot_id, - target_role = ?payload.target_role, - active_child_count = payload.active_child_count, - pending_wake_count = payload.pending_wake_count, - starting_child_count = payload.starting_child_count, - "team_run failed" - ); - self.emitter.broadcast_team_run(TEAM_RUN_FAILED_EVENT, payload.clone()); - Some(payload) - } - _ => { - info!( - team_id = %self.team_id, - team_run_id = %run.team_run_id, - slot_id = %child.slot_id, - role = ?child.role, - conversation_id = %child.conversation_id, - turn_id = %child.turn_id, - status = ?status, - active_child_count = run.active_child_turns.len(), - pending_wake_count = run.pending_wake_count(), - starting_child_count = run.starting_reservations.len(), - "team_child_turn completed" - ); - self.emitter - .broadcast_child_turn(TEAM_CHILD_TURN_COMPLETED_EVENT, child_payload); - let payload = maybe_complete_locked(run, &self.emitter); - if payload.is_some() { - *guard = None; - } - payload - } - } - } - - pub async fn maybe_complete(&self) -> Option { - let mut guard = self.state.lock().await; - let run = guard.as_mut()?; - let payload = maybe_complete_locked(run, &self.emitter)?; - *guard = None; - Some(payload) - } - - pub async fn begin_cancel(&self, target_slot_id: Option, reason: Option) -> Result<(), TeamError> { - let mut guard = self.state.lock().await; - let Some(run) = guard.as_mut().filter(|r| r.is_active()) else { - return Err(TeamError::InvalidRequest("no active team run to cancel".into())); - }; - if let Some(target) = target_slot_id.as_deref() - && target != run.target_slot_id - && !run.active_child_turns.contains_key(target) - && !run - .starting_reservations - .values() - .any(|reservation| reservation.slot_id == target) - { - return Err(TeamError::AgentNotFound(target.to_owned())); - } - run.status = TeamRunStatus::Cancelling; - run.cancel_reason = reason; - run.pending_wakes.clear(); - run.slot_wake_gate.clear(); - for reservation in run.starting_reservations.values_mut() { - reservation.state = StartingReservationState::Cancelling; - } - let payload = run.payload(); - info!( - team_id = %self.team_id, - team_run_id = %run.team_run_id, - target_slot_id = ?target_slot_id.as_deref(), - active_child_count = run.active_child_turns.len(), - starting_child_count = run.starting_reservations.len(), - pending_wake_count = payload.pending_wake_count, - "team_run cancel requested" - ); - drop(guard); - - self.emitter.broadcast_team_run(TEAM_RUN_UPDATED_EVENT, payload); - Ok(()) - } - - pub async fn try_complete_cancelled(&self) -> Option { - let mut guard = self.state.lock().await; - let run = guard.as_mut()?; - let payload = maybe_cancelled_locked(run, &self.emitter)?; - *guard = None; - Some(payload) - } - - pub async fn complete_failed(&self) -> Option { - let mut guard = self.state.lock().await; - let run = guard.as_mut()?; - run.status = TeamRunStatus::Failed; - run.completed_at = Some(now_ms()); - let team_run_id = run.team_run_id.clone(); - let payload = run.payload(); - *guard = None; - drop(guard); - - warn!( - team_id = %self.team_id, - team_run_id = %team_run_id, - active_child_count = payload.active_child_count, - pending_wake_count = payload.pending_wake_count, - starting_child_count = payload.starting_child_count, - "team_run failed" - ); - self.emitter.broadcast_team_run(TEAM_RUN_FAILED_EVENT, payload); - Some(team_run_id) - } - - pub async fn begin_cancel_child(&self, slot_id: &str) -> Result { - let mut guard = self.state.lock().await; - let Some(run) = guard.as_mut().filter(|r| r.is_active()) else { - return Err(TeamError::InvalidRequest("no active team run".into())); - }; - if let Some(child) = run.active_child_turns.get(slot_id).cloned() { - info!( - team_id = %self.team_id, - team_run_id = %run.team_run_id, - slot_id, - turn_id = %child.turn_id, - state = "active", - active_child_count = run.active_child_turns.len(), - pending_wake_count = run.pending_wake_count(), - starting_child_count = run.starting_reservations.len(), - "team_run child cancel requested" - ); - return Ok(ChildCancelTarget::Active(child)); - } - if let Some(reservation) = run - .starting_reservations - .values_mut() - .find(|reservation| reservation.slot_id == slot_id) - { - reservation.state = StartingReservationState::Cancelling; - let reservation = reservation.clone(); - let team_run_id = run.team_run_id.clone(); - let active_child_count = run.active_child_turns.len(); - let pending_wake_count = run.pending_wake_count(); - let starting_child_count = run.starting_reservations.len(); - info!( - team_id = %self.team_id, - team_run_id = %team_run_id, - slot_id, - reservation_id = %reservation.reservation_id, - state = "starting", - active_child_count, - pending_wake_count, - starting_child_count, - "team_run child cancel requested" - ); - return Ok(ChildCancelTarget::Starting(reservation)); - } - Err(TeamError::InvalidRequest(format!( - "agent {slot_id} has no active or starting child turn" - ))) - } - - pub async fn record_child_cancelled(&self, child: &ActiveChildTurn) { - let mut guard = self.state.lock().await; - let Some(run) = guard.as_mut() else { - return; - }; - run.active_child_turns.remove(&child.slot_id); - let payload = child_payload(&run.team_id, child, TeamRunStatus::Cancelled); - let run_payload = if matches!(run.status, TeamRunStatus::Cancelling) { - maybe_cancelled_locked(run, &self.emitter) - } else { - None - }; - let counts_payload = run.payload(); - if run_payload.is_some() { - *guard = None; - } - drop(guard); - - info!( - team_id = %self.team_id, - team_run_id = %child.team_run_id, - slot_id = %child.slot_id, - role = ?child.role, - conversation_id = %child.conversation_id, - turn_id = %child.turn_id, - active_child_count = counts_payload.active_child_count, - pending_wake_count = counts_payload.pending_wake_count, - starting_child_count = counts_payload.starting_child_count, - "team_child_turn cancelled" - ); - self.emitter - .broadcast_child_turn(TEAM_CHILD_TURN_CANCELLED_EVENT, payload); - } - - pub async fn complete_pause_after_child_cancelled( - &self, - child: &ActiveChildTurn, - reason: Option, - ) -> Option { - let mut guard = self.state.lock().await; - let run = guard.as_mut()?; - let reason = reason.unwrap_or_else(|| "user_stop".into()); - let pending_count = run - .pending_wakes - .remove(&child.slot_id) - .map(|wakes| wakes.len()) - .unwrap_or(0); - run.slot_wake_gate - .pause(&child.slot_id, child.role.clone(), reason.clone()); - run.slot_wake_gate.add_suppressed(&child.slot_id, pending_count); - run.active_child_turns.remove(&child.slot_id); - - let child_payload = child_payload(&run.team_id, child, TeamRunStatus::Cancelled); - let payload = run.payload(); - info!( - team_id = %self.team_id, - team_run_id = %child.team_run_id, - slot_id = %child.slot_id, - role = ?child.role, - conversation_id = %child.conversation_id, - turn_id = %child.turn_id, - pending_wake_count = pending_count, - reason = %reason, - "team slot paused" - ); - info!( - team_id = %self.team_id, - team_run_id = %child.team_run_id, - slot_id = %child.slot_id, - role = ?child.role, - conversation_id = %child.conversation_id, - turn_id = %child.turn_id, - active_child_count = payload.active_child_count, - pending_wake_count = payload.pending_wake_count, - starting_child_count = payload.starting_child_count, - "team_child_turn cancelled" - ); - drop(guard); - - self.emitter - .broadcast_child_turn(TEAM_CHILD_TURN_CANCELLED_EVENT, child_payload); - self.emitter.broadcast_team_run(TEAM_RUN_UPDATED_EVENT, payload.clone()); - Some(payload) - } - - pub(crate) async fn release_suppressed_wake_if_resumed( - &self, - slot_id: &str, - role: TeamRunTargetRole, - ) -> Option { - let mut guard = self.state.lock().await; - let run = guard.as_mut().filter(|r| r.is_active())?; - let source = run.slot_wake_gate.release_suppressed_if_resumed(slot_id)?; - push_pending_wake_locked(run, slot_id.to_owned(), role.clone(), source, None); - let payload = run.payload(); - let slot_work = payload.slot_work.iter().find(|work| work.slot_id == slot_id); - info!( - team_id = %self.team_id, - team_run_id = %payload.team_run_id, - slot_id, - role = ?role, - released_wake_source = %source, - pending_wake_count = payload.pending_wake_count, - slot_pending_wake_count = slot_work.map(|work| work.pending_wake_count).unwrap_or_default(), - suppressed_wake_count = slot_work.map(|work| work.suppressed_wake_count).unwrap_or_default(), - "team suppressed wake released" - ); - drop(guard); - self.emitter.broadcast_team_run(TEAM_RUN_UPDATED_EVENT, payload); - Some(source) - } -} - -pub fn target_role_for(role: TeammateRole) -> TeamRunTargetRole { - match role { - TeammateRole::Lead => TeamRunTargetRole::Lead, - TeammateRole::Teammate => TeamRunTargetRole::Teammate, - } -} - -fn child_payload(team_id: &str, child: &ActiveChildTurn, status: TeamRunStatus) -> TeamChildTurnPayload { - TeamChildTurnPayload { - team_id: team_id.to_owned(), - team_run_id: child.team_run_id.clone(), - slot_id: child.slot_id.clone(), - role: child.role.clone(), - conversation_id: child.conversation_id.clone(), - turn_id: child.turn_id.clone(), - status, - } -} - -fn take_starting_cancel_target_locked(run: &mut TeamRunRecord, slot_id: &str) -> Option { - let reservation_id = run - .starting_reservations - .iter() - .find_map(|(id, reservation)| (reservation.slot_id == slot_id).then(|| id.clone())); - reservation_id.and_then(|id| run.starting_reservations.remove(&id).map(ChildCancelTarget::Starting)) -} - -fn maybe_complete_locked(run: &mut TeamRunRecord, emitter: &TeamEventEmitter) -> Option { - if run.pending_wake_count() > 0 - || !run.starting_reservations.is_empty() - || !run.active_child_turns.is_empty() - || run.active_operation_lease_count() > 0 - || run.has_retained_wake_gate_work() - { - emitter.broadcast_team_run(TEAM_RUN_UPDATED_EVENT, run.payload()); - return None; - } - if !matches!(run.status, TeamRunStatus::Running | TeamRunStatus::Accepted) { - return None; - } - - run.status = TeamRunStatus::Completed; - run.completed_at = Some(now_ms()); - let payload = run.payload(); - info!( - team_id = %payload.team_id, - team_run_id = %payload.team_run_id, - target_slot_id = %payload.target_slot_id, - target_role = ?payload.target_role, - active_child_count = payload.active_child_count, - pending_wake_count = payload.pending_wake_count, - starting_child_count = payload.starting_child_count, - "team_run completed" - ); - emitter.broadcast_team_run(TEAM_RUN_COMPLETED_EVENT, payload.clone()); - Some(payload) -} - -fn maybe_cancelled_locked(run: &mut TeamRunRecord, emitter: &TeamEventEmitter) -> Option { - if !matches!(run.status, TeamRunStatus::Cancelling) { - return None; - } - if run.pending_wake_count() > 0 - || !run.starting_reservations.is_empty() - || !run.active_child_turns.is_empty() - || run.active_operation_lease_count() > 0 - || run.has_retained_wake_gate_work() - { - emitter.broadcast_team_run(TEAM_RUN_UPDATED_EVENT, run.payload()); - return None; - } - - run.status = TeamRunStatus::Cancelled; - run.cancelled_at = Some(now_ms()); - let payload = run.payload(); - info!( - team_id = %payload.team_id, - team_run_id = %payload.team_run_id, - target_slot_id = %payload.target_slot_id, - target_role = ?payload.target_role, - active_child_count = payload.active_child_count, - pending_wake_count = payload.pending_wake_count, - starting_child_count = payload.starting_child_count, - "team_run cancelled" - ); - emitter.broadcast_team_run(TEAM_RUN_CANCELLED_EVENT, payload.clone()); - Some(payload) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::wake::TeamWakeSource; - use aionui_api_types::WebSocketMessage; - use aionui_realtime::EventBroadcaster; - - #[derive(Default)] - struct RecordingBroadcaster { - events: std::sync::Mutex>>, - } - - impl RecordingBroadcaster { - fn events(&self) -> Vec> { - self.events.lock().unwrap().clone() - } - - fn names(&self) -> Vec { - self.events - .lock() - .unwrap() - .iter() - .map(|event| event.name.clone()) - .collect() - } - - fn run_payloads(&self) -> Vec { - self.events() - .into_iter() - .filter(|event| event.name.starts_with("team.run")) - .map(|event| serde_json::from_value(event.data).unwrap()) - .collect() - } - } - - impl EventBroadcaster for RecordingBroadcaster { - fn broadcast(&self, event: WebSocketMessage) { - self.events.lock().unwrap().push(event); - } - } - - fn manager() -> (TeamRunManager, Arc) { - let bc = Arc::new(RecordingBroadcaster::default()); - let emitter = Arc::new(TeamEventEmitter::new("team-1".into(), bc.clone())); - (TeamRunManager::new("team-1".into(), emitter), bc) - } - - fn slot_work<'a>(payload: &'a TeamRunPayload, slot_id: &str) -> &'a aionui_api_types::TeamSlotWorkPayload { - payload - .slot_work - .iter() - .find(|work| work.slot_id == slot_id) - .expect("slot work must exist") - } - - #[tokio::test] - async fn user_message_run_payload_has_user_source() { - let (manager, _bc) = manager(); - let (ack, lease) = manager - .acquire_user_message_wake("lead", TeamRunTargetRole::Lead) - .await - .expect("user message should create run"); - - assert_eq!(ack.source, TeamRunSource::UserMessage); - assert!(ack.has_user_intervention); - - manager - .commit_operation_lease(&lease.lease_id, Some("mailbox-1".into())) - .await - .expect("commit user wake"); - - let payload = manager.current_payload().await.expect("active payload"); - assert_eq!(payload.source, TeamRunSource::UserMessage); - assert!(payload.has_user_intervention); - } - - #[tokio::test] - async fn active_recovery_run_records_user_intervention_without_changing_source() { - let (manager, _bc) = manager(); - let result = manager - .recover_mailbox_backlog(vec![RecoveryWakeCandidate { - slot_id: "lead".into(), - role: TeamRunTargetRole::Lead, - unread_count: 2, - }]) - .await - .expect("recovery scan should succeed"); - - assert_eq!(result.source, TeamRunSource::RecoveryDrain); - assert_eq!( - manager.current_payload().await.unwrap().source, - TeamRunSource::RecoveryDrain - ); - - let (ack, lease) = manager - .acquire_user_message_wake("worker", TeamRunTargetRole::Teammate) - .await - .expect("user intervention should join recovery run"); - assert_eq!(ack.team_run_id, result.team_run_id); - assert_eq!(ack.source, TeamRunSource::RecoveryDrain); - assert!(ack.has_user_intervention); - assert!(!lease.accepted_as_new_run); - - let payload = manager.current_payload().await.expect("active payload"); - assert_eq!(payload.source, TeamRunSource::RecoveryDrain); - assert!(payload.has_user_intervention); - } - - #[tokio::test] - async fn recovery_drain_creates_run_with_pending_wakes() { - let (manager, _bc) = manager(); - let result = manager - .recover_mailbox_backlog(vec![ - RecoveryWakeCandidate { - slot_id: "lead".into(), - role: TeamRunTargetRole::Lead, - unread_count: 2, - }, - RecoveryWakeCandidate { - slot_id: "worker".into(), - role: TeamRunTargetRole::Teammate, - unread_count: 1, - }, - ]) - .await - .expect("recovery should create run"); - - assert_eq!(result.source, TeamRunSource::RecoveryDrain); - assert_eq!(result.recorded_wakes.len(), 2); - assert_eq!(result.pending_wake_count, 2); - - let lead = manager - .claim_wake_for_turn("lead", TeamRunTargetRole::Lead, "conv-lead") - .await - .expect("lead recovery wake should be claimable"); - assert_eq!(lead.team_run_id, result.team_run_id); - assert_eq!(lead.wake_source, TeamWakeSource::RecoveryDrain); - assert!(lead.message_id.is_none()); - } - - #[tokio::test] - async fn recovery_backlog_attaches_to_active_user_run() { - let (manager, _bc) = manager(); - let (ack, lease) = manager - .acquire_user_message_wake("lead", TeamRunTargetRole::Lead) - .await - .expect("user run"); - manager - .commit_operation_lease(&lease.lease_id, Some("mailbox-user".into())) - .await - .expect("commit user wake"); - - let result = manager - .recover_mailbox_backlog(vec![RecoveryWakeCandidate { - slot_id: "worker".into(), - role: TeamRunTargetRole::Teammate, - unread_count: 3, - }]) - .await - .expect("recovery should attach to active run"); - - assert_eq!(result.team_run_id, ack.team_run_id); - assert_eq!(result.source, TeamRunSource::UserMessage); - assert_eq!(result.recorded_wakes, vec!["worker".to_string()]); - - let worker = manager - .claim_wake_for_turn("worker", TeamRunTargetRole::Teammate, "conv-worker") - .await - .expect("attached recovery wake should be claimable"); - assert_eq!(worker.wake_source, TeamWakeSource::RecoveryDrain); - assert!(worker.message_id.is_none()); - } - - #[tokio::test] - async fn user_intervention_wake_prioritizes_over_recovery_backlog() { - let (manager, _bc) = manager(); - let result = manager - .recover_mailbox_backlog(vec![RecoveryWakeCandidate { - slot_id: "lead".into(), - role: TeamRunTargetRole::Lead, - unread_count: 2, - }]) - .await - .expect("recovery should create run"); - - let (ack, lease) = manager - .acquire_user_message_wake("lead", TeamRunTargetRole::Lead) - .await - .expect("user intervention should join recovery run"); - assert_eq!(ack.team_run_id, result.team_run_id); - assert_eq!(ack.source, TeamRunSource::RecoveryDrain); - assert!(ack.has_user_intervention); - assert!(!lease.accepted_as_new_run); - - manager - .commit_operation_lease(&lease.lease_id, Some("mailbox-user".into())) - .await - .expect("commit user intervention wake"); - - let foreground = manager - .claim_wake_for_turn("lead", TeamRunTargetRole::Lead, "conv-lead-user") - .await - .expect("foreground user wake should be claimed first"); - assert_eq!(foreground.team_run_id, result.team_run_id); - assert_eq!(foreground.wake_source, TeamWakeSource::UserIntervention); - assert_eq!(foreground.message_id.as_deref(), Some("mailbox-user")); - - let recovery = manager - .claim_wake_for_turn("lead", TeamRunTargetRole::Lead, "conv-lead-recovery") - .await - .expect("recovery backlog should remain after foreground wake"); - assert_eq!(recovery.team_run_id, result.team_run_id); - assert_eq!(recovery.wake_source, TeamWakeSource::RecoveryDrain); - assert!(recovery.message_id.is_none()); - } - - #[tokio::test] - async fn recovery_backlog_does_not_duplicate_pending_slot_work() { - let (manager, _bc) = manager(); - let (_ack, lease) = manager - .acquire_user_message_wake("worker", TeamRunTargetRole::Teammate) - .await - .expect("user run"); - manager - .commit_operation_lease(&lease.lease_id, Some("mailbox-worker".into())) - .await - .expect("commit pending user wake"); - - let result = manager - .recover_mailbox_backlog(vec![RecoveryWakeCandidate { - slot_id: "worker".into(), - role: TeamRunTargetRole::Teammate, - unread_count: 3, - }]) - .await; - - assert!(result.is_none(), "pending slot work already owns the unread backlog"); - - let user_wake = manager - .claim_wake_for_turn("worker", TeamRunTargetRole::Teammate, "conv-worker") - .await - .expect("original pending user wake should remain"); - assert_eq!(user_wake.wake_source, TeamWakeSource::UserMessage); - assert_eq!(user_wake.message_id.as_deref(), Some("mailbox-worker")); - - assert!( - manager - .claim_wake_for_turn("worker", TeamRunTargetRole::Teammate, "conv-worker-dup") - .await - .is_none(), - "recovery scan must not append a duplicate wake for represented work" - ); - } - - #[tokio::test] - async fn recovery_backlog_does_not_duplicate_paused_gate_work() { - let (manager, _bc) = manager(); - let (_ack, lease) = manager - .acquire_user_message_wake("lead", TeamRunTargetRole::Lead) - .await - .expect("user run"); - manager - .commit_operation_lease(&lease.lease_id, Some("mailbox-lead".into())) - .await - .expect("commit pending user wake"); - - let reservation = manager - .claim_wake_for_turn("lead", TeamRunTargetRole::Lead, "conv-lead") - .await - .expect("claim lead wake"); - let child = ActiveChildTurn { - team_run_id: reservation.team_run_id.clone(), - slot_id: "lead".into(), - role: TeamRunTargetRole::Lead, - conversation_id: "conv-lead".into(), - turn_id: "turn-lead".into(), - started_at_ms: now_ms(), - last_slow_notified_at_ms: None, - }; - assert_eq!( - manager - .record_child_started(&reservation.reservation_id, child.clone()) - .await, - ChildStartDecision::Accepted - ); - manager - .complete_pause_after_child_cancelled(&child, Some("test_pause".into())) - .await - .expect("pause slot"); - - let result = manager - .recover_mailbox_backlog(vec![RecoveryWakeCandidate { - slot_id: "lead".into(), - role: TeamRunTargetRole::Lead, - unread_count: 3, - }]) - .await; - - assert!(result.is_none(), "paused wake gate already represents retained work"); - assert!( - manager - .claim_wake_for_turn("lead", TeamRunTargetRole::Lead, "conv-lead-dup") - .await - .is_none(), - "recovery scan must not append a duplicate wake for paused gate work" - ); - } - - #[tokio::test] - async fn empty_recovery_backlog_does_not_create_run() { - let (manager, _bc) = manager(); - let result = manager.recover_mailbox_backlog(vec![]).await; - assert!(result.is_none()); - assert!(manager.active_run_id().await.is_none()); - } - - #[tokio::test] - async fn lease_keeps_run_active_until_commit() { - let (manager, _bc) = manager(); - let (ack, lease) = manager - .acquire_user_message_wake("lead", TeamRunTargetRole::Lead) - .await - .expect("user message should create run and lease"); - - assert_eq!(lease.team_run_id, ack.team_run_id); - assert_eq!(lease.slot_id, "lead"); - assert_eq!(lease.wake_source, TeamWakeSource::UserMessage); - assert!(lease.accepted_as_new_run); - - let completed = manager.maybe_complete().await; - assert!(completed.is_none(), "active lease must retain the run"); - assert_eq!(manager.active_run_id().await.as_deref(), Some(ack.team_run_id.as_str())); - - manager - .commit_operation_lease(&lease.lease_id, Some("mailbox-1".into())) - .await - .expect("commit should convert lease to pending wake"); - - let reservation = manager - .claim_wake_for_turn("lead", TeamRunTargetRole::Lead, "conv-lead") - .await - .expect("committed wake should be claimable"); - assert_eq!(reservation.message_id.as_deref(), Some("mailbox-1")); - } - - #[tokio::test] - async fn abort_lease_releases_completion_hold() { - let (manager, _bc) = manager(); - let (ack, lease) = manager - .acquire_user_message_wake("lead", TeamRunTargetRole::Lead) - .await - .expect("user message should create run and lease"); - - manager - .abort_operation_lease(&lease.lease_id, "mailbox_write_failed") - .await - .expect("abort should remove lease"); - - let completed = manager - .maybe_complete() - .await - .expect("run should complete after aborted only lease"); - assert_eq!(completed.team_run_id, ack.team_run_id); - assert_eq!(completed.status, TeamRunStatus::Completed); - } - - #[tokio::test] - async fn commit_missing_lease_returns_internal_consistency_error() { - let (manager, _bc) = manager(); - manager - .acquire_user_message_wake("lead", TeamRunTargetRole::Lead) - .await - .expect("create active run"); - - let err = manager - .commit_operation_lease("missing-lease", Some("mailbox-1".into())) - .await - .expect_err("missing lease is a contract violation"); - - assert!(matches!( - err, - TeamError::InvalidRequest(message) - if message == "team run operation lease missing: missing-lease" - )); - } - - #[tokio::test] - async fn run_scoped_wake_without_active_run_is_rejected() { - let (manager, _bc) = manager(); - - let err = manager - .acquire_run_scoped_wake("worker", TeamRunTargetRole::Teammate, TeamWakeSource::McpSendMessage) - .await - .expect_err("run-scoped wake must need active run"); - - assert!(matches!( - err, - TeamError::InvalidRequest(message) - if message == "no active team run for run-scoped wake" - )); - } - - #[tokio::test] - async fn user_message_busy_active_slot_is_rejected_without_lease() { - let (manager, _bc) = manager(); - let (ack, lease) = manager - .acquire_user_message_wake("lead", TeamRunTargetRole::Lead) - .await - .unwrap(); - manager - .commit_operation_lease(&lease.lease_id, Some("mailbox-1".into())) - .await - .unwrap(); - let reservation = manager - .claim_wake_for_turn("lead", TeamRunTargetRole::Lead, "conv-lead") - .await - .unwrap(); - manager - .record_child_started( - &reservation.reservation_id, - ActiveChildTurn { - team_run_id: ack.team_run_id, - slot_id: "lead".into(), - role: TeamRunTargetRole::Lead, - conversation_id: "conv-lead".into(), - turn_id: "turn-lead".into(), - started_at_ms: now_ms(), - last_slow_notified_at_ms: None, - }, - ) - .await; - - let err = manager - .acquire_user_message_wake("lead", TeamRunTargetRole::Lead) - .await - .expect_err("foreground user message must reject busy active slot"); - - assert!(matches!(err, TeamError::SlotBusy(slot) if slot == "lead")); - let payload = manager.current_payload().await.unwrap(); - assert_eq!(payload.pending_wake_count, 0); - } - - #[tokio::test] - async fn user_message_pending_slot_is_accepted_as_additional_foreground_wake() { - let (manager, _bc) = manager(); - let (_ack, first) = manager - .acquire_user_message_wake("lead", TeamRunTargetRole::Lead) - .await - .unwrap(); - manager - .commit_operation_lease(&first.lease_id, Some("mailbox-1".into())) - .await - .unwrap(); - - let (_ack, second) = manager - .acquire_user_message_wake("lead", TeamRunTargetRole::Lead) - .await - .expect("pending foreground message must be accepted"); - manager - .commit_operation_lease(&second.lease_id, Some("mailbox-2".into())) - .await - .unwrap(); - - let payload = manager.current_payload().await.unwrap(); - assert_eq!(payload.pending_wake_count, 2); - } - - #[tokio::test] - async fn mcp_send_message_busy_slot_is_accepted_and_queued() { - let (manager, _bc) = manager(); - let (ack, first) = manager - .acquire_user_message_wake("lead", TeamRunTargetRole::Lead) - .await - .unwrap(); - manager - .commit_operation_lease(&first.lease_id, Some("mailbox-1".into())) - .await - .unwrap(); - let reservation = manager - .claim_wake_for_turn("lead", TeamRunTargetRole::Lead, "conv-lead") - .await - .unwrap(); - manager - .record_child_started( - &reservation.reservation_id, - ActiveChildTurn { - team_run_id: ack.team_run_id, - slot_id: "lead".into(), - role: TeamRunTargetRole::Lead, - conversation_id: "conv-lead".into(), - turn_id: "turn-lead".into(), - started_at_ms: now_ms(), - last_slow_notified_at_ms: None, - }, - ) - .await; - - let outcome = manager - .acquire_run_scoped_wake("lead", TeamRunTargetRole::Lead, TeamWakeSource::McpSendMessage) - .await - .expect("MCP message to busy slot must be accepted"); - let TeamRunWakeAcquireOutcome::Accepted(lease) = outcome else { - panic!("MCP message should not be suppressed"); - }; - manager - .commit_operation_lease(&lease.lease_id, Some("mailbox-2".into())) - .await - .unwrap(); - - let payload = manager.current_payload().await.unwrap(); - assert_eq!(payload.pending_wake_count, 1); - } - - #[tokio::test] - async fn paused_slot_suppresses_mcp_send_message_without_lease() { - let (manager, _bc) = manager(); - let (ack, _lease) = manager - .acquire_user_message_wake("lead", TeamRunTargetRole::Lead) - .await - .unwrap(); - manager - .pause_slot_work("lead", Some("user stopped".into())) - .await - .unwrap(); - - let outcome = manager - .acquire_run_scoped_wake("lead", TeamRunTargetRole::Lead, TeamWakeSource::McpSendMessage) - .await - .expect("paused background wake suppresses cleanly"); - assert_eq!(outcome, TeamRunWakeAcquireOutcome::Suppressed); - assert_eq!(manager.active_run_id().await.as_deref(), Some(ack.team_run_id.as_str())); - } - - #[tokio::test] - async fn duplicate_spawn_welcome_is_suppressed() { - let (manager, _bc) = manager(); - manager - .acquire_user_message_wake("lead", TeamRunTargetRole::Lead) - .await - .unwrap(); - - let first = manager - .acquire_run_scoped_wake("new-worker", TeamRunTargetRole::Teammate, TeamWakeSource::SpawnWelcome) - .await - .unwrap(); - let TeamRunWakeAcquireOutcome::Accepted(first) = first else { - panic!("first spawn welcome should be accepted"); - }; - manager - .commit_operation_lease(&first.lease_id, Some("welcome-1".into())) - .await - .unwrap(); - - let second = manager - .acquire_run_scoped_wake("new-worker", TeamRunTargetRole::Teammate, TeamWakeSource::SpawnWelcome) - .await - .unwrap(); - assert_eq!(second, TeamRunWakeAcquireOutcome::Suppressed); - } - - #[tokio::test] - async fn idle_notification_idle_slot_is_suppressed_without_scheduler_wake_target() { - let (manager, _bc) = manager(); - manager - .acquire_user_message_wake("lead", TeamRunTargetRole::Lead) - .await - .unwrap(); - - let outcome = manager - .acquire_run_scoped_wake("worker", TeamRunTargetRole::Teammate, TeamWakeSource::IdleNotification) - .await - .expect("generic idle notification should suppress cleanly"); - - assert_eq!(outcome, TeamRunWakeAcquireOutcome::Suppressed); - } - - #[tokio::test] - async fn scheduler_idle_notification_idle_slot_is_accepted() { - let (manager, _bc) = manager(); - manager - .acquire_user_message_wake("lead", TeamRunTargetRole::Lead) - .await - .unwrap(); - - let outcome = manager - .acquire_scheduler_wake("worker", TeamRunTargetRole::Teammate, TeamWakeSource::IdleNotification) - .await - .expect("scheduler-produced wake target should be accepted"); - let TeamRunWakeAcquireOutcome::Accepted(lease) = outcome else { - panic!("scheduler wake target should produce a lease"); - }; - - manager - .commit_operation_lease(&lease.lease_id, None) - .await - .expect("scheduler wake should commit"); - let reservation = manager - .claim_wake_for_turn("worker", TeamRunTargetRole::Teammate, "conv-worker") - .await - .expect("scheduler wake should be claimable"); - assert_eq!(reservation.wake_source, TeamWakeSource::IdleNotification); - } - - #[tokio::test] - async fn active_slot_work_includes_backend_slow_fields_after_threshold() { - let (manager, _rx) = manager(); - let ack = manager - .accept_user_message("lead", TeamRunTargetRole::Lead, true, Some("msg-1".into())) - .await - .unwrap(); - manager - .record_pending_wake("lead", TeamRunTargetRole::Lead, TeamWakeSource::UserMessage) - .await - .unwrap(); - let reservation = manager - .claim_wake_for_turn("lead", TeamRunTargetRole::Lead, "conv-lead") - .await - .unwrap(); - - let child = ActiveChildTurn { - team_run_id: ack.team_run_id.clone(), - slot_id: "lead".into(), - role: TeamRunTargetRole::Lead, - conversation_id: "conv-lead".into(), - turn_id: "turn-lead".into(), - started_at_ms: 10_000, - last_slow_notified_at_ms: None, - }; - manager.record_child_started(&reservation.reservation_id, child).await; - - manager.observe_slow_child_turns(610_001).await; - let payload = manager.current_payload().await.unwrap(); - let lead = slot_work(&payload, "lead"); - - assert_eq!(lead.active_turn_started_at_ms, Some(10_000)); - assert_eq!(lead.active_turn_elapsed_ms, Some(600_001)); - assert_eq!(lead.active_turn_slow, Some(true)); - assert_eq!(lead.active_turn_slow_threshold_ms, Some(600_000)); - assert_eq!(lead.runtime_health, None); - } - - #[tokio::test] - async fn slow_observation_is_rate_limited_per_active_child() { - let (manager, _rx) = manager(); - let ack = manager - .accept_user_message("lead", TeamRunTargetRole::Lead, true, Some("msg-1".into())) - .await - .unwrap(); - manager - .record_pending_wake("lead", TeamRunTargetRole::Lead, TeamWakeSource::UserMessage) - .await - .unwrap(); - let reservation = manager - .claim_wake_for_turn("lead", TeamRunTargetRole::Lead, "conv-lead") - .await - .unwrap(); - - let child = ActiveChildTurn { - team_run_id: ack.team_run_id.clone(), - slot_id: "lead".into(), - role: TeamRunTargetRole::Lead, - conversation_id: "conv-lead".into(), - turn_id: "turn-lead".into(), - started_at_ms: 0, - last_slow_notified_at_ms: None, - }; - manager.record_child_started(&reservation.reservation_id, child).await; - - assert!(manager.observe_slow_child_turns(600_001).await.is_some()); - assert!(manager.observe_slow_child_turns(900_000).await.is_none()); - assert!(manager.observe_slow_child_turns(1_200_001).await.is_some()); - } - - #[tokio::test] - async fn cancel_run_clears_paused_gate_and_reaches_cancelled_terminal() { - let (manager, bc) = manager(); - manager - .accept_user_message("lead", TeamRunTargetRole::Lead, false, None) - .await - .unwrap(); - manager - .pause_slot_work("lead", Some("user stopped".into())) - .await - .unwrap(); - manager - .record_or_suppress_wake( - "lead", - TeamRunTargetRole::Lead, - TeamWakeSource::InterruptedNotification, - None, - ) - .await - .unwrap(); - - manager.begin_cancel(None, Some("stop all".into())).await.unwrap(); - let cancelled = manager - .try_complete_cancelled() - .await - .expect("cancel should clear retained gate work"); - - assert_eq!(cancelled.status, TeamRunStatus::Cancelled); - assert_eq!(cancelled.pending_wake_count, 0); - assert_eq!(cancelled.slot_work.len(), 0); - assert_eq!(manager.active_run_id().await, None); - assert!(bc.names().contains(&TEAM_RUN_CANCELLED_EVENT.to_owned())); - } - - #[tokio::test] - async fn run_payload_reports_pending_starting_and_active_work_by_slot() { - let (manager, bc) = manager(); - let ack = manager - .accept_user_message("lead", TeamRunTargetRole::Lead, false, None) - .await - .unwrap(); - - manager - .record_pending_wake("lead", TeamRunTargetRole::Lead, TeamWakeSource::UserMessage) - .await - .unwrap(); - manager - .record_pending_wake("worker", TeamRunTargetRole::Teammate, TeamWakeSource::McpSendMessage) - .await - .unwrap(); - let worker_reservation = manager - .claim_wake_for_turn("worker", TeamRunTargetRole::Teammate, "conv-worker") - .await - .unwrap(); - let lead_reservation = manager - .claim_wake_for_turn("lead", TeamRunTargetRole::Lead, "conv-lead") - .await - .unwrap(); - manager - .record_child_started( - &lead_reservation.reservation_id, - ActiveChildTurn { - team_run_id: ack.team_run_id.clone(), - slot_id: "lead".into(), - role: TeamRunTargetRole::Lead, - conversation_id: "conv-lead".into(), - turn_id: "turn-lead".into(), - started_at_ms: now_ms(), - last_slow_notified_at_ms: None, - }, - ) - .await; - - let latest = bc.run_payloads().last().cloned().expect("run payload"); - let lead = slot_work(&latest, "lead"); - assert_eq!(lead.pending_wake_count, 0); - assert_eq!(lead.starting_child_count, 0); - assert_eq!(lead.active_turn_id.as_deref(), Some("turn-lead")); - - let worker = slot_work(&latest, "worker"); - assert_eq!(worker.pending_wake_count, 0); - assert_eq!(worker.starting_child_count, 1); - assert_eq!(worker.active_turn_id, None); - assert_eq!(worker_reservation.slot_id, "worker"); - } - - #[tokio::test] - async fn pause_active_child_prepares_cancel_without_mutating_gate_or_active_child() { - let (manager, bc) = manager(); - let ack = manager - .accept_user_message("lead", TeamRunTargetRole::Lead, false, None) - .await - .unwrap(); - manager - .record_pending_wake("lead", TeamRunTargetRole::Lead, TeamWakeSource::UserMessage) - .await - .unwrap(); - let reservation = manager - .claim_wake_for_turn("lead", TeamRunTargetRole::Lead, "conv-lead") - .await - .unwrap(); - manager - .record_child_started( - &reservation.reservation_id, - ActiveChildTurn { - team_run_id: ack.team_run_id, - slot_id: "lead".into(), - role: TeamRunTargetRole::Lead, - conversation_id: "conv-lead".into(), - turn_id: "turn-lead".into(), - started_at_ms: now_ms(), - last_slow_notified_at_ms: None, - }, - ) - .await; - manager - .record_pending_wake("lead", TeamRunTargetRole::Lead, TeamWakeSource::McpSendMessage) - .await - .unwrap(); - - let outcome = manager - .pause_slot_work("lead", Some("user stopped".into())) - .await - .unwrap(); - - assert!(matches!(outcome.cancel_target, Some(ChildCancelTarget::Active(_)))); - let latest = bc.run_payloads().last().cloned().expect("run payload"); - let lead = slot_work(&latest, "lead"); - assert!(!lead.paused); - assert_eq!(lead.pending_wake_count, 1); - assert_eq!(lead.suppressed_wake_count, 0); - assert_eq!(lead.active_turn_id.as_deref(), Some("turn-lead")); - } - - #[tokio::test] - async fn complete_pause_after_active_child_cancel_marks_paused_moves_pending_and_removes_child() { - let (manager, bc) = manager(); - let ack = manager - .accept_user_message("lead", TeamRunTargetRole::Lead, false, None) - .await - .unwrap(); - manager - .record_pending_wake("lead", TeamRunTargetRole::Lead, TeamWakeSource::UserMessage) - .await - .unwrap(); - let reservation = manager - .claim_wake_for_turn("lead", TeamRunTargetRole::Lead, "conv-lead") - .await - .unwrap(); - let child = ActiveChildTurn { - team_run_id: ack.team_run_id, - slot_id: "lead".into(), - role: TeamRunTargetRole::Lead, - conversation_id: "conv-lead".into(), - turn_id: "turn-lead".into(), - started_at_ms: now_ms(), - last_slow_notified_at_ms: None, - }; - manager - .record_child_started(&reservation.reservation_id, child.clone()) - .await; - manager - .record_pending_wake("lead", TeamRunTargetRole::Lead, TeamWakeSource::McpSendMessage) - .await - .unwrap(); - manager - .pause_slot_work("lead", Some("user stopped".into())) - .await - .unwrap(); - - let payload = manager - .complete_pause_after_child_cancelled(&child, Some("user stopped".into())) - .await - .expect("pause completion payload"); - - let lead = slot_work(&payload, "lead"); - assert!(lead.paused); - assert_eq!(lead.pending_wake_count, 0); - assert_eq!(lead.suppressed_wake_count, 1); - assert_eq!(lead.active_turn_id, None); - assert!(bc.names().contains(&TEAM_CHILD_TURN_CANCELLED_EVENT.to_owned())); - } - - #[tokio::test] - async fn pause_slot_moves_pending_wakes_to_suppressed_slot_work() { - let (manager, bc) = manager(); - let ack = manager - .accept_user_message("lead", TeamRunTargetRole::Lead, false, None) - .await - .unwrap(); - manager - .record_or_suppress_wake( - "lead", - TeamRunTargetRole::Lead, - TeamWakeSource::UserMessage, - Some("msg-1".into()), - ) - .await - .unwrap(); - - let outcome = manager - .pause_slot_work("lead", Some("user stopped".into())) - .await - .unwrap(); - - assert_eq!(outcome.team_run_id, ack.team_run_id); - let latest = bc.run_payloads().pop().expect("run update"); - let lead = slot_work(&latest, "lead"); - assert!(lead.paused); - assert_eq!(lead.pending_wake_count, 0); - assert_eq!(lead.suppressed_wake_count, 1); - } - - #[tokio::test] - async fn paused_slot_suppresses_background_wake_without_pending_count() { - let (manager, bc) = manager(); - manager - .accept_user_message("lead", TeamRunTargetRole::Lead, false, None) - .await - .unwrap(); - manager - .pause_slot_work("lead", Some("user stopped".into())) - .await - .unwrap(); - - let decision = manager - .record_or_suppress_wake( - "lead", - TeamRunTargetRole::Lead, - TeamWakeSource::InterruptedNotification, - None, - ) - .await - .unwrap(); - - assert_eq!(decision, WakeRecordDecision::Suppressed); - let latest = bc.run_payloads().pop().expect("run update"); - let lead = slot_work(&latest, "lead"); - assert!(lead.paused); - assert_eq!(lead.pending_wake_count, 0); - assert_eq!(lead.suppressed_wake_count, 1); - } - - #[tokio::test] - async fn user_intervention_resumes_paused_slot_and_records_pending_wake() { - let (manager, bc) = manager(); - manager - .accept_user_message("lead", TeamRunTargetRole::Lead, false, None) - .await - .unwrap(); - manager - .pause_slot_work("lead", Some("user stopped".into())) - .await - .unwrap(); - - let decision = manager - .record_or_suppress_wake( - "lead", - TeamRunTargetRole::Lead, - TeamWakeSource::UserIntervention, - Some("msg-user".into()), - ) - .await - .unwrap(); - - assert_eq!(decision, WakeRecordDecision::Recorded); - let latest = bc.run_payloads().pop().expect("run update"); - let lead = slot_work(&latest, "lead"); - assert!(!lead.paused); - assert_eq!(lead.pending_wake_count, 1); - assert_eq!(lead.suppressed_wake_count, 0); - } - - #[tokio::test] - async fn resumed_user_message_releases_suppressed_wake_and_allows_completion() { - let (manager, bc) = manager(); - let ack = manager - .accept_user_message("lead", TeamRunTargetRole::Lead, false, None) - .await - .unwrap(); - - manager - .pause_slot_work("lead", Some("user stopped".into())) - .await - .unwrap(); - manager - .record_or_suppress_wake( - "lead", - TeamRunTargetRole::Lead, - TeamWakeSource::IdleNotification, - Some("idle-msg".into()), - ) - .await - .unwrap(); - - manager - .record_or_suppress_wake( - "lead", - TeamRunTargetRole::Lead, - TeamWakeSource::UserMessage, - Some("user-msg".into()), - ) - .await - .unwrap(); - - let user_reservation = manager - .claim_wake_for_turn("lead", TeamRunTargetRole::Lead, "conv-lead") - .await - .expect("foreground user message wake should run first"); - assert_eq!( - manager - .record_child_started( - &user_reservation.reservation_id, - ActiveChildTurn { - team_run_id: ack.team_run_id.clone(), - slot_id: "lead".into(), - role: TeamRunTargetRole::Lead, - conversation_id: "conv-lead".into(), - turn_id: "turn-user".into(), - started_at_ms: now_ms(), - last_slow_notified_at_ms: None, - }, - ) - .await, - ChildStartDecision::Accepted - ); - assert!( - manager - .record_child_completed("lead", "turn-user", TeamRunStatus::Completed) - .await - .is_none(), - "suppressed background wake should retain the run until released" - ); - - let source = manager - .release_suppressed_wake_if_resumed("lead", TeamRunTargetRole::Lead) - .await; - assert_eq!(source, Some(TeamWakeSource::IdleNotification)); - - let reservation = manager - .claim_wake_for_turn("lead", TeamRunTargetRole::Lead, "conv-lead") - .await - .expect("released suppressed wake should become pending wake"); - assert_eq!( - manager - .record_child_started( - &reservation.reservation_id, - ActiveChildTurn { - team_run_id: ack.team_run_id.clone(), - slot_id: "lead".into(), - role: TeamRunTargetRole::Lead, - conversation_id: "conv-lead".into(), - turn_id: "turn-background".into(), - started_at_ms: now_ms(), - last_slow_notified_at_ms: None, - }, - ) - .await, - ChildStartDecision::Accepted - ); - let completed = manager - .record_child_completed("lead", "turn-background", TeamRunStatus::Completed) - .await - .expect("run should complete after released background wake is consumed"); - - assert_eq!(completed.status, TeamRunStatus::Completed); - assert_eq!(completed.pending_wake_count, 0); - assert_eq!(completed.starting_child_count, 0); - assert_eq!(completed.active_child_count, 0); - assert_eq!( - completed.slot_work.len(), - 0, - "gate state should not leave empty slot work" - ); - assert!(bc.names().contains(&TEAM_RUN_COMPLETED_EVENT.to_owned())); - } - - #[tokio::test] - async fn active_intervention_ack_uses_accepted_slot_without_changing_initial_target() { - let (manager, _) = manager(); - let first = manager - .accept_user_message("lead", TeamRunTargetRole::Lead, false, None) - .await - .unwrap(); - - let second = manager - .accept_user_message("worker", TeamRunTargetRole::Teammate, true, Some("msg-2".into())) - .await - .unwrap(); - - assert_eq!(second.team_run_id, first.team_run_id); - assert_eq!(second.target_slot_id, "lead"); - assert_eq!(second.target_role, TeamRunTargetRole::Lead); - assert_eq!(second.accepted_slot_id, "worker"); - assert_eq!(second.accepted_role, TeamRunTargetRole::Teammate); - assert_eq!(second.message_id.as_deref(), Some("msg-2")); - } - - #[tokio::test] - async fn active_intervention_rejects_when_same_slot_has_pending_work() { - let (manager, _) = manager(); - manager - .accept_user_message("lead", TeamRunTargetRole::Lead, false, None) - .await - .unwrap(); - manager - .record_pending_wake("lead", TeamRunTargetRole::Lead, TeamWakeSource::UserMessage) - .await - .unwrap(); - - let err = manager - .accept_user_message("lead", TeamRunTargetRole::Lead, true, None) - .await - .expect_err("same slot pending work must be busy"); - - assert!(matches!(err, TeamError::SlotBusy(slot_id) if slot_id == "lead")); - } - - #[tokio::test] - async fn active_intervention_rejects_when_same_slot_has_starting_or_active_work() { - let (manager, _) = manager(); - let ack = manager - .accept_user_message("lead", TeamRunTargetRole::Lead, false, None) - .await - .unwrap(); - manager - .record_pending_wake("lead", TeamRunTargetRole::Lead, TeamWakeSource::UserMessage) - .await - .unwrap(); - let reservation = manager - .claim_wake_for_turn("lead", TeamRunTargetRole::Lead, "conv-lead") - .await - .unwrap(); - - let starting_err = manager - .accept_user_message("lead", TeamRunTargetRole::Lead, true, None) - .await - .expect_err("same slot starting work must be busy"); - assert!(matches!(starting_err, TeamError::SlotBusy(slot_id) if slot_id == "lead")); - - manager - .record_child_started( - &reservation.reservation_id, - ActiveChildTurn { - team_run_id: ack.team_run_id, - slot_id: "lead".into(), - role: TeamRunTargetRole::Lead, - conversation_id: "conv-lead".into(), - turn_id: "turn-lead".into(), - started_at_ms: now_ms(), - last_slow_notified_at_ms: None, - }, - ) - .await; - - let active_err = manager - .accept_user_message("lead", TeamRunTargetRole::Lead, true, None) - .await - .expect_err("same slot active work must be busy"); - assert!(matches!(active_err, TeamError::SlotBusy(slot_id) if slot_id == "lead")); - } - - #[tokio::test] - async fn record_pending_wake_requires_active_run() { - let (manager, _) = manager(); - - let err = manager - .record_pending_wake("worker-1", TeamRunTargetRole::Teammate, TeamWakeSource::McpSendMessage) - .await - .expect_err("run-scoped wake without active run must fail"); - - assert!(matches!( - err, - TeamError::InvalidRequest(message) - if message == "no active team run for run-scoped wake" - )); - } - - #[tokio::test] - async fn record_pending_wake_increments_active_run_count() { - let (manager, _) = manager(); - manager - .accept_user_message("lead-1", TeamRunTargetRole::Lead, false, None) - .await - .expect("accept run"); - - manager - .record_pending_wake("worker-1", TeamRunTargetRole::Teammate, TeamWakeSource::McpSendMessage) - .await - .expect("record pending wake"); - - let reservation = manager - .claim_wake_for_turn("worker-1", TeamRunTargetRole::Teammate, "conv-worker") - .await - .expect("pending wake should be claimable"); - - assert_eq!(reservation.slot_id, "worker-1"); - assert_eq!(reservation.role, TeamRunTargetRole::Teammate); - } - - #[tokio::test] - async fn empty_wake_for_one_slot_does_not_consume_another_slot_pending_wake() { - let (manager, _) = manager(); - let ack = manager - .accept_user_message("lead", TeamRunTargetRole::Lead, false, None) - .await - .expect("accept run"); - - manager - .record_pending_wake("worker", TeamRunTargetRole::Teammate, TeamWakeSource::McpSendMessage) - .await - .expect("record worker pending wake"); - - let completed = manager.record_empty_wake_observed("lead").await; - assert!( - completed.is_none(), - "leader empty wake must not complete while worker has pending wake" - ); - assert_eq!(manager.active_run_id().await.as_deref(), Some(ack.team_run_id.as_str())); - - let reservation = manager - .claim_wake_for_turn("worker", TeamRunTargetRole::Teammate, "conv-worker") - .await - .expect("worker pending wake must remain claimable"); - - assert_eq!(reservation.slot_id, "worker"); - assert_eq!(reservation.role, TeamRunTargetRole::Teammate); - } - - #[tokio::test] - async fn empty_wake_consumes_only_same_slot_pending_wake() { - let (manager, _) = manager(); - manager - .accept_user_message("lead", TeamRunTargetRole::Lead, false, None) - .await - .expect("accept run"); - - manager - .record_pending_wake("lead", TeamRunTargetRole::Lead, TeamWakeSource::UserMessage) - .await - .expect("record lead pending wake"); - manager - .record_pending_wake("worker", TeamRunTargetRole::Teammate, TeamWakeSource::McpSendMessage) - .await - .expect("record worker pending wake"); - - assert!( - manager.record_empty_wake_observed("lead").await.is_none(), - "worker wake should keep run active" - ); - - assert!( - manager - .claim_wake_for_turn("lead", TeamRunTargetRole::Lead, "conv-lead") - .await - .is_none(), - "lead pending wake was consumed by lead empty wake" - ); - - assert!( - manager - .claim_wake_for_turn("worker", TeamRunTargetRole::Teammate, "conv-worker") - .await - .is_some(), - "worker pending wake must remain" - ); - } - - #[tokio::test] - async fn leader_message_rejects_when_run_is_active() { - let (manager, _) = manager(); - manager - .accept_user_message("lead", TeamRunTargetRole::Lead, false, None) - .await - .unwrap(); - - let err = manager - .accept_user_message("lead", TeamRunTargetRole::Lead, false, None) - .await - .unwrap_err(); - - assert!(matches!(err, TeamError::InvalidRequest(message) if message.contains("already active"))); - } - - #[tokio::test] - async fn teammate_intervention_reuses_active_run() { - let (manager, _) = manager(); - let first = manager - .accept_user_message("lead", TeamRunTargetRole::Lead, false, None) - .await - .unwrap(); - - let second = manager - .accept_user_message("worker", TeamRunTargetRole::Teammate, true, Some("msg-1".into())) - .await - .unwrap(); - - assert_eq!(second.team_run_id, first.team_run_id); - assert_eq!(second.accepted_slot_id, "worker"); - assert_eq!(second.accepted_role, TeamRunTargetRole::Teammate); - assert_eq!(second.message_id.as_deref(), Some("msg-1")); - } - - #[tokio::test] - async fn child_start_and_completion_emit_lifecycle_events() { - let (manager, bc) = manager(); - let ack = manager - .accept_user_message("lead", TeamRunTargetRole::Lead, false, None) - .await - .unwrap(); - manager - .record_pending_wake("lead", TeamRunTargetRole::Lead, TeamWakeSource::UserMessage) - .await - .unwrap(); - let reservation = manager - .claim_wake_for_turn("lead", TeamRunTargetRole::Lead, "conv") - .await - .unwrap(); - - manager - .record_child_started( - &reservation.reservation_id, - ActiveChildTurn { - team_run_id: ack.team_run_id.clone(), - slot_id: "lead".into(), - role: TeamRunTargetRole::Lead, - conversation_id: "conv".into(), - turn_id: "turn".into(), - started_at_ms: now_ms(), - last_slow_notified_at_ms: None, - }, - ) - .await; - manager - .record_child_completed("lead", "turn", TeamRunStatus::Completed) - .await; - - assert_eq!(manager.active_run_id().await, None); - let names = bc.names(); - assert!(names.contains(&TEAM_RUN_ACCEPTED_EVENT.to_owned())); - assert!(names.contains(&TEAM_CHILD_TURN_STARTED_EVENT.to_owned())); - assert!(names.contains(&TEAM_RUN_COMPLETED_EVENT.to_owned())); - } - - #[tokio::test] - async fn claimed_wake_prevents_completion_until_child_finishes() { - let (manager, bc) = manager(); - let ack = manager - .accept_user_message("lead", TeamRunTargetRole::Lead, false, None) - .await - .unwrap(); - manager - .record_pending_wake("lead", TeamRunTargetRole::Lead, TeamWakeSource::UserMessage) - .await - .unwrap(); - - let reservation = manager - .claim_wake_for_turn("lead", TeamRunTargetRole::Lead, "conv") - .await - .expect("reservation should be claimed"); - - assert_eq!(reservation.team_run_id, ack.team_run_id); - assert_eq!(manager.maybe_complete().await, None); - assert_eq!(manager.active_run_id().await.as_deref(), Some(ack.team_run_id.as_str())); - - let payloads = bc.run_payloads(); - assert!(payloads.iter().any(|payload| payload.starting_child_count == 1)); - } - - #[tokio::test] - async fn child_start_promotes_matching_reservation_to_active() { - let (manager, bc) = manager(); - let ack = manager - .accept_user_message("lead", TeamRunTargetRole::Lead, false, None) - .await - .unwrap(); - manager - .record_pending_wake("lead", TeamRunTargetRole::Lead, TeamWakeSource::UserMessage) - .await - .unwrap(); - let reservation = manager - .claim_wake_for_turn("lead", TeamRunTargetRole::Lead, "conv") - .await - .unwrap(); - - let decision = manager - .record_child_started( - &reservation.reservation_id, - ActiveChildTurn { - team_run_id: ack.team_run_id.clone(), - slot_id: "lead".into(), - role: TeamRunTargetRole::Lead, - conversation_id: "conv".into(), - turn_id: "turn".into(), - started_at_ms: now_ms(), - last_slow_notified_at_ms: None, - }, - ) - .await; - - assert_eq!(decision, ChildStartDecision::Accepted); - let payloads = bc.run_payloads(); - assert!(payloads.iter().any(|payload| { - payload.status == TeamRunStatus::Running - && payload.starting_child_count == 0 - && payload.active_child_count == 1 - })); - } - - #[tokio::test] - async fn child_completion_completes_only_when_pending_starting_and_active_are_empty() { - let (manager, bc) = manager(); - let ack = manager - .accept_user_message("lead", TeamRunTargetRole::Lead, false, None) - .await - .unwrap(); - manager - .record_pending_wake("lead", TeamRunTargetRole::Lead, TeamWakeSource::UserMessage) - .await - .unwrap(); - let reservation = manager - .claim_wake_for_turn("lead", TeamRunTargetRole::Lead, "conv") - .await - .unwrap(); - assert_eq!( - manager - .record_child_started( - &reservation.reservation_id, - ActiveChildTurn { - team_run_id: ack.team_run_id, - slot_id: "lead".into(), - role: TeamRunTargetRole::Lead, - conversation_id: "conv".into(), - turn_id: "turn".into(), - started_at_ms: now_ms(), - last_slow_notified_at_ms: None, - }, - ) - .await, - ChildStartDecision::Accepted - ); - - let completed = manager - .record_child_completed("lead", "turn", TeamRunStatus::Completed) - .await - .expect("run should complete after last child"); - - assert_eq!(completed.status, TeamRunStatus::Completed); - assert_eq!(completed.pending_wake_count, 0); - assert_eq!(completed.starting_child_count, 0); - assert_eq!(completed.active_child_count, 0); - assert_eq!(manager.active_run_id().await, None); - assert!(bc.names().contains(&TEAM_RUN_COMPLETED_EVENT.to_owned())); - } - - #[tokio::test] - async fn completed_child_after_internal_conversation_replay_does_not_emit_failed_run() { - let (manager, bc) = manager(); - let ack = manager - .accept_user_message("lead", TeamRunTargetRole::Lead, false, Some("msg-user".into())) - .await - .unwrap(); - manager - .record_pending_wake("lead", TeamRunTargetRole::Lead, TeamWakeSource::UserMessage) - .await - .unwrap(); - let reservation = manager - .claim_wake_for_turn("lead", TeamRunTargetRole::Lead, "conv-lead") - .await - .unwrap(); - assert_eq!( - manager - .record_child_started( - &reservation.reservation_id, - ActiveChildTurn { - team_run_id: ack.team_run_id, - slot_id: "lead".into(), - role: TeamRunTargetRole::Lead, - conversation_id: "conv-lead".into(), - turn_id: "turn-replayed-in-conversation".into(), - started_at_ms: now_ms(), - last_slow_notified_at_ms: None, - }, - ) - .await, - ChildStartDecision::Accepted - ); - - let completed = manager - .record_child_completed("lead", "turn-replayed-in-conversation", TeamRunStatus::Completed) - .await - .expect("run should complete after conversation reports final completed outcome"); - - assert_eq!(completed.status, TeamRunStatus::Completed); - let names = bc.names(); - assert!(names.contains(&TEAM_RUN_COMPLETED_EVENT.to_owned())); - assert!( - !names.contains(&TEAM_RUN_FAILED_EVENT.to_owned()), - "internal conversation replay must not surface as an intermediate team_run failed event" - ); - } - - #[tokio::test] - async fn child_start_failed_releases_reservation_and_fails_run() { - let (manager, bc) = manager(); - manager - .accept_user_message("lead", TeamRunTargetRole::Lead, false, None) - .await - .unwrap(); - manager - .record_pending_wake("lead", TeamRunTargetRole::Lead, TeamWakeSource::UserMessage) - .await - .unwrap(); - let reservation = manager - .claim_wake_for_turn("lead", TeamRunTargetRole::Lead, "conv") - .await - .unwrap(); - - let failed = manager - .record_child_start_failed(&reservation.reservation_id, "spawn failed") - .await - .expect("run should fail before child start"); - - assert_eq!(failed.status, TeamRunStatus::Failed); - assert_eq!(failed.starting_child_count, 0); - assert_eq!(manager.active_run_id().await, None); - assert!(bc.names().contains(&TEAM_RUN_FAILED_EVENT.to_owned())); - } - - #[tokio::test] - async fn cancel_marks_starting_reservation_and_late_start_requests_immediate_cancel() { - let (manager, _) = manager(); - let ack = manager - .accept_user_message("worker", TeamRunTargetRole::Teammate, true, None) - .await - .unwrap(); - manager - .record_pending_wake("worker", TeamRunTargetRole::Teammate, TeamWakeSource::UserIntervention) - .await - .unwrap(); - let reservation = manager - .claim_wake_for_turn("worker", TeamRunTargetRole::Teammate, "conv-worker") - .await - .unwrap(); - - let target = manager.begin_cancel_child("worker").await.unwrap(); - assert!(matches!(target, ChildCancelTarget::Starting(_))); - - let decision = manager - .record_child_started( - &reservation.reservation_id, - ActiveChildTurn { - team_run_id: ack.team_run_id, - slot_id: "worker".into(), - role: TeamRunTargetRole::Teammate, - conversation_id: "conv-worker".into(), - turn_id: "turn-worker".into(), - started_at_ms: now_ms(), - last_slow_notified_at_ms: None, - }, - ) - .await; - - assert!(matches!(decision, ChildStartDecision::CancelImmediately(child) if child.turn_id == "turn-worker")); - } - - #[tokio::test] - async fn stale_child_start_does_not_revive_completed_run() { - let (manager, _) = manager(); - let ack = manager - .accept_user_message("lead", TeamRunTargetRole::Lead, false, None) - .await - .unwrap(); - manager - .record_pending_wake("lead", TeamRunTargetRole::Lead, TeamWakeSource::UserMessage) - .await - .unwrap(); - let reservation = manager - .claim_wake_for_turn("lead", TeamRunTargetRole::Lead, "conv") - .await - .unwrap(); - manager - .record_child_start_failed(&reservation.reservation_id, "failed") - .await - .unwrap(); - - let decision = manager - .record_child_started( - &reservation.reservation_id, - ActiveChildTurn { - team_run_id: ack.team_run_id, - slot_id: "lead".into(), - role: TeamRunTargetRole::Lead, - conversation_id: "conv".into(), - turn_id: "late-turn".into(), - started_at_ms: now_ms(), - last_slow_notified_at_ms: None, - }, - ) - .await; - - assert_eq!(decision, ChildStartDecision::Ignored); - assert_eq!(manager.active_run_id().await, None); - } -} diff --git a/crates/aionui-team/src/team_run/manager.rs b/crates/aionui-team/src/team_run/manager.rs new file mode 100644 index 000000000..d37d7c2f8 --- /dev/null +++ b/crates/aionui-team/src/team_run/manager.rs @@ -0,0 +1,297 @@ +use std::sync::{Arc, Mutex, MutexGuard}; + +use aionui_api_types::{ + TeamRunPayload, TeamRunSource, TeamRunStatus, TeamRunTargetRole, TeamSlotBlockedReason, TeamSlotWorkPayload, + TeamSlotWorkState, +}; +use aionui_common::{TimestampMs, generate_id, now_ms}; +use tracing::info; + +use crate::TeamError; +use crate::events::{ + TEAM_RUN_ACCEPTED_EVENT, TEAM_RUN_CANCELLED_EVENT, TEAM_RUN_COMPLETED_EVENT, TEAM_RUN_FAILED_EVENT, + TEAM_RUN_STARTED_EVENT, TEAM_RUN_UPDATED_EVENT, TeamEventEmitter, +}; +use crate::work_coordinator::{ + CausalBinding, CoordinatorSnapshot, EnqueueRequest, RunBinding, RunCausalityPort, RunWorkSummary, + RuntimeConstraint, SlotPhase, SlotWorkSnapshot, +}; + +const ACTIVE_TURN_SLOW_THRESHOLD_MS: u64 = 10 * 60 * 1000; + +#[derive(Clone)] +struct TeamRunRecord { + team_run_id: String, + source: TeamRunSource, + has_user_intervention: bool, + target_slot_id: String, + target_role: TeamRunTargetRole, + status: TeamRunStatus, + started_at_ms: Option, + completed_at_ms: Option, + cancel_reason: Option, + summary: Option, + accepted_emitted: bool, +} + +impl TeamRunRecord { + fn is_active(&self) -> bool { + matches!( + self.status, + TeamRunStatus::Accepted | TeamRunStatus::Running | TeamRunStatus::Cancelling + ) + } +} + +pub struct TeamRunManager { + team_id: String, + emitter: Arc, + state: Mutex>, +} + +impl TeamRunManager { + pub(crate) fn new(team_id: String, emitter: Arc) -> Self { + Self { + team_id, + emitter, + state: Mutex::new(None), + } + } + + pub(crate) fn bind_user_enqueue(&self, target_slot_id: &str, target_role: TeamRunTargetRole) -> RunBinding { + let mut state = self.lock_state(); + if let Some(run) = state.as_mut().filter(|run| run.is_active()) { + run.has_user_intervention = true; + return RunBinding { + team_run_id: Some(run.team_run_id.clone()), + created_new_run: false, + user_intervention: true, + }; + } + + let run = TeamRunRecord { + team_run_id: generate_id(), + source: TeamRunSource::UserMessage, + has_user_intervention: false, + target_slot_id: target_slot_id.to_owned(), + target_role, + status: TeamRunStatus::Accepted, + started_at_ms: None, + completed_at_ms: None, + cancel_reason: None, + summary: None, + accepted_emitted: false, + }; + let run_id = run.team_run_id.clone(); + *state = Some(run); + drop(state); + info!( + team_id = %self.team_id, + team_run_id = %run_id, + target_slot_id, + "team run accepted" + ); + RunBinding { + team_run_id: Some(run_id), + created_new_run: true, + user_intervention: false, + } + } + + pub fn current_active_run_id(&self) -> Option { + self.lock_state() + .as_ref() + .filter(|run| run.is_active()) + .map(|run| run.team_run_id.clone()) + } + + pub(crate) fn apply_work_summary(&self, summary: RunWorkSummary) { + let event = { + let mut state = self.lock_state(); + let Some(run) = state.as_mut().filter(|run| run.team_run_id == summary.team_run_id) else { + return; + }; + let previous_status = run.status.clone(); + run.summary = Some(summary.clone()); + let no_work = summary.queued_intent_count == 0 + && summary.starting_batch_count == 0 + && summary.running_batch_count == 0 + && summary.active_enqueue_lease_count == 0 + && summary.paused_intent_count == 0; + if summary.failed_intent_count > 0 { + run.status = TeamRunStatus::Failed; + run.completed_at_ms = Some(now_ms()); + } else if no_work { + run.status = if previous_status == TeamRunStatus::Cancelling { + TeamRunStatus::Cancelled + } else { + TeamRunStatus::Completed + }; + run.completed_at_ms = Some(now_ms()); + } else if summary.running_batch_count > 0 { + run.status = TeamRunStatus::Running; + run.started_at_ms.get_or_insert_with(now_ms); + } + let event_name = if !run.accepted_emitted { + run.accepted_emitted = true; + TEAM_RUN_ACCEPTED_EVENT + } else { + match (&previous_status, &run.status) { + (_, TeamRunStatus::Failed) if previous_status != TeamRunStatus::Failed => TEAM_RUN_FAILED_EVENT, + (_, TeamRunStatus::Cancelled) if previous_status != TeamRunStatus::Cancelled => { + TEAM_RUN_CANCELLED_EVENT + } + (_, TeamRunStatus::Completed) if previous_status != TeamRunStatus::Completed => { + TEAM_RUN_COMPLETED_EVENT + } + (_, TeamRunStatus::Running) if previous_status != TeamRunStatus::Running => TEAM_RUN_STARTED_EVENT, + _ => TEAM_RUN_UPDATED_EVENT, + } + }; + Some((event_name, self.payload_locked(run, &summary))) + }; + if let Some((event_name, payload)) = event { + self.emitter.broadcast_team_run(event_name, payload); + } + } + + pub(crate) fn begin_cancel(&self, team_run_id: &str, reason: Option) -> Result<(), TeamError> { + let payload = { + let mut state = self.lock_state(); + let run = state + .as_mut() + .filter(|run| run.team_run_id == team_run_id && run.is_active()) + .ok_or_else(|| TeamError::InvalidRequest(format!("team run is not active: {team_run_id}")))?; + run.status = TeamRunStatus::Cancelling; + run.cancel_reason = reason; + run.summary.as_ref().map(|summary| self.payload_locked(run, summary)) + }; + if let Some(payload) = payload { + self.emitter.broadcast_team_run(TEAM_RUN_UPDATED_EVENT, payload); + } + Ok(()) + } + + pub(crate) fn current_payload(&self, coordinator: &CoordinatorSnapshot) -> Option { + let state = self.lock_state(); + let run = state.as_ref()?; + let summary = run.summary.as_ref().or_else(|| { + coordinator + .active_run_summary + .as_ref() + .filter(|summary| summary.team_run_id == run.team_run_id) + }); + let empty_summary; + let summary = match summary { + Some(summary) => summary, + None => { + empty_summary = RunWorkSummary { + team_run_id: run.team_run_id.clone(), + queued_intent_count: 0, + starting_batch_count: 0, + running_batch_count: 0, + active_enqueue_lease_count: 0, + paused_intent_count: 0, + failed_intent_count: 0, + slots: Vec::new(), + }; + &empty_summary + } + }; + Some(self.payload_locked(run, summary)) + } + + pub(crate) fn publish_snapshot_update(&self, coordinator: &CoordinatorSnapshot) { + if let Some(payload) = self.current_payload(coordinator) + && matches!(payload.status, TeamRunStatus::Accepted | TeamRunStatus::Running) + { + self.emitter.broadcast_team_run(TEAM_RUN_UPDATED_EVENT, payload); + } + } + + fn payload_locked(&self, run: &TeamRunRecord, summary: &RunWorkSummary) -> TeamRunPayload { + TeamRunPayload { + team_id: self.team_id.clone(), + team_run_id: run.team_run_id.clone(), + source: run.source.clone(), + has_user_intervention: run.has_user_intervention, + target_slot_id: run.target_slot_id.clone(), + target_role: run.target_role.clone(), + status: run.status.clone(), + queued_intent_count: summary.queued_intent_count, + starting_batch_count: summary.starting_batch_count, + running_batch_count: summary.running_batch_count, + active_enqueue_lease_count: summary.active_enqueue_lease_count, + slot_work: summary.slots.iter().map(Self::slot_payload).collect(), + } + } + + pub(crate) fn slot_payload(slot: &SlotWorkSnapshot) -> TeamSlotWorkPayload { + let active_turn_elapsed_ms = slot + .active_turn_started_at_ms + .map(|started_at| now_ms().saturating_sub(started_at).max(0) as u64); + TeamSlotWorkPayload { + slot_id: slot.slot_id.clone(), + role: slot.role.clone(), + state: match slot.state { + SlotPhase::Idle => TeamSlotWorkState::Idle, + SlotPhase::Queued => TeamSlotWorkState::Queued, + SlotPhase::Starting => TeamSlotWorkState::Starting, + SlotPhase::Running => TeamSlotWorkState::Running, + SlotPhase::Paused => TeamSlotWorkState::Paused, + SlotPhase::Blocked => TeamSlotWorkState::Blocked, + }, + queued_foreground_count: slot.queued_foreground_count, + queued_background_count: slot.queued_background_count, + active_turn_id: slot.active_turn_id.clone(), + active_turn_started_at_ms: slot.active_turn_started_at_ms, + active_turn_elapsed_ms, + active_turn_slow: active_turn_elapsed_ms.map(|elapsed| elapsed >= ACTIVE_TURN_SLOW_THRESHOLD_MS), + active_turn_slow_threshold_ms: slot.active_turn_id.as_ref().map(|_| ACTIVE_TURN_SLOW_THRESHOLD_MS), + blocked_reason: match slot.runtime_constraint { + RuntimeConstraint::Ready => None, + RuntimeConstraint::Starting { .. } => Some(TeamSlotBlockedReason::RuntimeStarting), + RuntimeConstraint::Failed { .. } => Some(TeamSlotBlockedReason::RuntimeFailed), + RuntimeConstraint::Removing { .. } => Some(TeamSlotBlockedReason::Removing), + RuntimeConstraint::SessionStopped => Some(TeamSlotBlockedReason::SessionStopped), + }, + team_run_id: slot.team_run_id.clone(), + } + } + + fn lock_state(&self) -> MutexGuard<'_, Option> { + self.state.lock().unwrap_or_else(|poisoned| poisoned.into_inner()) + } +} + +impl RunCausalityPort for TeamRunManager { + fn bind_enqueue(&self, request: &EnqueueRequest) -> RunBinding { + match &request.binding { + CausalBinding::UserVisible => self.bind_user_enqueue(&request.slot_id, request.role.clone()), + CausalBinding::ActiveRunOrBackground => RunBinding { + team_run_id: self.current_active_run_id(), + created_new_run: false, + user_intervention: false, + }, + CausalBinding::InheritRunningBatch { .. } | CausalBinding::Background => RunBinding { + team_run_id: None, + created_new_run: false, + user_intervention: false, + }, + } + } + + fn abort_binding(&self, binding: &RunBinding) { + let mut state = self.lock_state(); + if state + .as_ref() + .is_some_and(|run| Some(&run.team_run_id) == binding.team_run_id.as_ref()) + { + *state = None; + } + } + + fn apply_work_summary(&self, summary: RunWorkSummary) { + TeamRunManager::apply_work_summary(self, summary); + } +} diff --git a/crates/aionui-team/src/team_run/mod.rs b/crates/aionui-team/src/team_run/mod.rs new file mode 100644 index 000000000..7f5e0fba3 --- /dev/null +++ b/crates/aionui-team/src/team_run/mod.rs @@ -0,0 +1,17 @@ +mod manager; + +pub use manager::TeamRunManager; + +use aionui_api_types::TeamRunTargetRole; + +use crate::types::TeammateRole; + +pub fn target_role_for(role: TeammateRole) -> TeamRunTargetRole { + match role { + TeammateRole::Lead => TeamRunTargetRole::Lead, + TeammateRole::Teammate => TeamRunTargetRole::Teammate, + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/aionui-team/src/team_run/tests.rs b/crates/aionui-team/src/team_run/tests.rs new file mode 100644 index 000000000..d54851b1d --- /dev/null +++ b/crates/aionui-team/src/team_run/tests.rs @@ -0,0 +1,290 @@ +use std::sync::Arc; + +use aionui_api_types::{TeamRunStatus, TeamRunTargetRole}; + +use crate::events::TeamEventEmitter; +use crate::team_run::TeamRunManager; +use crate::test_utils::workspace_harness::RecordingBroadcaster; +use crate::work_coordinator::{ + CausalBinding, EnqueueRequest, ReconcileDecision, RuntimeConstraint, SlotWorkCoordinator, +}; +use crate::work_source::WorkSource; + +fn coordinator_and_manager() -> (Arc, Arc) { + let broadcaster = Arc::new(RecordingBroadcaster::new()); + let emitter = Arc::new(TeamEventEmitter::new("team-1".into(), broadcaster)); + let manager = Arc::new(TeamRunManager::new("team-1".into(), emitter)); + let coordinator = Arc::new(SlotWorkCoordinator::new( + "team-1".into(), + "generation-1".into(), + manager.clone(), + )); + coordinator.set_runtime_constraint("lead-1", RuntimeConstraint::Ready); + (coordinator, manager) +} + +fn acquire(coordinator: &SlotWorkCoordinator, binding: CausalBinding) -> crate::work_coordinator::EnqueueLease { + coordinator + .acquire_enqueue(EnqueueRequest { + slot_id: "lead-1".into(), + role: TeamRunTargetRole::Lead, + source: WorkSource::UserMessage, + binding, + }) + .unwrap() +} + +#[test] +fn background_work_does_not_create_or_block_a_run() { + let (coordinator, manager) = coordinator_and_manager(); + let lease = coordinator + .acquire_enqueue(EnqueueRequest { + slot_id: "lead-1".into(), + role: TeamRunTargetRole::Lead, + source: WorkSource::RecoveryDrain, + binding: CausalBinding::Background, + }) + .unwrap(); + coordinator.commit_enqueue(&lease, Some("m1".into())).unwrap(); + let ReconcileDecision::Claim(batch) = coordinator.next("lead-1") else { + panic!("background intent must be claimable"); + }; + assert_eq!(manager.current_active_run_id(), None); + coordinator.complete_batch(&batch); + assert_eq!(manager.current_active_run_id(), None); +} + +#[test] +fn enqueue_lease_blocks_completion() { + let (coordinator, manager) = coordinator_and_manager(); + let lease = acquire(&coordinator, CausalBinding::UserVisible); + let run_id = lease.team_run_id.clone().unwrap(); + + let payload = manager.current_payload(&coordinator.snapshot()).unwrap(); + assert_eq!(payload.team_run_id, run_id); + assert_eq!(payload.active_enqueue_lease_count, 1); + assert_eq!(payload.status, TeamRunStatus::Accepted); +} + +#[test] +fn run_completes_after_its_last_intent_is_terminal() { + let (coordinator, manager) = coordinator_and_manager(); + let lease = acquire(&coordinator, CausalBinding::UserVisible); + coordinator.commit_enqueue(&lease, Some("m1".into())).unwrap(); + let ReconcileDecision::Claim(batch) = coordinator.next("lead-1") else { + panic!("user intent must be claimable"); + }; + coordinator.complete_batch(&batch); + + let payload = manager.current_payload(&coordinator.snapshot()).unwrap(); + assert_eq!(payload.status, TeamRunStatus::Completed); + assert_eq!(manager.current_active_run_id(), None); +} + +#[test] +fn runtime_failure_fails_the_related_run() { + let (coordinator, manager) = coordinator_and_manager(); + let lease = acquire(&coordinator, CausalBinding::UserVisible); + coordinator.commit_enqueue(&lease, Some("m1".into())).unwrap(); + coordinator.set_runtime_constraint( + "lead-1", + RuntimeConstraint::Failed { + operation_id: 7, + classification: "attach_failed", + }, + ); + + let payload = manager.current_payload(&coordinator.snapshot()).unwrap(); + assert_eq!(payload.status, TeamRunStatus::Failed); + assert_eq!(manager.current_active_run_id(), None); +} + +#[test] +fn aborted_first_enqueue_removes_the_empty_run() { + let (coordinator, manager) = coordinator_and_manager(); + let lease = acquire(&coordinator, CausalBinding::UserVisible); + coordinator.abort_enqueue(&lease, "mailbox_write_failed"); + + assert_eq!(manager.current_active_run_id(), None); + assert!(manager.current_payload(&coordinator.snapshot()).is_none()); +} + +#[test] +fn aborting_the_creator_lease_preserves_a_concurrent_committed_enqueue() { + let (coordinator, manager) = coordinator_and_manager(); + let creator = acquire(&coordinator, CausalBinding::UserVisible); + let concurrent = acquire(&coordinator, CausalBinding::UserVisible); + let run_id = creator.team_run_id.clone().unwrap(); + assert_eq!(concurrent.team_run_id.as_deref(), Some(run_id.as_str())); + + coordinator + .commit_enqueue(&concurrent, Some("m-concurrent".into())) + .unwrap(); + coordinator.abort_enqueue(&creator, "mailbox_write_failed"); + + let payload = manager + .current_payload(&coordinator.snapshot()) + .expect("the concurrent enqueue must retain its run"); + assert_eq!(payload.team_run_id, run_id); + assert_eq!(payload.queued_intent_count, 1); + assert_eq!(manager.current_active_run_id(), Some(payload.team_run_id)); +} + +#[test] +fn accepted_event_is_emitted_only_after_the_durable_enqueue_commits() { + let broadcaster = Arc::new(RecordingBroadcaster::new()); + let emitter = Arc::new(TeamEventEmitter::new("team-1".into(), broadcaster.clone())); + let manager = Arc::new(TeamRunManager::new("team-1".into(), emitter)); + let coordinator = SlotWorkCoordinator::new("team-1".into(), "generation-1".into(), manager); + coordinator.set_runtime_constraint("lead-1", RuntimeConstraint::Ready); + + let lease = acquire(&coordinator, CausalBinding::UserVisible); + assert!(broadcaster.events_by_name("team.runAccepted").is_empty()); + + coordinator.commit_enqueue(&lease, Some("m1".into())).unwrap(); + assert_eq!(broadcaster.events_by_name("team.runAccepted").len(), 1); +} + +#[test] +fn mcp_message_inherits_the_callers_running_batch_causality() { + let (coordinator, manager) = coordinator_and_manager(); + coordinator.set_runtime_constraint("worker-1", RuntimeConstraint::Ready); + coordinator.set_runtime_constraint("worker-2", RuntimeConstraint::Ready); + + let user = acquire(&coordinator, CausalBinding::UserVisible); + let run_id = user.team_run_id.clone().unwrap(); + coordinator.commit_enqueue(&user, Some("m-user".into())).unwrap(); + let ReconcileDecision::Claim(user_batch) = coordinator.next("lead-1") else { + panic!("user batch must be claimed"); + }; + coordinator.mark_started(&user_batch, "turn-user"); + let inherited = coordinator + .acquire_enqueue(EnqueueRequest { + slot_id: "worker-1".into(), + role: TeamRunTargetRole::Teammate, + source: WorkSource::McpSendMessage, + binding: CausalBinding::InheritRunningBatch { + caller_slot_id: "lead-1".into(), + }, + }) + .unwrap(); + assert_eq!(inherited.team_run_id.as_deref(), Some(run_id.as_str())); + coordinator.abort_enqueue(&inherited, "test_complete"); + + coordinator.complete_batch(&user_batch); + let background = coordinator + .acquire_enqueue(EnqueueRequest { + slot_id: "worker-1".into(), + role: TeamRunTargetRole::Teammate, + source: WorkSource::RecoveryDrain, + binding: CausalBinding::Background, + }) + .unwrap(); + coordinator + .commit_enqueue(&background, Some("m-background".into())) + .unwrap(); + let ReconcileDecision::Claim(background_batch) = coordinator.next("worker-1") else { + panic!("background batch must be claimed"); + }; + coordinator.mark_started(&background_batch, "turn-background"); + + let unrelated_user = coordinator + .acquire_enqueue(EnqueueRequest { + slot_id: "worker-2".into(), + role: TeamRunTargetRole::Teammate, + source: WorkSource::UserMessage, + binding: CausalBinding::UserVisible, + }) + .unwrap(); + assert!(unrelated_user.team_run_id.is_some()); + let background_child = coordinator + .acquire_enqueue(EnqueueRequest { + slot_id: "lead-1".into(), + role: TeamRunTargetRole::Lead, + source: WorkSource::McpSendMessage, + binding: CausalBinding::InheritRunningBatch { + caller_slot_id: "worker-1".into(), + }, + }) + .unwrap(); + assert_eq!(background_child.team_run_id, None); + + coordinator.abort_enqueue(&background_child, "test_complete"); + coordinator.abort_enqueue(&unrelated_user, "test_complete"); + coordinator.complete_batch(&background_batch); + assert!(manager.current_active_run_id().is_none()); +} + +#[test] +fn published_dynamic_attach_failure_blocks_only_related_work_and_preserves_healthy_runtimes() { + let (coordinator, manager) = coordinator_and_manager(); + coordinator.set_runtime_constraint("worker-1", RuntimeConstraint::Ready); + let user = coordinator + .acquire_enqueue(EnqueueRequest { + slot_id: "worker-1".into(), + role: TeamRunTargetRole::Teammate, + source: WorkSource::UserMessage, + binding: CausalBinding::UserVisible, + }) + .unwrap(); + coordinator.commit_enqueue(&user, Some("m-user".into())).unwrap(); + let healthy = coordinator + .acquire_enqueue(EnqueueRequest { + slot_id: "lead-1".into(), + role: TeamRunTargetRole::Lead, + source: WorkSource::RecoveryDrain, + binding: CausalBinding::Background, + }) + .unwrap(); + coordinator.commit_enqueue(&healthy, Some("m-healthy".into())).unwrap(); + + coordinator.set_runtime_constraint( + "worker-1", + RuntimeConstraint::Failed { + operation_id: 9, + classification: "attach_failed", + }, + ); + + let run = manager.current_payload(&coordinator.snapshot()).unwrap(); + assert_eq!(run.status, TeamRunStatus::Failed); + let healthy = coordinator.slot_snapshot("lead-1").unwrap(); + assert_eq!(healthy.runtime_constraint, RuntimeConstraint::Ready); + assert_eq!(healthy.queued_background_count, 1); +} + +#[test] +fn manual_or_leader_add_during_run_inherits_run_causality() { + let (coordinator, _manager) = coordinator_and_manager(); + coordinator.set_runtime_constraint("worker-1", RuntimeConstraint::Ready); + coordinator.set_runtime_constraint("worker-2", RuntimeConstraint::Ready); + let user = acquire(&coordinator, CausalBinding::UserVisible); + let run_id = user.team_run_id.clone().unwrap(); + coordinator.commit_enqueue(&user, Some("m-user".into())).unwrap(); + let ReconcileDecision::Claim(batch) = coordinator.next("lead-1") else { + panic!("lead batch must be claimed"); + }; + coordinator.mark_started(&batch, "turn-lead"); + + let manual = coordinator + .acquire_enqueue(EnqueueRequest { + slot_id: "worker-1".into(), + role: TeamRunTargetRole::Teammate, + source: WorkSource::SpawnWelcome, + binding: CausalBinding::ActiveRunOrBackground, + }) + .unwrap(); + let leader_spawn = coordinator + .acquire_enqueue(EnqueueRequest { + slot_id: "worker-2".into(), + role: TeamRunTargetRole::Teammate, + source: WorkSource::SpawnWelcome, + binding: CausalBinding::InheritRunningBatch { + caller_slot_id: "lead-1".into(), + }, + }) + .unwrap(); + + assert_eq!(manual.team_run_id.as_deref(), Some(run_id.as_str())); + assert_eq!(leader_spawn.team_run_id.as_deref(), Some(run_id.as_str())); +} diff --git a/crates/aionui-team/src/test_utils.rs b/crates/aionui-team/src/test_utils.rs index 92b80a92b..ff409483b 100644 --- a/crates/aionui-team/src/test_utils.rs +++ b/crates/aionui-team/src/test_utils.rs @@ -8,6 +8,7 @@ pub struct MockState { pub messages: Vec, pub tasks: Vec, pub fail_message_writes: bool, + pub fail_task_lists: bool, } pub struct MockTeamRepo { @@ -20,12 +21,6 @@ impl MockTeamRepo { state: Mutex::new(MockState::default()), } } - - pub fn with_message_write_failure() -> Self { - let repo = Self::new(); - repo.state.lock().unwrap().fail_message_writes = true; - repo - } } #[async_trait::async_trait] @@ -163,6 +158,9 @@ impl ITeamRepository for MockTeamRepo { async fn list_tasks(&self, team_id: &str) -> Result, DbError> { let state = self.state.lock().unwrap(); + if state.fail_task_lists { + return Err(DbError::Init("forced task list failure".into())); + } let tasks = state.tasks.iter().filter(|t| t.team_id == team_id).cloned().collect(); Ok(tasks) } @@ -204,7 +202,9 @@ pub(crate) mod workspace_harness { use std::sync::{Arc, Mutex}; use aionui_ai_agent::{AgentError, IWorkerTaskManager}; - use aionui_api_types::{CreateTeamRequest, WebSocketMessage}; + use aionui_api_types::{ + AcpConfigOptionDto, AcpConfigSelectOptionDto, CreateTeamRequest, GetConfigOptionsResponse, WebSocketMessage, + }; use aionui_common::{AgentKillReason, AgentType, PaginatedResult, now_ms}; use aionui_db::models::{ AgentMetadataRow, AssistantDefinitionRow, AssistantOverlayRow, ConversationRow, MessageRow, TeamRow, @@ -248,6 +248,18 @@ pub(crate) mod workspace_harness { .find(|c| c.id == id) .and_then(|c| serde_json::from_str(&c.extra).ok()) } + + pub(crate) fn mark_runtime_not_ready(&self, id: &str) { + let mut conversations = self.conversations.lock().unwrap(); + let conversation = conversations + .iter_mut() + .find(|c| c.id == id) + .expect("conversation exists"); + let mut extra: serde_json::Value = + serde_json::from_str(&conversation.extra).unwrap_or_else(|_| serde_json::json!({})); + extra["runtime_not_ready"] = serde_json::Value::Bool(true); + conversation.extra = serde_json::to_string(&extra).unwrap(); + } } #[async_trait] @@ -428,6 +440,9 @@ pub(crate) mod workspace_harness { if let Some(ref lead_id) = params.lead_agent_id { team.lead_agent_id = Some(lead_id.clone()); } + if let Some(ref session_mode) = params.session_mode { + team.session_mode = Some(session_mode.clone()); + } team.updated_at = now_ms(); Ok(()) } @@ -627,6 +642,44 @@ pub(crate) mod workspace_harness { .await } + async fn get_config_options(&self, conversation_id: &str) -> Result { + let extra = self + .repo + .get_extra(conversation_id) + .ok_or_else(|| TeamError::AgentNotFound(conversation_id.to_owned()))?; + if extra + .get("runtime_not_ready") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false) + { + return Err(TeamError::RuntimeNotReady { + conversation_id: conversation_id.to_owned(), + }); + } + let model = extra + .get("current_model_id") + .and_then(serde_json::Value::as_str) + .unwrap_or("mock-model") + .to_owned(); + Ok(GetConfigOptionsResponse { + config_options: vec![AcpConfigOptionDto { + id: "model".to_owned(), + name: None, + label: Some("Model".to_owned()), + description: None, + category: Some("model".to_owned()), + option_type: "select".to_owned(), + current_value: Some(model.clone()), + options: vec![AcpConfigSelectOptionDto { + value: model.clone(), + name: None, + label: Some(model), + description: None, + }], + }], + }) + } + async fn warmup_agent_process( &self, _user_id: &str, @@ -671,10 +724,44 @@ pub(crate) mod workspace_harness { } } - struct NullBroadcaster; + type BroadcastObserver = Arc) + Send + Sync>; + + pub(crate) struct RecordingBroadcaster { + events: std::sync::Mutex>>, + observer: std::sync::Mutex>, + } + + impl RecordingBroadcaster { + pub(crate) fn new() -> Self { + Self { + events: std::sync::Mutex::new(Vec::new()), + observer: std::sync::Mutex::new(None), + } + } + + pub(crate) fn events_by_name(&self, name: &str) -> Vec> { + self.events + .lock() + .unwrap() + .iter() + .filter(|event| event.name == name) + .cloned() + .collect() + } - impl EventBroadcaster for NullBroadcaster { - fn broadcast(&self, _event: WebSocketMessage) {} + pub(crate) fn set_observer(&self, observer: BroadcastObserver) { + *self.observer.lock().unwrap() = Some(observer); + } + } + + impl EventBroadcaster for RecordingBroadcaster { + fn broadcast(&self, event: WebSocketMessage) { + self.events.lock().unwrap().push(event.clone()); + let observer = self.observer.lock().unwrap().clone(); + if let Some(observer) = observer { + observer(&event); + } + } } struct NoopTaskManager; @@ -883,7 +970,7 @@ pub(crate) mod workspace_harness { async fn run_agent_turn(&self, request: AgentTurnRequest) -> Result { if let Some(on_started) = request.on_started.as_ref() { on_started(AgentTurnStarted { - team_run_id: request.team_run_id.clone().expect("team run id"), + team_run_id: request.team_run_id.clone(), slot_id: request.slot_id.clone(), role: request.role.clone(), conversation_id: request.conversation_id.clone(), @@ -920,14 +1007,35 @@ pub(crate) mod workspace_harness { Arc, Arc, ) { + let (svc, team_repo, task_manager, conv_repo, _broadcaster) = + setup_with_factory_metadata_team_repo_conversation_repo_and_broadcaster(); + (svc, team_repo, task_manager, conv_repo) + } + + type ServiceHarness = ( + Arc, + Arc, + Arc, + Arc, + Arc, + ); + + pub(crate) fn setup_with_factory_metadata_team_repo_conversation_repo_and_broadcaster() -> ServiceHarness { + let task_manager: Arc = Arc::new(NoopTaskManager); + setup_with_factory_metadata_team_repo_conversation_repo_broadcaster_and_task_manager(task_manager) + } + + pub(crate) fn setup_with_factory_metadata_team_repo_conversation_repo_broadcaster_and_task_manager( + task_manager: Arc, + ) -> ServiceHarness { let team_repo = Arc::new(FullMockTeamRepo::new()); let team_repo_dyn: Arc = team_repo.clone(); let conv_repo = Arc::new(MockConversationRepo::new()); - let broadcaster: Arc = Arc::new(NullBroadcaster); + let broadcaster = Arc::new(RecordingBroadcaster::new()); + let broadcaster_dyn: Arc = broadcaster.clone(); let conversation_ports = Arc::new(FakeConversationPorts::new(conv_repo.clone())); let conversation_port: Arc = conversation_ports.clone(); let projection_store: Arc = conversation_ports.clone(); - let task_manager: Arc = Arc::new(NoopTaskManager); let svc = TeamSessionService::new( team_repo_dyn, Arc::new(EmptyAgentMetadataRepo), @@ -937,13 +1045,13 @@ pub(crate) mod workspace_harness { Arc::new(EmptyProviderRepo), conversation_port, projection_store, - broadcaster, + broadcaster_dyn, task_manager.clone(), Arc::new(NoopTurnPort), Arc::new(NoopCancellationPort), Arc::new(std::path::PathBuf::from("/tmp/aioncore-test")), ); - (svc, team_repo, task_manager, conv_repo) + (svc, team_repo, task_manager, conv_repo, broadcaster) } pub(crate) async fn force_team_workspace(repo: &Arc, team_id: &str, workspace: &str) { diff --git a/crates/aionui-team/src/wake.rs b/crates/aionui-team/src/wake.rs deleted file mode 100644 index 2db79833c..000000000 --- a/crates/aionui-team/src/wake.rs +++ /dev/null @@ -1,72 +0,0 @@ -use std::fmt; - -#[allow(dead_code)] -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum TeamWakeSource { - UserMessage, - UserIntervention, - McpSendMessage, - McpShutdownRequest, - SpawnWelcome, - SpawnAttachFailure, - IdleNotification, - InterruptedNotification, - CrashNotification, - InactivityTimeout, - ShutdownRejected, - RecoveryDrain, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum TeamWakeClass { - Foreground, - Background, - SystemRecovery, - Lifecycle, -} - -impl TeamWakeSource { - pub(crate) fn as_str(self) -> &'static str { - match self { - Self::UserMessage => "user_message", - Self::UserIntervention => "user_intervention", - Self::McpSendMessage => "mcp_send_message", - Self::McpShutdownRequest => "mcp_shutdown_request", - Self::SpawnWelcome => "spawn_welcome", - Self::SpawnAttachFailure => "spawn_attach_failure", - Self::IdleNotification => "idle_notification", - Self::InterruptedNotification => "interrupted_notification", - Self::CrashNotification => "crash_notification", - Self::InactivityTimeout => "inactivity_timeout", - Self::ShutdownRejected => "shutdown_rejected", - Self::RecoveryDrain => "recovery_drain", - } - } - - pub(crate) fn class(self) -> TeamWakeClass { - match self { - Self::UserMessage | Self::UserIntervention => TeamWakeClass::Foreground, - Self::McpSendMessage | Self::IdleNotification | Self::InterruptedNotification => TeamWakeClass::Background, - Self::CrashNotification - | Self::InactivityTimeout - | Self::SpawnAttachFailure - | Self::ShutdownRejected - | Self::RecoveryDrain => TeamWakeClass::SystemRecovery, - Self::SpawnWelcome | Self::McpShutdownRequest => TeamWakeClass::Lifecycle, - } - } - - pub(crate) fn bypasses_pause(self) -> bool { - matches!(self.class(), TeamWakeClass::Foreground | TeamWakeClass::Lifecycle) - } - - pub(crate) fn resumes_paused_slot(self) -> bool { - matches!(self, Self::UserMessage | Self::UserIntervention) - } -} - -impl fmt::Display for TeamWakeSource { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.write_str(self.as_str()) - } -} diff --git a/crates/aionui-team/src/work_coordinator/coordinator.rs b/crates/aionui-team/src/work_coordinator/coordinator.rs new file mode 100644 index 000000000..0d31dd029 --- /dev/null +++ b/crates/aionui-team/src/work_coordinator/coordinator.rs @@ -0,0 +1,978 @@ +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque}; +use std::sync::{Arc, Mutex}; + +use aionui_api_types::TeamRunTargetRole; +use aionui_common::{generate_id, now_ms}; +use tracing::{debug, info, warn}; + +use super::model::*; +use crate::TeamError; +use crate::work_source::WorkSource; + +#[derive(Debug, Clone)] +pub(super) struct EnqueueLeaseRecord { + pub(super) lease: EnqueueLease, + binding: RunBinding, +} + +#[derive(Debug, Clone)] +pub(super) struct ActiveBatch { + pub(super) batch: WorkBatch, + pub(super) turn_id: Option, + pub(super) started_at_ms: Option, +} + +#[derive(Debug)] +pub(super) struct SlotState { + pub(super) role: TeamRunTargetRole, + pub(super) foreground: VecDeque, + pub(super) control: VecDeque, + pub(super) background: VecDeque, + pub(super) active: Option, + pub(super) paused: bool, + pub(super) runtime_constraint: RuntimeConstraint, + known_unread_message_ids: HashSet, + removed: bool, +} + +impl SlotState { + fn new(role: TeamRunTargetRole) -> Self { + Self { + role, + foreground: VecDeque::new(), + control: VecDeque::new(), + background: VecDeque::new(), + active: None, + paused: false, + runtime_constraint: RuntimeConstraint::Starting { operation_id: 0 }, + known_unread_message_ids: HashSet::new(), + removed: false, + } + } + + fn queue(&self, priority: WorkPriority) -> &VecDeque { + match priority { + WorkPriority::Foreground => &self.foreground, + WorkPriority::Control => &self.control, + WorkPriority::Background => &self.background, + } + } + + fn queue_mut(&mut self, priority: WorkPriority) -> &mut VecDeque { + match priority { + WorkPriority::Foreground => &mut self.foreground, + WorkPriority::Control => &mut self.control, + WorkPriority::Background => &mut self.background, + } + } + + pub(super) fn queued_ids(&self) -> impl Iterator { + self.foreground + .iter() + .chain(self.control.iter()) + .chain(self.background.iter()) + } + + fn remove_queued(&mut self, intent_id: &str) { + for queue in [&mut self.foreground, &mut self.control, &mut self.background] { + queue.retain(|candidate| candidate != intent_id); + } + } +} + +#[derive(Default)] +pub(super) struct CoordinatorState { + pub(super) slots: BTreeMap, + pub(super) intents: HashMap, + pub(super) enqueue_leases: HashMap, + next_operation_id: u64, +} + +pub(crate) struct SlotWorkCoordinator { + team_id: String, + pub(super) session_generation: String, + pub(super) run_causality: Arc, + pub(super) state: Mutex, +} + +impl SlotWorkCoordinator { + pub(crate) fn new(team_id: String, session_generation: String, run_causality: Arc) -> Self { + Self { + team_id, + session_generation, + run_causality, + state: Mutex::new(CoordinatorState::default()), + } + } + + pub(crate) fn acquire_enqueue(&self, request: EnqueueRequest) -> Result { + let mut state = self.lock_state(); + let slot = state + .slots + .entry(request.slot_id.clone()) + .or_insert_with(|| SlotState::new(request.role.clone())); + slot.role = request.role.clone(); + Self::validate_enqueue_constraint(&request.slot_id, slot)?; + + let binding = match &request.binding { + CausalBinding::Background => RunBinding { + team_run_id: None, + created_new_run: false, + user_intervention: false, + }, + CausalBinding::InheritRunningBatch { caller_slot_id } => RunBinding { + team_run_id: state + .slots + .get(caller_slot_id) + .and_then(|caller| caller.active.as_ref()) + .and_then(|active| active.batch.team_run_ids.first().cloned()), + created_new_run: false, + user_intervention: false, + }, + CausalBinding::UserVisible | CausalBinding::ActiveRunOrBackground => { + self.run_causality.bind_enqueue(&request) + } + }; + let lease = EnqueueLease { + lease_id: generate_id(), + session_generation: self.session_generation.clone(), + slot_id: request.slot_id, + role: request.role, + source: request.source, + team_run_id: binding.team_run_id.clone(), + }; + state.enqueue_leases.insert( + lease.lease_id.clone(), + EnqueueLeaseRecord { + lease: lease.clone(), + binding, + }, + ); + Ok(lease) + } + + pub(crate) fn commit_enqueue( + &self, + lease: &EnqueueLease, + mailbox_message_id: Option, + ) -> Result { + let mut state = self.lock_state(); + let record = state + .enqueue_leases + .get(&lease.lease_id) + .cloned() + .ok_or_else(|| TeamError::InvalidRequest("enqueue lease is no longer active".into()))?; + if record.lease != *lease || lease.session_generation != self.session_generation { + return Err(TeamError::InvalidRequest("stale enqueue lease owner".into())); + } + if lease.source.requires_mailbox_message() && mailbox_message_id.is_none() { + return Err(TeamError::InvalidRequest(format!( + "{:?} requires a mailbox message", + lease.source + ))); + } + + let disposition = { + let slot = state + .slots + .get(&lease.slot_id) + .ok_or_else(|| TeamError::AgentNotFound(lease.slot_id.clone()))?; + Self::validate_enqueue_constraint(&lease.slot_id, slot)?; + match slot.runtime_constraint { + RuntimeConstraint::Starting { .. } => EnqueueDisposition::BlockedRuntimeStarting, + RuntimeConstraint::Ready if slot.active.is_none() && slot.queued_ids().next().is_none() => { + EnqueueDisposition::Accepted + } + RuntimeConstraint::Ready => EnqueueDisposition::Queued, + RuntimeConstraint::Failed { .. } + | RuntimeConstraint::Removing { .. } + | RuntimeConstraint::SessionStopped => unreachable!("validated above"), + } + }; + + state.enqueue_leases.remove(&lease.lease_id); + let intent_id = generate_id(); + let intent = WorkIntent { + intent_id: intent_id.clone(), + session_generation: self.session_generation.clone(), + slot_id: lease.slot_id.clone(), + role: lease.role.clone(), + source: lease.source, + priority: lease.source.priority(), + mailbox_message_id: mailbox_message_id.clone(), + team_run_id: lease.team_run_id.clone(), + created_at_ms: now_ms(), + state: WorkIntentState::Queued, + }; + state.intents.insert(intent_id.clone(), intent.clone()); + let slot = state.slots.get_mut(&lease.slot_id).expect("validated slot exists"); + if lease.source.resumes_paused_slot() { + slot.paused = false; + } + if let Some(message_id) = mailbox_message_id { + slot.known_unread_message_ids.insert(message_id); + } + slot.queue_mut(intent.priority).push_back(intent_id.clone()); + let slot_snapshot = Self::slot_snapshot_locked(&state, &lease.slot_id).expect("committed slot exists"); + let summaries = Self::run_summaries_locked(&state, lease.team_run_id.iter().cloned()); + drop(state); + self.publish_run_summaries(summaries); + + info!( + team_id = %self.team_id, + session_generation = %self.session_generation, + slot_id = %lease.slot_id, + intent_id = %intent_id, + team_run_id = lease.team_run_id.as_deref().unwrap_or("background"), + source = ?lease.source, + priority = ?lease.source.priority(), + "team work enqueue committed" + ); + Ok(EnqueueCommit { + intent_id, + team_run_id: lease.team_run_id.clone(), + disposition, + slot: slot_snapshot, + }) + } + + pub(crate) fn abort_enqueue(&self, lease: &EnqueueLease, classification: &'static str) -> CommitResult { + let mut state = self.lock_state(); + let Some(record) = state.enqueue_leases.remove(&lease.lease_id) else { + return CommitResult::StaleOwner; + }; + if record.lease != *lease || lease.session_generation != self.session_generation { + state.enqueue_leases.insert(record.lease.lease_id.clone(), record); + return CommitResult::StaleOwner; + } + let remove_empty_run = lease.team_run_id.as_ref().is_some_and(|team_run_id| { + !state + .intents + .values() + .any(|intent| intent.team_run_id.as_ref() == Some(team_run_id)) + && !state + .enqueue_leases + .values() + .any(|candidate| candidate.lease.team_run_id.as_ref() == Some(team_run_id)) + }); + let summaries = Self::run_summaries_locked(&state, lease.team_run_id.iter().cloned()); + drop(state); + if remove_empty_run { + self.run_causality.abort_binding(&record.binding); + } + self.publish_run_summaries(summaries); + debug!( + team_id = %self.team_id, + session_generation = %self.session_generation, + slot_id = %lease.slot_id, + lease_id = %lease.lease_id, + classification, + "team work enqueue aborted" + ); + CommitResult::Committed + } + + pub(crate) fn reconcile_mailbox( + &self, + slot_id: &str, + unread_message_ids: &[String], + role: TeamRunTargetRole, + ) -> ReconcileProjection { + let mut state = self.lock_state(); + let unread = unread_message_ids.iter().cloned().collect::>(); + let slot = state + .slots + .entry(slot_id.to_owned()) + .or_insert_with(|| SlotState::new(role.clone())); + slot.role = role.clone(); + slot.known_unread_message_ids = unread.clone(); + + let candidates = slot.queued_ids().cloned().collect::>(); + let mut retained_intent_ids = Vec::new(); + let mut cleared_stale_intent_ids = Vec::new(); + for intent_id in candidates { + let Some(intent) = state.intents.get_mut(&intent_id) else { + continue; + }; + let Some(message_id) = intent.mailbox_message_id.as_ref() else { + retained_intent_ids.push(intent_id); + continue; + }; + if unread.contains(message_id) { + retained_intent_ids.push(intent_id); + } else { + intent.state = WorkIntentState::Completed; + cleared_stale_intent_ids.push(intent_id); + } + } + if let Some(slot) = state.slots.get_mut(slot_id) { + for intent_id in &cleared_stale_intent_ids { + slot.remove_queued(intent_id); + } + } + + let mut created_recovery_intent_ids = Vec::new(); + for message_id in unread_message_ids { + let already_owned = state.intents.values().any(|intent| { + intent.slot_id == slot_id + && intent.mailbox_message_id.as_deref() == Some(message_id.as_str()) + && !intent.state.is_terminal() + }); + if already_owned { + continue; + } + let intent_id = generate_id(); + state.intents.insert( + intent_id.clone(), + WorkIntent { + intent_id: intent_id.clone(), + session_generation: self.session_generation.clone(), + slot_id: slot_id.to_owned(), + role: role.clone(), + source: WorkSource::RecoveryDrain, + priority: WorkPriority::Background, + mailbox_message_id: Some(message_id.clone()), + team_run_id: None, + created_at_ms: now_ms(), + state: WorkIntentState::Queued, + }, + ); + state + .slots + .get_mut(slot_id) + .expect("reconciled slot exists") + .background + .push_back(intent_id.clone()); + created_recovery_intent_ids.push(intent_id); + } + debug!( + team_id = %self.team_id, + session_generation = %self.session_generation, + slot_id, + unread_count = unread_message_ids.len(), + created_recovery_count = created_recovery_intent_ids.len(), + cleared_stale_count = cleared_stale_intent_ids.len(), + "team work mailbox reconciled" + ); + ReconcileProjection { + created_recovery_intent_ids, + retained_intent_ids, + cleared_stale_intent_ids, + } + } + + pub(crate) fn next(&self, slot_id: &str) -> ReconcileDecision { + let mut state = self.lock_state(); + let Some(slot) = state.slots.get(slot_id) else { + return ReconcileDecision::Quiescent; + }; + if slot.active.is_some() { + return ReconcileDecision::WaitingForCompletion; + } + if !matches!(slot.runtime_constraint, RuntimeConstraint::Ready) { + return ReconcileDecision::Blocked(slot.runtime_constraint.clone()); + } + if slot.paused { + return ReconcileDecision::Quiescent; + } + + let selected_priority = [ + WorkPriority::Foreground, + WorkPriority::Control, + WorkPriority::Background, + ] + .into_iter() + .find(|priority| !slot.queue(*priority).is_empty()); + let Some(priority) = selected_priority else { + info!( + team_id = %self.team_id, + session_generation = %self.session_generation, + slot_id, + "team work slot quiescent" + ); + return ReconcileDecision::Quiescent; + }; + + let queued_ids = slot.queue(priority).iter().cloned().collect::>(); + let message_intent_ids = queued_ids + .iter() + .filter(|intent_id| { + state + .intents + .get(*intent_id) + .is_some_and(|intent| intent.mailbox_message_id.is_some()) + }) + .cloned() + .collect::>(); + if message_intent_ids.is_empty() { + return ReconcileDecision::SettleSignals(queued_ids); + } + + state.next_operation_id = state.next_operation_id.saturating_add(1); + let operation_id = state.next_operation_id; + let batch_id = generate_id(); + let mut mailbox_message_ids = Vec::new(); + let mut team_run_ids = Vec::new(); + for intent_id in &message_intent_ids { + let intent = state.intents.get_mut(intent_id).expect("queued intent exists"); + intent.state = WorkIntentState::Starting { + batch_id: batch_id.clone(), + operation_id, + }; + if let Some(message_id) = &intent.mailbox_message_id { + mailbox_message_ids.push(message_id.clone()); + } + if let Some(team_run_id) = &intent.team_run_id + && !team_run_ids.contains(team_run_id) + { + team_run_ids.push(team_run_id.clone()); + } + } + let batch = WorkBatch { + batch_id, + session_generation: self.session_generation.clone(), + slot_id: slot_id.to_owned(), + intent_ids: message_intent_ids.clone(), + mailbox_message_ids, + highest_priority: priority, + team_run_ids: team_run_ids.clone(), + operation_id, + }; + let slot = state.slots.get_mut(slot_id).expect("selected slot exists"); + for intent_id in &message_intent_ids { + slot.remove_queued(intent_id); + } + slot.active = Some(ActiveBatch { + batch: batch.clone(), + turn_id: None, + started_at_ms: None, + }); + let summaries = Self::run_summaries_locked(&state, team_run_ids); + drop(state); + self.publish_run_summaries(summaries); + info!( + team_id = %self.team_id, + session_generation = %self.session_generation, + slot_id, + batch_id = %batch.batch_id, + operation_id, + intent_count = batch.intent_ids.len(), + message_count = batch.mailbox_message_ids.len(), + priority = ?batch.highest_priority, + "team work batch claimed" + ); + ReconcileDecision::Claim(batch) + } + + pub(crate) fn mark_started(&self, batch: &WorkBatch, turn_id: &str) -> StartCommitResult { + let mut state = self.lock_state(); + if !self.is_current_batch(&state, batch) { + let was_cancelled = batch.intent_ids.iter().all(|intent_id| { + state + .intents + .get(intent_id) + .is_some_and(|intent| matches!(intent.state, WorkIntentState::Cancelled { .. })) + }); + if was_cancelled { + return StartCommitResult::CancelImmediately; + } + self.log_stale_batch(batch, "mark_started"); + return StartCommitResult::StaleOwner; + } + let cancel_immediately = state + .slots + .get(&batch.slot_id) + .is_some_and(|slot| slot.paused || !matches!(slot.runtime_constraint, RuntimeConstraint::Ready)); + for intent_id in &batch.intent_ids { + if let Some(intent) = state.intents.get_mut(intent_id) { + intent.state = WorkIntentState::Running { + batch_id: batch.batch_id.clone(), + operation_id: batch.operation_id, + turn_id: turn_id.to_owned(), + }; + } + } + if let Some(active) = state + .slots + .get_mut(&batch.slot_id) + .and_then(|slot| slot.active.as_mut()) + { + active.turn_id = Some(turn_id.to_owned()); + active.started_at_ms = Some(now_ms()); + } + let summaries = Self::run_summaries_locked(&state, batch.team_run_ids.iter().cloned()); + drop(state); + self.publish_run_summaries(summaries); + if cancel_immediately { + StartCommitResult::CancelImmediately + } else { + StartCommitResult::Accepted + } + } + + pub(crate) fn retry_start(&self, batch: &WorkBatch, classification: &'static str) -> CommitResult { + let mut state = self.lock_state(); + if !self.is_current_batch(&state, batch) { + self.log_stale_batch(batch, "retry_start"); + return CommitResult::StaleOwner; + } + for intent_id in &batch.intent_ids { + if let Some(intent) = state.intents.get_mut(intent_id) { + intent.state = WorkIntentState::Queued; + } + } + let slot = state.slots.get_mut(&batch.slot_id).expect("current batch slot exists"); + slot.active = None; + for intent_id in batch.intent_ids.iter().rev() { + slot.queue_mut(batch.highest_priority).push_front(intent_id.clone()); + } + let summaries = Self::run_summaries_locked(&state, batch.team_run_ids.iter().cloned()); + drop(state); + self.publish_run_summaries(summaries); + debug!( + team_id = %self.team_id, + session_generation = %self.session_generation, + slot_id = %batch.slot_id, + batch_id = %batch.batch_id, + operation_id = batch.operation_id, + classification, + "team work batch start will retry" + ); + CommitResult::Committed + } + + pub(crate) fn complete_batch(&self, batch: &WorkBatch) -> CommitResult { + self.terminalize_batch(batch, WorkIntentState::Completed, "completed") + } + + pub(crate) fn fail_batch(&self, batch: &WorkBatch, classification: &'static str) -> CommitResult { + self.terminalize_batch(batch, WorkIntentState::Failed { classification }, classification) + } + + pub(crate) fn cancel_batch(&self, batch: &WorkBatch, classification: &'static str) -> CommitResult { + self.terminalize_batch(batch, WorkIntentState::Cancelled { classification }, classification) + } + + pub(crate) fn cancel_run(&self, team_run_id: &str) -> CancelRunWorkResult { + let mut state = self.lock_state(); + let mut cancel_targets = Vec::new(); + let mut terminal_message_ids = Vec::new(); + let active_slots = state + .slots + .iter() + .filter_map(|(slot_id, slot)| { + let active = slot.active.as_ref()?; + active + .batch + .intent_ids + .iter() + .any(|intent_id| { + state + .intents + .get(intent_id) + .is_some_and(|intent| intent.team_run_id.as_deref() == Some(team_run_id)) + }) + .then(|| (slot_id.clone(), active.clone())) + }) + .collect::>(); + for (slot_id, active) in active_slots { + cancel_targets.push(BatchCancelTarget { + batch: active.batch.clone(), + turn_id: active.turn_id, + }); + let mut retained = Vec::new(); + for intent_id in &active.batch.intent_ids { + let intent = state.intents.get_mut(intent_id).expect("active intent exists"); + if intent.team_run_id.as_deref() == Some(team_run_id) { + if let Some(message_id) = &intent.mailbox_message_id { + terminal_message_ids.push(message_id.clone()); + } + intent.state = WorkIntentState::Cancelled { + classification: "run_cancelled", + }; + } else { + intent.state = WorkIntentState::Queued; + retained.push((intent.priority, intent_id.clone())); + } + } + let slot = state.slots.get_mut(&slot_id).expect("active slot exists"); + slot.active = None; + for (priority, intent_id) in retained.into_iter().rev() { + slot.queue_mut(priority).push_front(intent_id); + } + } + + let queued_ids = state + .intents + .values() + .filter(|intent| { + intent.team_run_id.as_deref() == Some(team_run_id) && matches!(intent.state, WorkIntentState::Queued) + }) + .map(|intent| intent.intent_id.clone()) + .collect::>(); + for intent_id in queued_ids { + let slot_id = { + let intent = state.intents.get_mut(&intent_id).expect("queued intent exists"); + if let Some(message_id) = &intent.mailbox_message_id { + terminal_message_ids.push(message_id.clone()); + } + intent.state = WorkIntentState::Cancelled { + classification: "run_cancelled", + }; + intent.slot_id.clone() + }; + if let Some(slot) = state.slots.get_mut(&slot_id) { + slot.remove_queued(&intent_id); + } + } + let lease_ids = state + .enqueue_leases + .values() + .filter(|record| record.lease.team_run_id.as_deref() == Some(team_run_id)) + .map(|record| record.lease.lease_id.clone()) + .collect::>(); + for lease_id in lease_ids { + state.enqueue_leases.remove(&lease_id); + } + let summary = Self::run_summary_locked(&state, team_run_id); + drop(state); + self.run_causality.apply_work_summary(summary.clone()); + CancelRunWorkResult { + cancel_targets, + terminal_message_ids, + summary, + } + } + + pub(crate) fn complete_signals(&self, slot_id: &str, intent_ids: &[String]) -> CommitResult { + let mut state = self.lock_state(); + let valid = intent_ids.iter().all(|intent_id| { + state.intents.get(intent_id).is_some_and(|intent| { + intent.slot_id == slot_id + && intent.mailbox_message_id.is_none() + && intent.state == WorkIntentState::Queued + }) + }); + if !valid { + return CommitResult::Rejected; + } + let run_ids = intent_ids + .iter() + .filter_map(|intent_id| state.intents.get(intent_id)?.team_run_id.clone()) + .collect::>(); + for intent_id in intent_ids { + if let Some(intent) = state.intents.get_mut(intent_id) { + intent.state = WorkIntentState::Completed; + } + if let Some(slot) = state.slots.get_mut(slot_id) { + slot.remove_queued(intent_id); + } + } + let summaries = Self::run_summaries_locked(&state, run_ids); + drop(state); + self.publish_run_summaries(summaries); + CommitResult::Committed + } + + pub(crate) fn pause_slot(&self, slot_id: &str) -> PauseWorkResult { + let mut state = self.lock_state(); + let slot = state + .slots + .entry(slot_id.to_owned()) + .or_insert_with(|| SlotState::new(TeamRunTargetRole::Teammate)); + slot.paused = true; + let cancel_target = slot.active.as_ref().map(|active| BatchCancelTarget { + batch: active.batch.clone(), + turn_id: active.turn_id.clone(), + }); + let snapshot = Self::slot_snapshot_locked(&state, slot_id).expect("paused slot exists"); + PauseWorkResult { + cancel_target, + slot: snapshot, + } + } + + pub(crate) fn set_runtime_constraint( + &self, + slot_id: &str, + constraint: RuntimeConstraint, + ) -> RuntimeConstraintUpdate { + let mut state = self.lock_state(); + let slot = state + .slots + .entry(slot_id.to_owned()) + .or_insert_with(|| SlotState::new(TeamRunTargetRole::Teammate)); + slot.runtime_constraint = constraint.clone(); + + let mut terminal_message_ids = Vec::new(); + let mut affected_run_ids = BTreeSet::new(); + if matches!( + constraint, + RuntimeConstraint::Failed { .. } | RuntimeConstraint::SessionStopped + ) { + let intent_ids = state + .intents + .values() + .filter(|intent| intent.slot_id == slot_id && !intent.state.is_terminal()) + .map(|intent| intent.intent_id.clone()) + .collect::>(); + let classification = match constraint { + RuntimeConstraint::Failed { classification, .. } => classification, + RuntimeConstraint::SessionStopped => "session_stopped", + _ => unreachable!(), + }; + let mut requeue_background = Vec::new(); + for intent_id in intent_ids { + let intent = state.intents.get_mut(&intent_id).expect("selected intent exists"); + if intent.team_run_id.is_none() && !matches!(constraint, RuntimeConstraint::SessionStopped) { + intent.state = WorkIntentState::Queued; + requeue_background.push((intent.priority, intent_id)); + continue; + } + if let Some(team_run_id) = &intent.team_run_id { + affected_run_ids.insert(team_run_id.clone()); + } + if let Some(message_id) = &intent.mailbox_message_id { + terminal_message_ids.push(message_id.clone()); + } + intent.state = if matches!(constraint, RuntimeConstraint::SessionStopped) { + WorkIntentState::Cancelled { classification } + } else { + WorkIntentState::Failed { classification } + }; + } + let slot = state.slots.get_mut(slot_id).expect("constrained slot exists"); + slot.foreground.clear(); + slot.control.clear(); + slot.background.clear(); + slot.active = None; + for (priority, intent_id) in requeue_background { + slot.queue_mut(priority).push_back(intent_id); + } + } + let affected_run_summaries = affected_run_ids + .iter() + .map(|run_id| Self::run_summary_locked(&state, run_id)) + .collect::>(); + let slot = Self::slot_snapshot_locked(&state, slot_id).expect("constrained slot exists"); + drop(state); + self.publish_run_summaries(affected_run_summaries.clone()); + RuntimeConstraintUpdate { + slot, + terminal_message_ids, + affected_run_summaries, + } + } + + pub(crate) fn remove_slot(&self, slot_id: &str) -> RemoveWorkResult { + let mut state = self.lock_state(); + let cancel_target = state + .slots + .get(slot_id) + .and_then(|slot| slot.active.as_ref()) + .map(|active| BatchCancelTarget { + batch: active.batch.clone(), + turn_id: active.turn_id.clone(), + }); + let mut intent_ids = state + .slots + .get(slot_id) + .into_iter() + .flat_map(|slot| { + slot.active + .iter() + .flat_map(|active| active.batch.intent_ids.iter()) + .chain(slot.queued_ids()) + }) + .cloned() + .collect::>(); + let already_ordered = intent_ids.iter().cloned().collect::>(); + let mut remaining = state + .intents + .values() + .filter(|intent| { + intent.slot_id == slot_id && !intent.state.is_terminal() && !already_ordered.contains(&intent.intent_id) + }) + .collect::>(); + remaining.sort_by_key(|intent| (intent.created_at_ms, intent.intent_id.as_str())); + intent_ids.extend(remaining.into_iter().map(|intent| intent.intent_id.clone())); + let mut terminal_message_ids = Vec::new(); + let mut affected_run_ids = BTreeSet::new(); + for intent_id in intent_ids { + let intent = state.intents.get_mut(&intent_id).expect("selected intent exists"); + if let Some(message_id) = &intent.mailbox_message_id { + terminal_message_ids.push(message_id.clone()); + } + if let Some(team_run_id) = &intent.team_run_id { + affected_run_ids.insert(team_run_id.clone()); + } + intent.state = WorkIntentState::Cancelled { + classification: "slot_removed", + }; + } + if let Some(slot) = state.slots.get_mut(slot_id) { + slot.foreground.clear(); + slot.control.clear(); + slot.background.clear(); + slot.active = None; + slot.removed = true; + } + let affected_run_summaries = affected_run_ids + .iter() + .map(|run_id| Self::run_summary_locked(&state, run_id)) + .collect::>(); + drop(state); + self.publish_run_summaries(affected_run_summaries.clone()); + RemoveWorkResult { + cancel_target, + terminal_message_ids, + affected_run_summaries, + } + } + + pub(crate) fn stop(&self) -> Vec { + let mut state = self.lock_state(); + let run_ids = Self::all_run_ids(&state); + for intent in state.intents.values_mut().filter(|intent| !intent.state.is_terminal()) { + intent.state = WorkIntentState::Cancelled { + classification: "session_stopped", + }; + } + for slot in state.slots.values_mut() { + slot.foreground.clear(); + slot.control.clear(); + slot.background.clear(); + slot.active = None; + slot.runtime_constraint = RuntimeConstraint::SessionStopped; + } + state.enqueue_leases.clear(); + let summaries = run_ids + .iter() + .map(|run_id| Self::run_summary_locked(&state, run_id)) + .collect::>(); + drop(state); + self.publish_run_summaries(summaries.clone()); + summaries + } + + pub(crate) fn slot_snapshot(&self, slot_id: &str) -> Option { + Self::slot_snapshot_locked(&self.lock_state(), slot_id) + } + + pub(crate) fn snapshot(&self) -> CoordinatorSnapshot { + let state = self.lock_state(); + let slots = state + .slots + .keys() + .filter_map(|slot_id| Self::slot_snapshot_locked(&state, slot_id)) + .collect(); + let active_run_summary = Self::all_run_ids(&state) + .into_iter() + .map(|run_id| Self::run_summary_locked(&state, &run_id)) + .find(|summary| { + summary.queued_intent_count > 0 + || summary.starting_batch_count > 0 + || summary.running_batch_count > 0 + || summary.active_enqueue_lease_count > 0 + || summary.paused_intent_count > 0 + }); + CoordinatorSnapshot { + session_generation: self.session_generation.clone(), + slots, + active_run_summary, + } + } + + #[cfg(test)] + pub(super) fn intents_for_slot(&self, slot_id: &str) -> Vec { + let state = self.lock_state(); + let mut intents = state + .intents + .values() + .filter(|intent| intent.slot_id == slot_id) + .cloned() + .collect::>(); + intents.sort_by_key(|intent| intent.created_at_ms); + intents + } + + fn terminalize_batch( + &self, + batch: &WorkBatch, + terminal_state: WorkIntentState, + classification: &'static str, + ) -> CommitResult { + let mut state = self.lock_state(); + if !self.is_current_batch(&state, batch) { + self.log_stale_batch(batch, "terminalize_batch"); + return CommitResult::StaleOwner; + } + for intent_id in &batch.intent_ids { + if let Some(intent) = state.intents.get_mut(intent_id) { + intent.state = terminal_state.clone(); + } + } + state + .slots + .get_mut(&batch.slot_id) + .expect("current batch slot exists") + .active = None; + let summaries = Self::run_summaries_locked(&state, batch.team_run_ids.iter().cloned()); + drop(state); + self.publish_run_summaries(summaries); + info!( + team_id = %self.team_id, + session_generation = %self.session_generation, + slot_id = %batch.slot_id, + batch_id = %batch.batch_id, + operation_id = batch.operation_id, + classification, + "team work batch terminal" + ); + CommitResult::Committed + } + + fn validate_enqueue_constraint(slot_id: &str, slot: &SlotState) -> Result<(), TeamError> { + if slot.removed { + return Err(TeamError::InvalidRequest(format!("Team slot was removed: {slot_id}"))); + } + match slot.runtime_constraint { + RuntimeConstraint::Ready | RuntimeConstraint::Starting { .. } => Ok(()), + RuntimeConstraint::Failed { classification, .. } => Err(TeamError::InvalidRequest(format!( + "Team slot runtime failed ({classification}): {slot_id}" + ))), + RuntimeConstraint::Removing { .. } => Err(TeamError::InvalidRequest(format!( + "Team slot is being removed: {slot_id}" + ))), + RuntimeConstraint::SessionStopped => Err(TeamError::InvalidRequest(format!( + "Team session stopped for slot: {slot_id}" + ))), + } + } + + fn is_current_batch(&self, state: &CoordinatorState, batch: &WorkBatch) -> bool { + batch.session_generation == self.session_generation + && state + .slots + .get(&batch.slot_id) + .and_then(|slot| slot.active.as_ref()) + .is_some_and(|active| { + active.batch.batch_id == batch.batch_id + && active.batch.operation_id == batch.operation_id + && active.batch.session_generation == batch.session_generation + }) + } + + fn log_stale_batch(&self, batch: &WorkBatch, operation: &'static str) { + warn!( + team_id = %self.team_id, + session_generation = %self.session_generation, + submitted_generation = %batch.session_generation, + slot_id = %batch.slot_id, + batch_id = %batch.batch_id, + operation_id = batch.operation_id, + operation, + "stale team work ownership rejected" + ); + } +} diff --git a/crates/aionui-team/src/work_coordinator/mod.rs b/crates/aionui-team/src/work_coordinator/mod.rs new file mode 100644 index 000000000..81be06690 --- /dev/null +++ b/crates/aionui-team/src/work_coordinator/mod.rs @@ -0,0 +1,9 @@ +mod coordinator; +mod model; +mod snapshot; + +pub(crate) use coordinator::SlotWorkCoordinator; +pub(crate) use model::*; + +#[cfg(test)] +mod tests; diff --git a/crates/aionui-team/src/work_coordinator/model.rs b/crates/aionui-team/src/work_coordinator/model.rs new file mode 100644 index 000000000..f300895a8 --- /dev/null +++ b/crates/aionui-team/src/work_coordinator/model.rs @@ -0,0 +1,240 @@ +use aionui_api_types::TeamRunTargetRole; +use aionui_common::TimestampMs; + +use crate::work_source::WorkSource; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum WorkPriority { + Foreground, + Control, + Background, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum RuntimeConstraint { + Ready, + Starting { + operation_id: u64, + }, + Failed { + operation_id: u64, + classification: &'static str, + }, + Removing { + operation_id: u64, + }, + SessionStopped, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum SlotPhase { + Idle, + Queued, + Starting, + Running, + Paused, + Blocked, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum WorkIntentState { + Queued, + Starting { + batch_id: String, + operation_id: u64, + }, + Running { + batch_id: String, + operation_id: u64, + turn_id: String, + }, + Completed, + Failed { + classification: &'static str, + }, + Cancelled { + classification: &'static str, + }, +} + +impl WorkIntentState { + pub(super) fn is_terminal(&self) -> bool { + matches!(self, Self::Completed | Self::Failed { .. } | Self::Cancelled { .. }) + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct WorkIntent { + pub(crate) intent_id: String, + pub(crate) session_generation: String, + pub(crate) slot_id: String, + pub(crate) role: TeamRunTargetRole, + pub(crate) source: WorkSource, + pub(crate) priority: WorkPriority, + pub(crate) mailbox_message_id: Option, + pub(crate) team_run_id: Option, + pub(crate) created_at_ms: TimestampMs, + pub(crate) state: WorkIntentState, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum CausalBinding { + UserVisible, + InheritRunningBatch { caller_slot_id: String }, + ActiveRunOrBackground, + Background, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct EnqueueRequest { + pub(crate) slot_id: String, + pub(crate) role: TeamRunTargetRole, + pub(crate) source: WorkSource, + pub(crate) binding: CausalBinding, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct EnqueueLease { + pub(crate) lease_id: String, + pub(crate) session_generation: String, + pub(crate) slot_id: String, + pub(crate) role: TeamRunTargetRole, + pub(crate) source: WorkSource, + pub(crate) team_run_id: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct WorkBatch { + pub(crate) batch_id: String, + pub(crate) session_generation: String, + pub(crate) slot_id: String, + pub(crate) intent_ids: Vec, + pub(crate) mailbox_message_ids: Vec, + pub(crate) highest_priority: WorkPriority, + pub(crate) team_run_ids: Vec, + pub(crate) operation_id: u64, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum ReconcileDecision { + Claim(WorkBatch), + WaitingForCompletion, + Blocked(RuntimeConstraint), + SettleSignals(Vec), + Quiescent, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum CommitResult { + Committed, + StaleOwner, + Rejected, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum StartCommitResult { + Accepted, + CancelImmediately, + StaleOwner, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum EnqueueDisposition { + Accepted, + Queued, + BlockedRuntimeStarting, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct EnqueueCommit { + pub(crate) intent_id: String, + pub(crate) team_run_id: Option, + pub(crate) disposition: EnqueueDisposition, + pub(crate) slot: SlotWorkSnapshot, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct SlotWorkSnapshot { + pub(crate) slot_id: String, + pub(crate) role: TeamRunTargetRole, + pub(crate) state: SlotPhase, + pub(crate) queued_foreground_count: usize, + pub(crate) queued_background_count: usize, + pub(crate) active_batch: Option, + pub(crate) active_turn_id: Option, + pub(crate) active_turn_started_at_ms: Option, + pub(crate) runtime_constraint: RuntimeConstraint, + pub(crate) team_run_id: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct RunWorkSummary { + pub(crate) team_run_id: String, + pub(crate) queued_intent_count: usize, + pub(crate) starting_batch_count: usize, + pub(crate) running_batch_count: usize, + pub(crate) active_enqueue_lease_count: usize, + pub(crate) paused_intent_count: usize, + pub(crate) failed_intent_count: usize, + pub(crate) slots: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct CoordinatorSnapshot { + pub(crate) session_generation: String, + pub(crate) slots: Vec, + pub(crate) active_run_summary: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct ReconcileProjection { + pub(crate) created_recovery_intent_ids: Vec, + pub(crate) retained_intent_ids: Vec, + pub(crate) cleared_stale_intent_ids: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct BatchCancelTarget { + pub(crate) batch: WorkBatch, + pub(crate) turn_id: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct PauseWorkResult { + pub(crate) cancel_target: Option, + pub(crate) slot: SlotWorkSnapshot, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct RemoveWorkResult { + pub(crate) cancel_target: Option, + pub(crate) terminal_message_ids: Vec, + pub(crate) affected_run_summaries: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct CancelRunWorkResult { + pub(crate) cancel_targets: Vec, + pub(crate) terminal_message_ids: Vec, + pub(crate) summary: RunWorkSummary, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct RuntimeConstraintUpdate { + pub(crate) slot: SlotWorkSnapshot, + pub(crate) terminal_message_ids: Vec, + pub(crate) affected_run_summaries: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct RunBinding { + pub(crate) team_run_id: Option, + pub(crate) created_new_run: bool, + pub(crate) user_intervention: bool, +} + +pub(crate) trait RunCausalityPort: Send + Sync { + fn bind_enqueue(&self, request: &EnqueueRequest) -> RunBinding; + fn abort_binding(&self, binding: &RunBinding); + fn apply_work_summary(&self, summary: RunWorkSummary); +} diff --git a/crates/aionui-team/src/work_coordinator/snapshot.rs b/crates/aionui-team/src/work_coordinator/snapshot.rs new file mode 100644 index 000000000..aef8aa61a --- /dev/null +++ b/crates/aionui-team/src/work_coordinator/snapshot.rs @@ -0,0 +1,153 @@ +use std::collections::{BTreeSet, HashSet}; +use std::sync::MutexGuard; + +use super::coordinator::{CoordinatorState, SlotWorkCoordinator}; +use super::model::{RunWorkSummary, RuntimeConstraint, SlotPhase, SlotWorkSnapshot, WorkIntentState}; + +impl SlotWorkCoordinator { + pub(super) fn run_summaries_locked( + state: &CoordinatorState, + run_ids: impl IntoIterator, + ) -> Vec { + run_ids + .into_iter() + .collect::>() + .into_iter() + .map(|run_id| Self::run_summary_locked(state, &run_id)) + .collect() + } + + pub(super) fn publish_run_summaries(&self, summaries: Vec) { + for summary in summaries { + self.run_causality.apply_work_summary(summary); + } + } + + pub(super) fn all_run_ids(state: &CoordinatorState) -> BTreeSet { + state + .intents + .values() + .filter_map(|intent| intent.team_run_id.clone()) + .chain( + state + .enqueue_leases + .values() + .filter_map(|record| record.lease.team_run_id.clone()), + ) + .collect() + } + + pub(super) fn run_summary_locked(state: &CoordinatorState, team_run_id: &str) -> RunWorkSummary { + let associated = state + .intents + .values() + .filter(|intent| intent.team_run_id.as_deref() == Some(team_run_id)) + .collect::>(); + let queued_intent_count = associated + .iter() + .filter(|intent| matches!(intent.state, WorkIntentState::Queued)) + .count(); + let starting_batch_count = associated + .iter() + .filter_map(|intent| match &intent.state { + WorkIntentState::Starting { batch_id, .. } => Some(batch_id), + _ => None, + }) + .collect::>() + .len(); + let running_batch_count = associated + .iter() + .filter_map(|intent| match &intent.state { + WorkIntentState::Running { batch_id, .. } => Some(batch_id), + _ => None, + }) + .collect::>() + .len(); + let paused_intent_count = associated + .iter() + .filter(|intent| { + matches!(intent.state, WorkIntentState::Queued) + && state.slots.get(&intent.slot_id).is_some_and(|slot| slot.paused) + }) + .count(); + let failed_intent_count = associated + .iter() + .filter(|intent| matches!(intent.state, WorkIntentState::Failed { .. })) + .count(); + let slot_ids = associated + .iter() + .map(|intent| intent.slot_id.as_str()) + .chain( + state + .enqueue_leases + .values() + .filter(|record| record.lease.team_run_id.as_deref() == Some(team_run_id)) + .map(|record| record.lease.slot_id.as_str()), + ) + .collect::>(); + RunWorkSummary { + team_run_id: team_run_id.to_owned(), + queued_intent_count, + starting_batch_count, + running_batch_count, + active_enqueue_lease_count: state + .enqueue_leases + .values() + .filter(|record| record.lease.team_run_id.as_deref() == Some(team_run_id)) + .count(), + paused_intent_count, + failed_intent_count, + slots: slot_ids + .into_iter() + .filter_map(|slot_id| Self::slot_snapshot_locked(state, slot_id)) + .collect(), + } + } + + pub(super) fn slot_snapshot_locked(state: &CoordinatorState, slot_id: &str) -> Option { + let slot = state.slots.get(slot_id)?; + let active_batch = slot.active.as_ref().map(|active| active.batch.clone()); + let active_turn_id = slot.active.as_ref().and_then(|active| active.turn_id.clone()); + let active_turn_started_at_ms = slot.active.as_ref().and_then(|active| active.started_at_ms); + let state_phase = if slot.paused { + SlotPhase::Paused + } else if active_turn_id.is_some() { + SlotPhase::Running + } else if active_batch.is_some() { + SlotPhase::Starting + } else if !matches!(slot.runtime_constraint, RuntimeConstraint::Ready) { + SlotPhase::Blocked + } else if slot.queued_ids().next().is_some() { + SlotPhase::Queued + } else { + SlotPhase::Idle + }; + let team_run_id = active_batch + .as_ref() + .and_then(|batch| batch.team_run_ids.first().cloned()) + .or_else(|| { + slot.queued_ids().find_map(|intent_id| { + state + .intents + .get(intent_id) + .and_then(|intent| intent.team_run_id.clone()) + }) + }); + Some(SlotWorkSnapshot { + slot_id: slot_id.to_owned(), + role: slot.role.clone(), + state: state_phase, + queued_foreground_count: slot.foreground.len(), + queued_background_count: slot.control.len() + slot.background.len(), + active_batch, + active_turn_id, + active_turn_started_at_ms, + runtime_constraint: slot.runtime_constraint.clone(), + team_run_id, + }) + } + + pub(super) fn lock_state(&self) -> MutexGuard<'_, CoordinatorState> { + self.state.lock().unwrap_or_else(|poisoned| poisoned.into_inner()) + } +} diff --git a/crates/aionui-team/src/work_coordinator/tests.rs b/crates/aionui-team/src/work_coordinator/tests.rs new file mode 100644 index 000000000..8b8c2bc16 --- /dev/null +++ b/crates/aionui-team/src/work_coordinator/tests.rs @@ -0,0 +1,421 @@ +use std::sync::{Arc, Mutex}; + +use aionui_api_types::TeamRunTargetRole; + +use super::*; +use crate::{TeamError, work_source::WorkSource}; + +#[derive(Default)] +struct RecordingRunCausality { + summaries: Mutex>, +} + +impl RunCausalityPort for RecordingRunCausality { + fn bind_enqueue(&self, request: &EnqueueRequest) -> RunBinding { + let user_visible = matches!(request.binding, CausalBinding::UserVisible); + RunBinding { + team_run_id: user_visible.then(|| "run-1".to_owned()), + created_new_run: user_visible, + user_intervention: false, + } + } + + fn abort_binding(&self, _binding: &RunBinding) {} + + fn apply_work_summary(&self, summary: RunWorkSummary) { + self.summaries.lock().unwrap().push(summary); + } +} + +fn coordinator() -> SlotWorkCoordinator { + SlotWorkCoordinator::new( + "team-1".into(), + "generation-1".into(), + Arc::new(RecordingRunCausality::default()), + ) +} + +fn enqueue(coordinator: &SlotWorkCoordinator, source: WorkSource, message_id: &str) { + let lease = coordinator + .acquire_enqueue(EnqueueRequest { + slot_id: "lead-1".into(), + role: TeamRunTargetRole::Lead, + source, + binding: CausalBinding::Background, + }) + .unwrap(); + coordinator.commit_enqueue(&lease, Some(message_id.into())).unwrap(); +} + +#[test] +fn foreground_precedes_background_and_fifo_is_stable() { + let coordinator = coordinator(); + coordinator.set_runtime_constraint("lead-1", RuntimeConstraint::Ready); + enqueue(&coordinator, WorkSource::McpSendMessage, "background-1"); + enqueue(&coordinator, WorkSource::UserMessage, "foreground-1"); + enqueue(&coordinator, WorkSource::UserIntervention, "foreground-2"); + enqueue(&coordinator, WorkSource::TeamMembershipChanged, "background-2"); + + let ReconcileDecision::Claim(foreground) = coordinator.next("lead-1") else { + panic!("foreground batch must be claimable"); + }; + assert_eq!(foreground.highest_priority, WorkPriority::Foreground); + assert_eq!( + foreground.mailbox_message_ids, + vec!["foreground-1".to_owned(), "foreground-2".to_owned()] + ); + assert_eq!(coordinator.complete_batch(&foreground), CommitResult::Committed); + + let ReconcileDecision::Claim(background) = coordinator.next("lead-1") else { + panic!("background batch must follow foreground"); + }; + assert_eq!(background.highest_priority, WorkPriority::Background); + assert_eq!( + background.mailbox_message_ids, + vec!["background-1".to_owned(), "background-2".to_owned()] + ); +} + +#[test] +fn five_enqueues_require_one_reconcile_not_five_signals() { + let coordinator = coordinator(); + coordinator.set_runtime_constraint("lead-1", RuntimeConstraint::Ready); + for index in 1..=5 { + enqueue(&coordinator, WorkSource::UserMessage, &format!("message-{index}")); + } + + loop { + match coordinator.next("lead-1") { + ReconcileDecision::Claim(batch) => { + assert_eq!(coordinator.complete_batch(&batch), CommitResult::Committed); + } + ReconcileDecision::Quiescent => break, + other => panic!("unexpected reconcile decision: {other:?}"), + } + } + + let snapshot = coordinator.slot_snapshot("lead-1").unwrap(); + assert_eq!(snapshot.queued_foreground_count, 0); + assert_eq!(snapshot.queued_background_count, 0); + assert!(snapshot.active_batch.is_none()); + assert_eq!(snapshot.state, SlotPhase::Idle); +} + +#[test] +fn messages_consumed_by_one_turn_are_claimed_in_one_batch() { + let coordinator = coordinator(); + coordinator.set_runtime_constraint("lead-1", RuntimeConstraint::Ready); + for message_id in ["m1", "m2", "m3"] { + enqueue(&coordinator, WorkSource::UserMessage, message_id); + } + coordinator.reconcile_mailbox( + "lead-1", + &["m1".into(), "m2".into(), "m3".into()], + TeamRunTargetRole::Lead, + ); + + let ReconcileDecision::Claim(batch) = coordinator.next("lead-1") else { + panic!("all unread messages must be claimed together"); + }; + assert_eq!(batch.intent_ids.len(), 3); + assert_eq!(batch.mailbox_message_ids, vec!["m1", "m2", "m3"]); + assert_eq!(coordinator.next("lead-1"), ReconcileDecision::WaitingForCompletion); +} + +#[test] +fn enqueue_during_running_batch_waits_for_next_batch() { + let coordinator = coordinator(); + coordinator.set_runtime_constraint("lead-1", RuntimeConstraint::Ready); + enqueue(&coordinator, WorkSource::UserMessage, "m1"); + let ReconcileDecision::Claim(first) = coordinator.next("lead-1") else { + panic!("first message must be claimable"); + }; + assert_eq!(coordinator.mark_started(&first, "turn-1"), StartCommitResult::Accepted); + + enqueue(&coordinator, WorkSource::UserMessage, "m2"); + let active = coordinator.slot_snapshot("lead-1").unwrap(); + assert_eq!(active.active_batch.unwrap().mailbox_message_ids, vec!["m1"]); + + assert_eq!(coordinator.complete_batch(&first), CommitResult::Committed); + let ReconcileDecision::Claim(second) = coordinator.next("lead-1") else { + panic!("second message must follow the active batch"); + }; + assert_eq!(second.mailbox_message_ids, vec!["m2"]); +} + +#[test] +fn retryable_start_returns_the_same_intents_to_the_queue() { + let coordinator = coordinator(); + coordinator.set_runtime_constraint("lead-1", RuntimeConstraint::Ready); + enqueue(&coordinator, WorkSource::UserMessage, "m1"); + let ReconcileDecision::Claim(first) = coordinator.next("lead-1") else { + panic!("message must be claimable"); + }; + + assert_eq!( + coordinator.retry_start(&first, "already_running"), + CommitResult::Committed + ); + let queued = coordinator.intents_for_slot("lead-1"); + assert_eq!(queued[0].intent_id, first.intent_ids[0]); + assert_eq!(queued[0].state, WorkIntentState::Queued); + + let ReconcileDecision::Claim(second) = coordinator.next("lead-1") else { + panic!("retried message must be claimable"); + }; + assert_eq!(second.intent_ids, first.intent_ids); + assert!(second.operation_id > first.operation_id); +} + +#[test] +fn foreground_message_resumes_paused_slot() { + let coordinator = coordinator(); + coordinator.set_runtime_constraint("lead-1", RuntimeConstraint::Ready); + enqueue(&coordinator, WorkSource::McpSendMessage, "background"); + coordinator.pause_slot("lead-1"); + enqueue(&coordinator, WorkSource::UserIntervention, "intervention"); + + let ReconcileDecision::Claim(first) = coordinator.next("lead-1") else { + panic!("foreground intervention must resume the paused slot"); + }; + assert_eq!(first.mailbox_message_ids, vec!["intervention"]); + assert_eq!(coordinator.complete_batch(&first), CommitResult::Committed); + + let ReconcileDecision::Claim(second) = coordinator.next("lead-1") else { + panic!("work retained across pause must remain queued"); + }; + assert_eq!(second.mailbox_message_ids, vec!["background"]); +} + +#[test] +fn cancelled_batch_rejects_a_late_start_by_cancelling_it_immediately() { + let coordinator = coordinator(); + coordinator.set_runtime_constraint("lead-1", RuntimeConstraint::Ready); + enqueue(&coordinator, WorkSource::UserMessage, "m1"); + let ReconcileDecision::Claim(batch) = coordinator.next("lead-1") else { + panic!("message must be claimable"); + }; + + assert_eq!( + coordinator.cancel_batch(&batch, "run_cancelled"), + CommitResult::Committed + ); + assert_eq!( + coordinator.mark_started(&batch, "turn-late"), + StartCommitResult::CancelImmediately + ); +} + +#[test] +fn runtime_starting_blocks_and_runtime_ready_releases_work() { + let coordinator = coordinator(); + coordinator.set_runtime_constraint("lead-1", RuntimeConstraint::Starting { operation_id: 7 }); + enqueue(&coordinator, WorkSource::UserMessage, "m1"); + + assert_eq!( + coordinator.next("lead-1"), + ReconcileDecision::Blocked(RuntimeConstraint::Starting { operation_id: 7 }) + ); + coordinator.set_runtime_constraint("lead-1", RuntimeConstraint::Ready); + let ReconcileDecision::Claim(batch) = coordinator.next("lead-1") else { + panic!("ready runtime must release the queued intent"); + }; + assert_eq!(batch.mailbox_message_ids, vec!["m1"]); +} + +#[test] +fn remove_cancels_queued_and_running_work_and_rejects_new_enqueue() { + let coordinator = coordinator(); + coordinator.set_runtime_constraint("lead-1", RuntimeConstraint::Ready); + enqueue(&coordinator, WorkSource::UserMessage, "m1"); + let ReconcileDecision::Claim(first) = coordinator.next("lead-1") else { + panic!("first message must be claimable"); + }; + assert_eq!(coordinator.mark_started(&first, "turn-1"), StartCommitResult::Accepted); + enqueue(&coordinator, WorkSource::UserMessage, "m2"); + + let removed = coordinator.remove_slot("lead-1"); + assert_eq!(removed.cancel_target.unwrap().batch, first); + assert_eq!(removed.terminal_message_ids, vec!["m1", "m2"]); + assert!(coordinator.intents_for_slot("lead-1").iter().all(|intent| { + intent.state + == WorkIntentState::Cancelled { + classification: "slot_removed", + } + })); + + let result = coordinator.acquire_enqueue(EnqueueRequest { + slot_id: "lead-1".into(), + role: TeamRunTargetRole::Lead, + source: WorkSource::UserMessage, + binding: CausalBinding::UserVisible, + }); + assert!(matches!(result, Err(TeamError::InvalidRequest(message)) if message.contains("removed"))); +} + +#[test] +fn stale_generation_and_operation_cannot_commit() { + let coordinator = coordinator(); + coordinator.set_runtime_constraint("lead-1", RuntimeConstraint::Ready); + enqueue(&coordinator, WorkSource::UserMessage, "m1"); + let ReconcileDecision::Claim(batch) = coordinator.next("lead-1") else { + panic!("message must be claimable"); + }; + let mut stale_generation = batch.clone(); + stale_generation.session_generation = "generation-2".into(); + let mut stale_operation = batch.clone(); + stale_operation.operation_id += 1; + + assert_eq!( + coordinator.mark_started(&stale_generation, "turn-stale"), + StartCommitResult::StaleOwner + ); + assert_eq!( + coordinator.mark_started(&stale_operation, "turn-stale"), + StartCommitResult::StaleOwner + ); + assert_eq!(coordinator.slot_snapshot("lead-1").unwrap().active_batch, Some(batch)); +} + +#[test] +fn unread_without_projection_creates_one_recovery_intent() { + let coordinator = coordinator(); + coordinator.set_runtime_constraint("lead-1", RuntimeConstraint::Ready); + coordinator.reconcile_mailbox("lead-1", &["m1".into()], TeamRunTargetRole::Lead); + coordinator.reconcile_mailbox("lead-1", &["m1".into()], TeamRunTargetRole::Lead); + + let intents = coordinator.intents_for_slot("lead-1"); + assert_eq!(intents.len(), 1); + assert_eq!(intents[0].source, WorkSource::RecoveryDrain); + assert_eq!(intents[0].mailbox_message_id.as_deref(), Some("m1")); +} + +#[test] +fn active_batch_prevents_unread_projection_from_being_rebuilt() { + let coordinator = coordinator(); + coordinator.set_runtime_constraint("lead-1", RuntimeConstraint::Ready); + coordinator.reconcile_mailbox("lead-1", &["m1".into()], TeamRunTargetRole::Lead); + let ReconcileDecision::Claim(batch) = coordinator.next("lead-1") else { + panic!("recovery message must be claimable"); + }; + coordinator.reconcile_mailbox("lead-1", &["m1".into()], TeamRunTargetRole::Lead); + + assert_eq!(coordinator.intents_for_slot("lead-1").len(), 1); + assert_eq!(coordinator.slot_snapshot("lead-1").unwrap().active_batch, Some(batch)); +} + +#[test] +fn pause_cancels_running_batch_and_retains_queued_work() { + let coordinator = coordinator(); + coordinator.set_runtime_constraint("lead-1", RuntimeConstraint::Ready); + enqueue(&coordinator, WorkSource::UserMessage, "running"); + let ReconcileDecision::Claim(running) = coordinator.next("lead-1") else { + panic!("first message must be claimable"); + }; + assert_eq!( + coordinator.mark_started(&running, "turn-1"), + StartCommitResult::Accepted + ); + enqueue(&coordinator, WorkSource::McpSendMessage, "retained"); + + let paused = coordinator.pause_slot("lead-1"); + assert_eq!(paused.cancel_target.unwrap().batch, running); + assert_eq!( + coordinator.cancel_batch(&running, "slot_paused"), + CommitResult::Committed + ); + assert_eq!(coordinator.next("lead-1"), ReconcileDecision::Quiescent); + assert_eq!(coordinator.slot_snapshot("lead-1").unwrap().queued_background_count, 1); +} + +#[test] +fn cancel_run_terminalizes_every_associated_intent_and_lease() { + let coordinator = coordinator(); + coordinator.set_runtime_constraint("lead-1", RuntimeConstraint::Ready); + let first = coordinator + .acquire_enqueue(EnqueueRequest { + slot_id: "lead-1".into(), + role: TeamRunTargetRole::Lead, + source: WorkSource::UserMessage, + binding: CausalBinding::UserVisible, + }) + .unwrap(); + coordinator.commit_enqueue(&first, Some("running".into())).unwrap(); + let ReconcileDecision::Claim(running) = coordinator.next("lead-1") else { + panic!("first message must be claimable"); + }; + let queued = coordinator + .acquire_enqueue(EnqueueRequest { + slot_id: "lead-1".into(), + role: TeamRunTargetRole::Lead, + source: WorkSource::UserIntervention, + binding: CausalBinding::UserVisible, + }) + .unwrap(); + coordinator.commit_enqueue(&queued, Some("queued".into())).unwrap(); + let lease = coordinator + .acquire_enqueue(EnqueueRequest { + slot_id: "lead-1".into(), + role: TeamRunTargetRole::Lead, + source: WorkSource::UserIntervention, + binding: CausalBinding::UserVisible, + }) + .unwrap(); + + let cancelled = coordinator.cancel_run("run-1"); + assert_eq!(cancelled.cancel_targets[0].batch, running); + assert_eq!(cancelled.terminal_message_ids, vec!["running", "queued"]); + assert_eq!(cancelled.summary.active_enqueue_lease_count, 0); + assert_eq!(cancelled.summary.queued_intent_count, 0); + assert_eq!( + coordinator.abort_enqueue(&lease, "late_abort"), + CommitResult::StaleOwner + ); +} + +#[test] +fn background_work_continues_after_unrelated_run_completion() { + let coordinator = coordinator(); + coordinator.set_runtime_constraint("lead-1", RuntimeConstraint::Ready); + let user = coordinator + .acquire_enqueue(EnqueueRequest { + slot_id: "lead-1".into(), + role: TeamRunTargetRole::Lead, + source: WorkSource::UserMessage, + binding: CausalBinding::UserVisible, + }) + .unwrap(); + coordinator.commit_enqueue(&user, Some("user".into())).unwrap(); + enqueue(&coordinator, WorkSource::McpSendMessage, "background"); + + let ReconcileDecision::Claim(user_batch) = coordinator.next("lead-1") else { + panic!("user work must run first"); + }; + coordinator.complete_batch(&user_batch); + let ReconcileDecision::Claim(background_batch) = coordinator.next("lead-1") else { + panic!("background work must remain runnable"); + }; + assert_eq!(background_batch.mailbox_message_ids, vec!["background"]); + assert!(background_batch.team_run_ids.is_empty()); +} + +#[test] +fn late_terminal_after_cancel_is_rejected() { + let coordinator = coordinator(); + coordinator.set_runtime_constraint("lead-1", RuntimeConstraint::Ready); + let user = coordinator + .acquire_enqueue(EnqueueRequest { + slot_id: "lead-1".into(), + role: TeamRunTargetRole::Lead, + source: WorkSource::UserMessage, + binding: CausalBinding::UserVisible, + }) + .unwrap(); + coordinator.commit_enqueue(&user, Some("m1".into())).unwrap(); + let ReconcileDecision::Claim(batch) = coordinator.next("lead-1") else { + panic!("user work must be claimable"); + }; + coordinator.cancel_run("run-1"); + + assert_eq!(coordinator.complete_batch(&batch), CommitResult::StaleOwner); +} diff --git a/crates/aionui-team/src/work_source.rs b/crates/aionui-team/src/work_source.rs new file mode 100644 index 000000000..739bda785 --- /dev/null +++ b/crates/aionui-team/src/work_source.rs @@ -0,0 +1,71 @@ +use crate::work_coordinator::WorkPriority; +use std::fmt; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum WorkSource { + UserMessage, + UserIntervention, + McpSendMessage, + McpShutdownRequest, + SpawnWelcome, + TeamMembershipChanged, + SpawnAttachFailure, + IdleNotification, + InterruptedNotification, + ShutdownRejected, + RecoveryDrain, +} + +impl WorkSource { + pub(crate) fn priority(self) -> WorkPriority { + match self { + Self::UserMessage | Self::UserIntervention => WorkPriority::Foreground, + Self::McpShutdownRequest | Self::ShutdownRejected => WorkPriority::Control, + Self::McpSendMessage + | Self::SpawnWelcome + | Self::TeamMembershipChanged + | Self::SpawnAttachFailure + | Self::IdleNotification + | Self::InterruptedNotification + | Self::RecoveryDrain => WorkPriority::Background, + } + } + + pub(crate) fn resumes_paused_slot(self) -> bool { + matches!(self, Self::UserMessage | Self::UserIntervention) + } + + pub(crate) fn requires_mailbox_message(self) -> bool { + matches!( + self, + Self::UserMessage + | Self::UserIntervention + | Self::McpSendMessage + | Self::McpShutdownRequest + | Self::SpawnWelcome + | Self::SpawnAttachFailure + | Self::InterruptedNotification + | Self::ShutdownRejected + | Self::RecoveryDrain + ) + } +} + +impl fmt::Display for WorkSource { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + let value = match self { + Self::UserMessage => "user_message", + Self::UserIntervention => "user_intervention", + Self::McpSendMessage => "mcp_send_message", + Self::McpShutdownRequest => "mcp_shutdown_request", + Self::SpawnWelcome => "spawn_welcome", + Self::TeamMembershipChanged => "team_membership_changed", + Self::SpawnAttachFailure => "spawn_attach_failure", + Self::IdleNotification => "idle_notification", + Self::InterruptedNotification => "interrupted_notification", + Self::ShutdownRejected => "shutdown_rejected", + Self::RecoveryDrain => "recovery_drain", + }; + formatter.write_str(value) + } +} diff --git a/crates/aionui-team/tests/e2e_team_flow.rs b/crates/aionui-team/tests/e2e_team_flow.rs index c6aff8d7c..670ed5fb9 100644 --- a/crates/aionui-team/tests/e2e_team_flow.rs +++ b/crates/aionui-team/tests/e2e_team_flow.rs @@ -29,8 +29,8 @@ use aionui_ai_agent::agent_task::{AgentInstance, IAgentTask, IMockAgent}; use aionui_ai_agent::protocol::events::{AgentStreamEvent, FinishEventData}; use aionui_ai_agent::shared_kernel::approval_key; use aionui_ai_agent::types::{BuildTaskOptions, SendMessageData}; +use aionui_api_types::AgentModeResponse; use aionui_api_types::WebSocketMessage; -use aionui_api_types::{AgentModeResponse, TeamRunTargetRole}; use aionui_common::{AgentKillReason, AgentType, Confirmation, ConversationStatus, TimestampMs, now_ms}; use aionui_db::ITeamRepository; use aionui_realtime::EventBroadcaster; @@ -79,7 +79,7 @@ impl AgentTurnExecutionPort for RecordingTurnPort { let turn_id = format!("turn-test-{turn_index}"); if let Some(on_started) = request.on_started.as_ref() { on_started(AgentTurnStarted { - team_run_id: request.team_run_id.clone().expect("team run id"), + team_run_id: request.team_run_id.clone(), slot_id: request.slot_id.clone(), role: request.role.clone(), conversation_id: request.conversation_id.clone(), @@ -149,7 +149,7 @@ impl AgentTurnExecutionPort for StartedThenFailedTurnPort { async fn run_agent_turn(&self, request: AgentTurnRequest) -> Result { if let Some(on_started) = request.on_started.as_ref() { on_started(AgentTurnStarted { - team_run_id: request.team_run_id.clone().expect("team run id"), + team_run_id: request.team_run_id.clone(), slot_id: request.slot_id.clone(), role: request.role.clone(), conversation_id: request.conversation_id.clone(), @@ -171,6 +171,63 @@ struct BlockingStartTurnPort { release_rx: Mutex>>, } +struct HoldFirstRunningTurnPort { + requests: Arc>>, + first_started_tx: Mutex>>, + first_release_rx: Mutex>>, +} + +impl HoldFirstRunningTurnPort { + fn new(first_started_tx: oneshot::Sender<()>, first_release_rx: oneshot::Receiver<()>) -> Self { + Self { + requests: Arc::new(Mutex::new(Vec::new())), + first_started_tx: Mutex::new(Some(first_started_tx)), + first_release_rx: Mutex::new(Some(first_release_rx)), + } + } + + fn requests(&self) -> Arc>> { + self.requests.clone() + } +} + +#[async_trait] +impl AgentTurnExecutionPort for HoldFirstRunningTurnPort { + async fn run_agent_turn(&self, request: AgentTurnRequest) -> Result { + let turn_index = { + let mut requests = self.requests.lock().unwrap(); + requests.push(request.clone()); + requests.len() + }; + let turn_id = format!("turn-coalesced-{turn_index}"); + if let Some(on_started) = request.on_started.as_ref() { + on_started(AgentTurnStarted { + team_run_id: request.team_run_id.clone(), + slot_id: request.slot_id.clone(), + role: request.role.clone(), + conversation_id: request.conversation_id.clone(), + turn_id: turn_id.clone(), + }) + .await; + } + if turn_index == 1 { + if let Some(tx) = self.first_started_tx.lock().unwrap().take() { + let _ = tx.send(()); + } + let release_rx = self.first_release_rx.lock().unwrap().take(); + if let Some(rx) = release_rx { + let _ = rx.await; + } + } + Ok(AgentTurnOutcome { + conversation_id: request.conversation_id, + turn_id, + status: AgentTurnStatus::Completed, + runtime: None, + }) + } +} + impl BlockingStartTurnPort { fn new(claimed_tx: oneshot::Sender<()>, release_rx: oneshot::Receiver<()>) -> Self { Self { @@ -192,7 +249,7 @@ impl AgentTurnExecutionPort for BlockingStartTurnPort { } if let Some(on_started) = request.on_started.as_ref() { on_started(AgentTurnStarted { - team_run_id: request.team_run_id.clone().expect("team run id"), + team_run_id: request.team_run_id.clone(), slot_id: request.slot_id.clone(), role: request.role.clone(), conversation_id: request.conversation_id.clone(), @@ -729,6 +786,16 @@ async fn wait_for_turn_request_count(requests: &Arc> panic!("timed out waiting for {count} turn requests"); } +async fn wait_for_cancelled_turn_count(cancelled: &Arc>>, count: usize) { + for _ in 0..100 { + if cancelled.lock().unwrap().len() >= count { + return; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + panic!("timed out waiting for {count} cancelled turns"); +} + fn register_test_event_loops(session: &Arc) { let registry = session.event_loops().clone(); for agent in two_agents() { @@ -742,7 +809,7 @@ fn register_test_event_loops(session: &Arc) { turn_port: session.turn_port().clone(), registry: registry.clone(), }; - registry.spawn(&agent.slot_id, ctx); + registry.spawn(&agent.slot_id, ctx).expect("register test event loop"); } } @@ -765,7 +832,7 @@ async fn wait_until_turn_count(turn_requests: &Arc>> /// Verifies: /// - TeamSession::start succeeds /// - MCP TCP server is reachable -/// - tools/list returns all 11 expected tools +/// - tools/list returns all expected tools #[tokio::test] async fn s1a_mcp_server_starts_and_tools_available() { let (session, _tm, _repo, _sent) = setup_session().await; @@ -781,13 +848,17 @@ async fn s1a_mcp_server_starts_and_tools_available() { let resp = tcp_recv(&mut stream).await; let tools = resp["result"]["tools"].as_array().expect("tools array"); - assert_eq!(tools.len(), 11, "expected exactly 11 MCP tools, got {}", tools.len()); + assert_eq!(tools.len(), 10, "expected exactly 10 MCP tools, got {}", tools.len()); let names: Vec<&str> = tools.iter().map(|t| t["name"].as_str().unwrap()).collect(); assert!(names.contains(&"team_send_message"), "missing team_send_message"); assert!(names.contains(&"team_members"), "missing team_members"); assert!(names.contains(&"team_task_create"), "missing team_task_create"); assert!(names.contains(&"team_list_assistants"), "missing team_list_assistants"); + assert!( + !names.contains(&"team_list_models"), + "team_list_models should not be exposed" + ); session.stop(); } @@ -1092,19 +1163,21 @@ async fn s4_dynamic_agent_added_then_finish_propagates() { // Add the agent to the session's scheduler session.add_agent(&new_agent).await; let registry = session.event_loops().clone(); - registry.spawn( - &new_agent.slot_id, - AgentLoopContext { - team_id: session.team_id().to_owned(), - slot_id: new_agent.slot_id.clone(), - user_id: session.user_id().to_owned(), - session: session.clone(), - scheduler: session.scheduler().clone(), - mailbox: session.mailbox().clone(), - turn_port: session.turn_port().clone(), - registry: registry.clone(), - }, - ); + registry + .spawn( + &new_agent.slot_id, + AgentLoopContext { + team_id: session.team_id().to_owned(), + slot_id: new_agent.slot_id.clone(), + user_id: session.user_id().to_owned(), + session: session.clone(), + scheduler: session.scheduler().clone(), + mailbox: session.mailbox().clone(), + turn_port: session.turn_port().clone(), + registry: registry.clone(), + }, + ) + .expect("register test event loop"); // Verify the agent is in the roster let agents = session.scheduler().list_agents().await; @@ -1144,15 +1217,8 @@ async fn s4_dynamic_agent_added_then_finish_propagates() { ); drop(log); - // Simulate helper finishing → IdleNotification → lead should be wake target - let wake_target = session.on_agent_finish("conv-helper", false).await.unwrap(); - assert_eq!( - wake_target.as_deref(), - Some("lead-1"), - "helper finish must return lead as wake target" - ); - - // Lead mailbox should have IdleNotification from helper + // EventLoop finalization owns the finish transition and writes the leader + // notification; calling on_agent_finish here would duplicate that terminal. let state = repo.state.lock().unwrap(); let lead_notifs: Vec<_> = state .messages @@ -1168,83 +1234,6 @@ async fn s4_dynamic_agent_added_then_finish_propagates() { session.stop(); } -#[tokio::test] -async fn s4b_pending_wake_for_unregistered_dynamic_agent_survives_leader_empty_wake() { - let (session, task_manager, _repo, sent, turn_requests) = setup_session_with_turn_recorder().await; - - let ack = session - .team_run_manager() - .accept_user_message("lead-1", TeamRunTargetRole::Lead, false, None) - .await - .expect("accept active Team Run"); - - let new_agent = TeamAgent { - slot_id: "helper-1".into(), - name: "Helper".into(), - role: TeammateRole::Teammate, - conversation_id: "conv-helper".into(), - backend: "acp".into(), - model: "claude".into(), - assistant_id: None, - status: None, - conversation_type: None, - cli_path: None, - }; - - let handle = AgentInstance::Mock(Arc::new(RecordingAgent::new("conv-helper", sent.clone()))); - task_manager.insert("conv-helper", handle); - session.add_agent(&new_agent).await; - - session - .send_message_to_agent("helper-1", "start your task", None) - .await - .expect("send to unregistered helper should record pending wake"); - - assert!( - session - .team_run_manager() - .record_empty_wake_observed("lead-1") - .await - .is_none(), - "leader empty wake must not complete while helper pending wake is retained" - ); - assert_eq!( - session.team_run_manager().active_run_id().await.as_deref(), - Some(ack.team_run_id.as_str()) - ); - - let registry = session.event_loops().clone(); - registry.spawn( - &new_agent.slot_id, - AgentLoopContext { - team_id: session.team_id().to_owned(), - slot_id: new_agent.slot_id.clone(), - user_id: session.user_id().to_owned(), - session: session.clone(), - scheduler: session.scheduler().clone(), - mailbox: session.mailbox().clone(), - turn_port: session.turn_port().clone(), - registry: registry.clone(), - }, - ); - registry.notify("helper-1"); - - tokio::time::sleep(std::time::Duration::from_millis(100)).await; - - let log = turn_requests.lock().unwrap(); - assert!( - log.iter().any(|request| { - request.slot_id == "helper-1" - && request.conversation_id == "conv-helper" - && request.team_id == "e2e-team" - && request.team_run_id.as_deref() == Some(ack.team_run_id.as_str()) - }), - "helper turn port must be called after delayed registration; requests: {log:?}" - ); - - session.stop(); -} - // =========================================================================== // Scenario 5: Rapid consecutive messages – dedup window after clear // =========================================================================== @@ -1518,12 +1507,12 @@ async fn s8b_worker_cannot_call_spawn_agent() { // =========================================================================== #[tokio::test] -async fn s8_paused_slot_user_intervention_runs_first_then_releases_background_backlog() { +async fn turn_completion_reconciles_without_another_notify() { let (session, _task_manager, _repo, _sent, turn_requests) = setup_session_with_turn_recorder_without_loops().await; let ack = session.send_message("start team", None).await.unwrap(); session - .pause_slot_work(&ack.team_run_id, "lead-1", Some("user stopped".into())) + .pause_slot_work(&ack.run.team_run_id, "lead-1", Some("user stopped".into())) .await .unwrap(); session @@ -1548,6 +1537,7 @@ async fn s8_paused_slot_user_intervention_runs_first_then_releases_background_ba wait_until_turn_count(&turn_requests, 1).await; let first = turn_requests.lock().unwrap()[0].clone(); + assert!(first.content.contains("start team")); assert!(first.content.contains("user priority")); assert!(!first.content.contains("background backlog")); match first.source { @@ -1555,8 +1545,8 @@ async fn s8_paused_slot_user_intervention_runs_first_then_releases_background_ba unread_message_ids, unread_count, } => { - assert_eq!(unread_count, 1); - assert_eq!(unread_message_ids, vec![user_ack.message_id.unwrap()]); + assert_eq!(unread_count, 2); + assert_eq!(unread_message_ids, vec![ack.message_id, user_ack.message_id]); } } @@ -1605,7 +1595,88 @@ async fn s9_session_send_message_wakes_lead() { } #[tokio::test] -async fn s9b_event_loop_fails_run_when_turn_errors_before_started() { +async fn one_notify_drains_five_intents_after_busy_turn() { + let (first_started_tx, first_started_rx) = oneshot::channel(); + let (first_release_tx, first_release_rx) = oneshot::channel(); + let broadcaster = Arc::new(RecordingBroadcaster::new()); + let turn_port = Arc::new(HoldFirstRunningTurnPort::new(first_started_tx, first_release_rx)); + let requests = turn_port.requests(); + let session = + setup_session_with_runtime_ports(turn_port, Arc::new(NoopCancellationPort), broadcaster.clone()).await; + + let first = session.send_message("first", None).await.unwrap(); + tokio::time::timeout(Duration::from_secs(2), first_started_rx) + .await + .expect("first turn should start") + .expect("first start signal should be sent"); + + let mut queued_message_ids = Vec::new(); + for index in 1..=5 { + let ack = session.send_message(&format!("queued-{index}"), None).await.unwrap(); + assert_eq!(ack.enqueue_status, aionui_api_types::TeamMessageEnqueueStatus::Queued); + assert_eq!(ack.run.team_run_id, first.run.team_run_id); + queued_message_ids.push(ack.message_id); + } + + first_release_tx.send(()).unwrap(); + wait_for_event(&broadcaster, "team.runCompleted").await; + wait_for_turn_request_count(&requests, 2).await; + + let requests = requests.lock().unwrap(); + assert_eq!(requests.len(), 2); + match &requests[1].source { + AgentTurnSource::Mailbox { + unread_message_ids, + unread_count, + } => { + assert_eq!(*unread_count, 5); + assert_eq!(unread_message_ids, &queued_message_ids); + } + } + drop(requests); + assert!( + session + .mailbox() + .peek_unread("e2e-team", "lead-1") + .await + .unwrap() + .is_empty() + ); + assert!(session.team_run_manager().current_active_run_id().is_none()); + session.stop(); +} + +#[tokio::test] +async fn one_turn_claims_every_message_it_reads() { + let (session, _task_manager, _repo, _sent, requests) = setup_session_with_turn_recorder_without_loops().await; + let first = session.send_message("batch-1", None).await.unwrap(); + let second = session.send_message("batch-2", None).await.unwrap(); + let third = session.send_message("batch-3", None).await.unwrap(); + register_test_event_loops(&session); + + session.event_loops().notify("lead-1"); + wait_for_turn_request_count(&requests, 1).await; + + let requests = requests.lock().unwrap(); + assert_eq!(requests.len(), 1); + match &requests[0].source { + AgentTurnSource::Mailbox { + unread_message_ids, + unread_count, + } => { + assert_eq!(*unread_count, 3); + assert_eq!( + unread_message_ids, + &vec![first.message_id, second.message_id, third.message_id] + ); + } + } + drop(requests); + session.stop(); +} + +#[tokio::test] +async fn nonretryable_start_failure_is_terminal_and_not_recovered() { let broadcaster = Arc::new(RecordingBroadcaster::new()); let session = setup_session_with_runtime_ports( Arc::new(ErrorBeforeStartTurnPort), @@ -1620,13 +1691,21 @@ async fn s9b_event_loop_fails_run_when_turn_errors_before_started() { .expect("send_message must succeed"); wait_for_event(&broadcaster, "team.runFailed").await; - assert_eq!(session.team_run_manager().active_run_id().await, None); + assert_eq!(session.team_run_manager().current_active_run_id(), None); + assert!( + session + .mailbox() + .peek_unread("e2e-team", "lead-1") + .await + .unwrap() + .is_empty() + ); session.stop(); } #[tokio::test] -async fn s9b_retryable_busy_before_start_retains_team_run_without_timer_retry() { +async fn retryable_start_requeues_without_marking_mailbox_read() { let broadcaster = Arc::new(RecordingBroadcaster::new()); let turn_port = Arc::new(SkippedBusyTurnPort::default()); let requests = turn_port.requests(); @@ -1654,18 +1733,14 @@ async fn s9b_retryable_busy_before_start_retains_team_run_without_timer_retry() drop(request_log); assert_eq!( - session.team_run_manager().active_run_id().await, + session.team_run_manager().current_active_run_id(), Some(first_team_run_id.clone()), "run should remain active for a state-driven retry" ); - let payload = session - .team_run_manager() - .current_payload() - .await - .expect("retryable busy skip should keep the active run"); - assert_eq!(payload.team_run_id, first_team_run_id); - assert_eq!(payload.pending_wake_count, 1); - assert_eq!(payload.starting_child_count, 0); + assert_eq!( + session.mailbox().peek_unread("e2e-team", "lead-1").await.unwrap().len(), + 1 + ); session.stop(); } @@ -1685,9 +1760,8 @@ async fn s9c_event_loop_fails_run_when_started_turn_returns_failed() { .await .expect("send_message must succeed"); - wait_for_event(&broadcaster, "team.childTurnCompleted").await; wait_for_event(&broadcaster, "team.runFailed").await; - assert_eq!(session.team_run_manager().active_run_id().await, None); + assert_eq!(session.team_run_manager().current_active_run_id(), None); session.stop(); } @@ -1716,15 +1790,16 @@ async fn s9d_late_child_start_after_team_cancel_is_cancelled_without_reviving_ru .expect("claim signal should be sent"); session - .cancel_run(&ack.team_run_id, None, Some("stop".into())) + .cancel_run(&ack.run.team_run_id, None, Some("stop".into())) .await .expect("team run cancel must succeed"); release_tx.send(()).expect("release signal should be sent"); wait_for_event(&broadcaster, "team.runCancelled").await; + wait_for_cancelled_turn_count(&cancelled, 1).await; let cancelled_turns = cancelled.lock().unwrap().clone(); assert_eq!(cancelled_turns, vec!["turn-late-start".to_owned()]); - assert_eq!(session.team_run_manager().active_run_id().await, None); + assert_eq!(session.team_run_manager().current_active_run_id(), None); session.stop(); } @@ -1807,71 +1882,3 @@ async fn s11_shutdown_approved_interception() { session.stop(); } - -// =========================================================================== -// Scenario 12: compute_wake_input — role prompt injection on cold start -// =========================================================================== - -/// Scenario 12a: Cold-start lead gets role prompt injected into first_message. -#[tokio::test] -async fn s12a_cold_start_lead_gets_role_prompt() { - let (session, _tm, _repo, _sent, _turn_requests) = setup_session_with_turn_recorder_without_loops().await; - session.send_message("kick off", None).await.unwrap(); - - let input = session - .compute_wake_input("lead-1") - .await - .unwrap() - .expect("WakeInput must be Some"); - - assert!(input.should_send, "should_send must be true when mailbox has messages"); - assert!( - input.first_message.contains("You are the Team Leader"), - "cold-start lead must get role prompt; got: {}", - &input.first_message[..input.first_message.len().min(200)] - ); - assert!(input.first_message.contains("kick off")); - - session.stop(); -} - -/// Scenario 12b: Warm lead (role prompt already consumed) does not get role prompt again. -/// -/// The cold-start flag is a one-shot: `take_needs_role_prompt` returns true -/// only on the first call. To simulate a "warm" agent, we consume the flag -/// on a first compute_wake_input call, then assert the second call omits it. -#[tokio::test] -async fn s12b_warm_lead_skips_role_prompt() { - let (session, _tm, _repo, _sent, _turn_requests) = setup_session_with_turn_recorder_without_loops().await; - - // First: consume the cold-start role-prompt flag from a TeamRun-owned wake. - session.send_message("initial kick", None).await.unwrap(); - let first_input = session - .compute_wake_input("lead-1") - .await - .unwrap() - .expect("first WakeInput"); - assert!( - first_input.first_message.contains("You are the Team Leader"), - "first (cold) compute_wake_input must include role prompt" - ); - // The flag is now consumed (take_needs_role_prompt returned true, set to false). - - // Second call: write another message and compute again — no role prompt this time. - session.send_message("follow-up message", None).await.unwrap(); - - let input = session - .compute_wake_input("lead-1") - .await - .unwrap() - .expect("second WakeInput must be Some"); - - assert!( - !input.first_message.contains("You are the Team Leader"), - "warm lead (flag consumed) must NOT get role prompt on second call; got: {}", - &input.first_message[..input.first_message.len().min(300)] - ); - assert!(input.first_message.contains("follow-up message")); - - session.stop(); -} diff --git a/crates/aionui-team/tests/mcp_server_integration.rs b/crates/aionui-team/tests/mcp_server_integration.rs index 927c0a51b..2e6892ef7 100644 --- a/crates/aionui-team/tests/mcp_server_integration.rs +++ b/crates/aionui-team/tests/mcp_server_integration.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use aionui_api_types::WebSocketMessage; use aionui_realtime::EventBroadcaster; use aionui_team::mcp::protocol::{read_frame, write_frame}; -use aionui_team::{Mailbox, TaskBoard, TeamAgent, TeamMcpServer, TeammateManager, TeammateRole}; +use aionui_team::{Mailbox, TaskBoard, TeamAgent, TeamMcpServer, TeamPromptDumpConfig, TeammateManager, TeammateRole}; use common::MockTeamRepo; use serde_json::{Value, json}; use tokio::net::TcpStream; @@ -24,10 +24,6 @@ impl RecordingBroadcaster { events: std::sync::Mutex::new(vec![]), } } - - fn events(&self) -> Vec> { - self.events.lock().unwrap().clone() - } } impl EventBroadcaster for RecordingBroadcaster { @@ -49,7 +45,7 @@ fn make_agents() -> Vec { conversation_id: "conv-lead".into(), backend: "acp".into(), model: "claude".into(), - assistant_id: None, + assistant_id: Some("lead-assistant".into()), status: None, conversation_type: None, cli_path: None, @@ -61,7 +57,7 @@ fn make_agents() -> Vec { conversation_id: "conv-worker".into(), backend: "acp".into(), model: "claude".into(), - assistant_id: None, + assistant_id: Some("worker-assistant".into()), status: None, conversation_type: None, cli_path: None, @@ -72,10 +68,13 @@ fn make_agents() -> Vec { struct TestEnv { server: TeamMcpServer, _repo: Arc, - broadcaster: Arc, } async fn setup() -> TestEnv { + setup_with_prompt_dump(None).await +} + +async fn setup_with_prompt_dump(prompt_dump: Option) -> TestEnv { let repo = Arc::new(MockTeamRepo::new()); let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo.clone())); @@ -94,21 +93,18 @@ async fn setup() -> TestEnv { // the Weak cannot upgrade, so `team_spawn_agent` will surface the // service-unavailable error. Non-spawn tools still exercise scheduler // flows directly and do not hit this path. - let server = TeamMcpServer::start( + let server = TeamMcpServer::start_with_prompt_dump( "test-token-123".into(), scheduler, "team-1".into(), broadcaster, std::sync::Weak::new(), + prompt_dump, ) .await .unwrap(); - TestEnv { - server, - _repo: repo, - broadcaster: recorder, - } + TestEnv { server, _repo: repo } } async fn connect_and_init(port: u16, token: &str, slot_id: &str) -> TcpStream { @@ -206,6 +202,34 @@ fn is_error_response(resp: &Value) -> bool { resp["result"]["isError"].as_bool().unwrap_or(false) } +async fn create_task(stream: &mut TcpStream, id: u64, subject: &str, owner: Option<&str>) -> String { + let mut args = json!({ "subject": subject }); + if let Some(owner) = owner { + args["owner"] = json!(owner); + } + let resp = call_tool(stream, id, "team_task_create", args).await; + assert!(!is_error_response(&resp), "team_task_create failed: {resp}"); + let payload: Value = serde_json::from_str(&extract_text(&resp)).unwrap(); + payload["task"]["task_id"].as_str().unwrap().to_owned() +} + +async fn update_task_status(stream: &mut TcpStream, id: u64, task_id: &str, status: &str) { + let resp = call_tool( + stream, + id, + "team_task_update", + json!({ "task_id": task_id, "status": status }), + ) + .await; + assert!(!is_error_response(&resp), "team_task_update failed: {resp}"); +} + +async fn list_tasks_with_args(stream: &mut TcpStream, id: u64, args: Value) -> Vec { + let resp = call_tool(stream, id, "team_task_list", args).await; + assert!(!is_error_response(&resp), "team_task_list failed: {resp}"); + serde_json::from_str(&extract_text(&resp)).unwrap() +} + // --------------------------------------------------------------------------- // Tests: Connection & Authentication (MC-1, MC-2, MC-3) // --------------------------------------------------------------------------- @@ -223,7 +247,9 @@ async fn mc1_correct_token_connects() { send_request(&mut stream, &req).await; let resp = read_response(&mut stream).await; let tools = resp["result"]["tools"].as_array().unwrap(); - assert_eq!(tools.len(), 11); + assert_eq!(tools.len(), 10); + let names: Vec<&str> = tools.iter().filter_map(|tool| tool["name"].as_str()).collect(); + assert!(!names.contains(&"team_list_models")); env.server.stop(); } @@ -283,12 +309,12 @@ async fn mc3_no_token_rejected() { // --------------------------------------------------------------------------- #[tokio::test] -async fn tools_list_returns_all_11_tools() { +async fn tools_list_returns_all_10_tools() { let env = setup().await; let mut stream = connect_and_init(env.server.port(), "test-token-123", "lead-1").await; let names = list_tools(&mut stream, 10).await; - assert_eq!(names.len(), 11); + assert_eq!(names.len(), 10); assert!(names.contains(&"team_send_message".to_owned())); assert!(names.contains(&"team_spawn_agent".to_owned())); @@ -299,8 +325,8 @@ async fn tools_list_returns_all_11_tools() { assert!(names.contains(&"team_rename_agent".to_owned())); assert!(names.contains(&"team_shutdown_agent".to_owned())); assert!(names.contains(&"team_list_assistants".to_owned())); - assert!(names.contains(&"team_list_models".to_owned())); assert!(names.contains(&"team_describe_assistant".to_owned())); + assert!(!names.contains(&"team_list_models".to_owned())); env.server.stop(); } @@ -321,8 +347,8 @@ async fn mcp_tools_list_filters_lead_only_tools() { assert!(names.contains(&"team_task_list".to_owned())); assert!(names.contains(&"team_members".to_owned())); assert!(names.contains(&"team_list_assistants".to_owned())); - assert!(names.contains(&"team_list_models".to_owned())); assert!(names.contains(&"team_describe_assistant".to_owned())); + assert!(!names.contains(&"team_list_models".to_owned())); env.server.stop(); } @@ -386,7 +412,27 @@ async fn ts3_send_message_to_nonexistent_agent() { assert!(is_error_response(&resp)); let text = extract_text(&resp); - assert!(text.contains("No agent matches 'nonexistent'")); + assert!(text.contains("expected slot_id or \"*\""), "unexpected error: {text}"); + + env.server.stop(); +} + +#[tokio::test] +async fn team_send_message_rejects_display_name_target() { + let env = setup().await; + let mut stream = connect_and_init(env.server.port(), "test-token-123", "lead-1").await; + + let resp = call_tool( + &mut stream, + 2, + "team_send_message", + json!({"to": "Worker", "message": "Hello by name"}), + ) + .await; + + assert!(is_error_response(&resp)); + let text = extract_text(&resp); + assert!(text.contains("expected slot_id or \"*\""), "unexpected error: {text}"); env.server.stop(); } @@ -427,7 +473,10 @@ async fn team_send_message_shutdown_rejected_intercepted() { assert!(!is_error_response(&resp)); let text = extract_text(&resp); - assert_eq!(text, "shutdown_rejected: still finishing task"); + let payload: Value = serde_json::from_str(&text).expect("shutdown_rejected response must be JSON"); + assert_eq!(payload["status"], "ok"); + assert_eq!(payload["action"], "shutdown_rejected"); + assert_eq!(payload["reason"], "still finishing task"); env.server.stop(); } @@ -472,7 +521,7 @@ async fn sp1_lead_spawn_requires_live_session_service() { &mut stream, 2, "team_spawn_agent", - json!({"name": "Helper", "role": "worker", "assistant_id": "word-creator"}), + json!({"name": "Helper", "assistant_id": "word-creator"}), ) .await; @@ -548,7 +597,10 @@ async fn ttc1_create_basic_task() { assert!(!is_error_response(&resp)); let text = extract_text(&resp); - assert!(text.contains("Implement feature X")); + let payload: Value = serde_json::from_str(&text).expect("team_task_create must return JSON"); + assert_eq!(payload["status"], "ok"); + assert_eq!(payload["task"]["subject"], "Implement feature X"); + assert!(payload["task"]["task_id"].as_str().is_some()); env.server.stop(); } @@ -573,6 +625,10 @@ async fn ttc2_create_task_with_dependency() { .await; assert!(!is_error_response(&resp)); + let text = extract_text(&resp); + let payload: Value = serde_json::from_str(&text).expect("team_task_create must return JSON"); + assert_eq!(payload["status"], "ok"); + assert_eq!(payload["task"]["subject"], "Task B"); let list_resp2 = call_tool(&mut stream, 5, "team_task_list", json!({})).await; let tasks2: Vec = serde_json::from_str(&extract_text(&list_resp2)).unwrap(); @@ -616,6 +672,137 @@ async fn ttl1_task_list_after_create() { env.server.stop(); } +#[tokio::test] +async fn ttl3_task_list_empty_args_still_returns_full_board() { + let env = setup().await; + let mut stream = connect_and_init(env.server.port(), "test-token-123", "lead-1").await; + let keep_id = create_task(&mut stream, 2, "Keep", Some("worker-1")).await; + let deleted_id = create_task(&mut stream, 3, "Deleted", Some("worker-2")).await; + update_task_status(&mut stream, 4, &deleted_id, "deleted").await; + + let tasks = list_tasks_with_args(&mut stream, 5, json!({})).await; + assert_eq!(tasks.len(), 2); + assert!(tasks.iter().any(|task| task["id"] == keep_id)); + assert!(tasks.iter().any(|task| task["id"] == deleted_id)); + + env.server.stop(); +} + +#[tokio::test] +async fn ttl4_task_list_filters_owner_status_include_deleted_and_limit() { + let env = setup().await; + let mut stream = connect_and_init(env.server.port(), "test-token-123", "lead-1").await; + let worker_pending = create_task(&mut stream, 2, "Worker pending", Some("worker-1")).await; + let lead_pending = create_task(&mut stream, 3, "Lead pending", Some("lead-1")).await; + let worker_progress = create_task(&mut stream, 4, "Worker progress", Some("worker-1")).await; + let deleted = create_task(&mut stream, 5, "Deleted", Some("worker-1")).await; + update_task_status(&mut stream, 6, &worker_progress, "in_progress").await; + update_task_status(&mut stream, 7, &deleted, "deleted").await; + + let owner_tasks = list_tasks_with_args(&mut stream, 8, json!({"owner": "worker-1"})).await; + assert_eq!(owner_tasks.len(), 3); + assert!(owner_tasks.iter().all(|task| task["owner"] == "worker-1")); + + let pending_tasks = list_tasks_with_args(&mut stream, 9, json!({"status": "pending"})).await; + assert_eq!(pending_tasks.len(), 2); + assert!(pending_tasks.iter().all(|task| task["status"] == "pending")); + + let active_tasks = list_tasks_with_args(&mut stream, 10, json!({"status": ["pending", "in_progress"]})).await; + assert_eq!(active_tasks.len(), 3); + assert!(active_tasks.iter().any(|task| task["id"] == worker_progress)); + assert!(!active_tasks.iter().any(|task| task["id"] == deleted)); + + let no_deleted = list_tasks_with_args(&mut stream, 11, json!({"include_deleted": false})).await; + assert_eq!(no_deleted.len(), 3); + assert!(!no_deleted.iter().any(|task| task["id"] == deleted)); + + let status_with_deleted_flag = + list_tasks_with_args(&mut stream, 12, json!({"status": "pending", "include_deleted": true})).await; + assert_eq!(status_with_deleted_flag.len(), 2); + assert!(status_with_deleted_flag.iter().all(|task| task["status"] == "pending")); + assert!(!status_with_deleted_flag.iter().any(|task| task["id"] == deleted)); + + let limited = list_tasks_with_args(&mut stream, 13, json!({"owner": "worker-1", "limit": 1})).await; + assert_eq!(limited.len(), 1); + assert_eq!(limited[0]["id"], worker_pending); + assert_ne!(limited[0]["id"], lead_pending); + + env.server.stop(); +} + +#[tokio::test] +async fn ttl5_task_list_clamps_large_limit_after_filtering() { + let env = setup().await; + let mut stream = connect_and_init(env.server.port(), "test-token-123", "lead-1").await; + for index in 0..205 { + create_task(&mut stream, 2 + index, &format!("Task {index}"), Some("worker-1")).await; + } + + let tasks = list_tasks_with_args(&mut stream, 300, json!({"limit": 10000})).await; + assert_eq!(tasks.len(), 200); + assert_eq!(tasks[0]["subject"], "Task 0"); + assert_eq!(tasks[199]["subject"], "Task 199"); + + env.server.stop(); +} + +#[tokio::test] +async fn ttl6_task_list_rejects_invalid_filter_arguments() { + let env = setup().await; + let mut stream = connect_and_init(env.server.port(), "test-token-123", "lead-1").await; + + for (id, args, expected) in [ + (2, json!({"slot_id": "worker-1"}), "Invalid params"), + (3, json!({"status": "blocked"}), "Invalid params"), + (4, json!({"status": []}), "Invalid params"), + (5, json!({"limit": 0}), "Invalid params"), + ] { + let resp = call_tool(&mut stream, id, "team_task_list", args).await; + assert!(is_error_response(&resp), "expected error response: {resp}"); + assert!(extract_text(&resp).contains(expected)); + } + + env.server.stop(); +} + +#[tokio::test] +async fn tools_list_dumps_team_tool_schema_when_enabled() { + let temp = tempfile::TempDir::new().unwrap(); + let dump_config = TeamPromptDumpConfig::enabled(temp.path()); + let env = setup_with_prompt_dump(Some(dump_config)).await; + + let mut stream = connect_and_init(env.server.port(), "test-token-123", "lead-1").await; + send_request( + &mut stream, + &json!({ + "jsonrpc": "2.0", + "id": 41, + "method": "tools/list" + }), + ) + .await; + let response = read_response(&mut stream).await; + assert!(response["result"]["tools"].as_array().unwrap().len() > 1); + + let dumps: Vec<_> = std::fs::read_dir(temp.path()) + .unwrap() + .map(|entry| entry.unwrap().path()) + .collect(); + assert_eq!(dumps.len(), 1); + let dump: Value = serde_json::from_str(&std::fs::read_to_string(&dumps[0]).unwrap()).unwrap(); + assert_eq!(dump["kind"], "team-tools-list"); + assert_eq!(dump["scope"], "team-mcp-server-tools-only"); + assert_eq!(dump["not_final_agent_tools"], true); + assert_eq!(dump["team_id"], "team-1"); + assert_eq!(dump["caller_slot_id"], "lead-1"); + assert_eq!(dump["caller_role"], "lead"); + assert!(dump["tools"].as_array().unwrap().iter().any(|tool| { + tool["name"] == "team_spawn_agent" && tool["inputSchema"]["properties"]["assistant_id"].is_object() + })); + + env.server.stop(); +} + // --------------------------------------------------------------------------- // Tests: team_task_update (TTU-1, TTU-2, TTU-3) // --------------------------------------------------------------------------- @@ -640,6 +827,11 @@ async fn ttu1_update_task_status() { .await; assert!(!is_error_response(&resp)); + let text = extract_text(&resp); + let payload: Value = serde_json::from_str(&text).expect("team_task_update must return JSON"); + assert_eq!(payload["status"], "ok"); + assert_eq!(payload["task"]["task_id"], task_id); + assert_eq!(payload["task"]["status"], "completed"); let list_resp2 = call_tool(&mut stream, 5, "team_task_list", json!({})).await; let tasks2: Vec = serde_json::from_str(&extract_text(&list_resp2)).unwrap(); @@ -682,9 +874,16 @@ async fn tm1_list_all_members() { let members: Vec = serde_json::from_str(&text).unwrap(); assert_eq!(members.len(), 2); - let names: Vec<&str> = members.iter().map(|m| m["name"].as_str().unwrap()).collect(); - assert!(names.contains(&"Leader")); - assert!(names.contains(&"Worker")); + assert_eq!(members[0]["name"], "Leader"); + assert_eq!(members[0]["role"], "lead"); + assert_eq!(members[0]["assistant_id"], "lead-assistant"); + assert_eq!(members[0]["model"], "claude"); + assert!(members[0].get("conversation_id").is_none()); + assert_eq!(members[1]["name"], "Worker"); + assert_eq!(members[1]["role"], "teammate"); + assert_eq!(members[1]["assistant_id"], "worker-assistant"); + assert_eq!(members[1]["model"], "claude"); + assert!(members[1].get("conversation_id").is_none()); // Regression: cold-start agents (including the lead before its first // wake) must report an explicit `idle` status — never `null` — so MCP @@ -720,7 +919,31 @@ async fn tra1_rename_existing_agent() { assert!(!is_error_response(&resp)); let text = extract_text(&resp); - assert!(text.contains("renamed")); + let payload: Value = serde_json::from_str(&text).expect("team_rename_agent must return JSON"); + assert_eq!(payload["status"], "ok"); + assert_eq!(payload["action"], "agent_renamed"); + assert_eq!(payload["agent"]["slot_id"], "worker-1"); + assert_eq!(payload["agent"]["name"], "Senior Worker"); + + env.server.stop(); +} + +#[tokio::test] +async fn team_rename_agent_rejects_display_name_target() { + let env = setup().await; + let mut stream = connect_and_init(env.server.port(), "test-token-123", "lead-1").await; + + let resp = call_tool( + &mut stream, + 2, + "team_rename_agent", + json!({"slot_id": "Worker", "new_name": "Senior Worker"}), + ) + .await; + + assert!(is_error_response(&resp)); + let text = extract_text(&resp); + assert!(text.contains("expected slot_id"), "unexpected error: {text}"); env.server.stop(); } @@ -1056,39 +1279,6 @@ async fn sb3_different_agents_get_different_slot_ids() { ); } -// --------------------------------------------------------------------------- -// Tests: mcpStatus broadcast (W5-D31b-1) -// --------------------------------------------------------------------------- - -#[tokio::test] -async fn mcp_status_tcp_ready_is_broadcast_on_successful_bind() { - use aionui_api_types::{TeamMcpPhase, TeamMcpStatusPayload}; - - let env = setup().await; - let port = env.server.port(); - - let events = env.broadcaster.events(); - let status_events: Vec<_> = events.iter().filter(|e| e.name == "team.mcpStatus").collect(); - assert_eq!( - status_events.len(), - 1, - "expected exactly one team.mcpStatus event after bind, got {}", - status_events.len() - ); - - let payload: TeamMcpStatusPayload = serde_json::from_value(status_events[0].data.clone()).unwrap(); - assert_eq!(payload.team_id, "team-1"); - assert_eq!(payload.slot_id, ""); - assert!(matches!(payload.phase, TeamMcpPhase::TcpReady)); - assert_eq!(payload.port, Some(port)); - assert!(payload.server_count.is_none()); - assert!(payload.error.is_none()); - - env.server.stop(); - - env.server.stop(); -} - // --------------------------------------------------------------------------- // Tests: W5-D30b — shutdown_rejected detection in team_send_message // --------------------------------------------------------------------------- @@ -1108,14 +1298,10 @@ async fn tsr1_shutdown_rejected_notifies_lead_and_preserves_agent() { assert!(!is_error_response(&resp)); let text = extract_text(&resp); - assert!( - text.contains("shutdown_rejected"), - "response should echo the sentinel, got: {text}" - ); - assert!( - text.contains("still working"), - "response should echo the reason, got: {text}" - ); + let payload: Value = serde_json::from_str(&text).expect("shutdown_rejected response must be JSON"); + assert_eq!(payload["status"], "ok"); + assert_eq!(payload["action"], "shutdown_rejected"); + assert_eq!(payload["reason"], "still working"); // Leader mailbox contains the notification, worker did not receive a // literal copy of the sentinel. diff --git a/crates/aionui-team/tests/prompts_events_integration.rs b/crates/aionui-team/tests/prompts_events_integration.rs index 8d3924c7e..d3f68a61b 100644 --- a/crates/aionui-team/tests/prompts_events_integration.rs +++ b/crates/aionui-team/tests/prompts_events_integration.rs @@ -1,5 +1,6 @@ mod common; +use std::collections::HashSet; use std::sync::Arc; use aionui_api_types::{ @@ -47,6 +48,10 @@ impl EventBroadcaster for RecordingBroadcaster { } } +fn roster(ids: &[&str]) -> HashSet { + ids.iter().map(|id| (*id).to_owned()).collect() +} + #[derive(Default)] struct RecordingProjectionStore { rows: std::sync::Mutex>, @@ -139,6 +144,24 @@ fn projection_request_for_teammate_mirror_uses_stable_mailbox_dedupe_key() { assert!(!req.visibility.allow_hidden_conversation_message); } +#[test] +fn projection_request_for_team_system_message_uses_stable_mailbox_dedupe_key() { + let req = TeamProjectionRequest::team_system_visible( + "team-1", + "lead-1", + "conv-lead", + "Roster changed", + "mailbox-system-123", + ); + + assert_eq!( + req.dedupe_key.as_deref(), + Some("team:team-1:mailbox:mailbox-system-123:conversation:conv-lead") + ); + assert!(req.visibility.insert_teammate_visible_bubble); + assert!(!req.visibility.allow_hidden_conversation_message); +} + #[tokio::test] async fn projection_inserts_user_visible_bubble_with_stripped_system_notes() { let store = Arc::new(RecordingProjectionStore::default()); @@ -209,6 +232,43 @@ async fn projection_dedupes_teammate_mirror_and_broadcasts_persisted_msg_id() { assert_eq!(events[0].data["from_name"], "Worker"); } +#[tokio::test] +async fn projection_inserts_team_system_bubble_and_broadcasts_it_like_existing_mirror() { + let store = Arc::new(RecordingProjectionStore::default()); + let bc = Arc::new(RecordingBroadcaster::new()); + let projection = TeamMessageProjection::new(store.clone(), bc.clone()); + let req = TeamProjectionRequest::team_system_visible( + "team-1", + "lead-1", + "conv-lead", + "Roster changed", + "mailbox-system-123", + ); + + let projected = projection.project(req).await.unwrap(); + + let msg_id = match projected { + ProjectedTeamMessage::Inserted { msg_id } => msg_id, + other => panic!("expected insert, got {other:?}"), + }; + let rows = store.rows(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].position.as_deref(), Some("left")); + assert_eq!(rows[0].msg_id.as_deref(), Some(msg_id.as_str())); + let content: serde_json::Value = serde_json::from_str(&rows[0].content).unwrap(); + assert_eq!(content["content"], "Roster changed"); + assert_eq!(content["teammate_message"], true); + assert_eq!(content["sender_name"], "team_system"); + + let events = bc.events(); + assert_eq!(events.len(), 1); + assert_eq!(events[0].name, "team.teammateMessage"); + assert_eq!(events[0].data["conversation_id"], "conv-lead"); + assert_eq!(events[0].data["msg_id"], msg_id); + assert_eq!(events[0].data["from_slot_id"], "team_system"); + assert_eq!(events[0].data["from_name"], "team_system"); +} + // =========================================================================== // Test-plan §9: Prompt Templates // =========================================================================== @@ -239,29 +299,33 @@ fn default_assistants() -> Vec { ] } -// -- LP-1: Lead prompt contains member list ---------------------------------- +// -- LP-1: Lead prompt relies on team_members for roster ---------------------- #[test] -fn lp1_lead_prompt_contains_member_list() { +fn lp1_lead_prompt_does_not_contain_member_snapshot() { let members = vec![ make_agent("lead-1", "Lead", TeammateRole::Lead), make_agent("w1", "Alice", TeammateRole::Teammate), make_agent("w2", "Bob", TeammateRole::Teammate), ]; let assistants = default_assistants(); - let prompt = build_lead_prompt("Alpha", &members, &assistants); - - // AionUi bullet format: `- {name} ({backend}, status: {status})` - assert!(prompt.contains("- Lead ("), "lead name missing"); - assert!(prompt.contains("- Alice ("), "teammate Alice missing"); - assert!(prompt.contains("- Bob ("), "teammate Bob missing"); + let lead = make_agent("lead-1", "Lead", TeammateRole::Lead); + let prompt = build_lead_prompt(&lead, "Alpha", &members, &assistants); + + assert!(!prompt.contains("## Your Teammates")); + assert!(!prompt.contains("- Lead ("), "lead snapshot leaked"); + assert!(!prompt.contains("- Alice ("), "teammate Alice snapshot leaked"); + assert!(!prompt.contains("- Bob ("), "teammate Bob snapshot leaked"); + assert!(prompt.to_lowercase().contains("first team turn")); + assert!(prompt.contains("team_members")); } // -- LP-2: Lead prompt contains tool descriptions ---------------------------- #[test] fn lp2_lead_prompt_contains_tool_descriptions() { - let prompt = build_lead_prompt("Beta", &[], &default_assistants()); + let lead = make_agent("lead-1", "Lead", TeammateRole::Lead); + let prompt = build_lead_prompt(&lead, "Beta", &[], &default_assistants()); // AionUi lead prompt references the `team_*` coordination tools that the // leader must use; the MCP layer enumerates them with arguments, so the @@ -274,18 +338,20 @@ fn lp2_lead_prompt_contains_tool_descriptions() { "team_members", "team_rename_agent", "team_shutdown_agent", - "team_list_models", ]; for tool in expected_tools { assert!(prompt.contains(tool), "missing tool: {tool}"); } + assert!(prompt.contains("team_list_assistants")); + assert!(!prompt.contains("## Available Assistants for Spawning")); } // -- LP-3: Lead prompt contains task management guidance --------------------- #[test] fn lp3_lead_prompt_contains_task_management_guidance() { - let prompt = build_lead_prompt("Gamma", &[], &default_assistants()); + let lead = make_agent("lead-1", "Lead", TeammateRole::Lead); + let prompt = build_lead_prompt(&lead, "Gamma", &[], &default_assistants()); assert!( prompt.contains("Break the work into tasks"), @@ -318,6 +384,11 @@ fn tp1_teammate_prompt_contains_execution_guidance() { assert!(prompt.contains("shutdown_request"), "missing shutdown protocol"); assert!(prompt.contains("shutdown_approved"), "missing shutdown_approved"); assert!(prompt.contains("STOP GENERATING"), "missing stop protocol"); + assert!(prompt.contains("Slot ID: w1"), "missing teammate slot id"); + assert!( + !prompt.contains("Teammates:"), + "static teammate list must not be injected" + ); } // -- TP-2: Teammate prompt contains team name -------------------------------- @@ -362,7 +433,7 @@ fn wp1_wake_payload_includes_unread_messages() { created_at: 0, }, ]; - let payload = build_wake_payload(&agent, &[], &messages); + let payload = build_wake_payload(&agent, &[], &messages, &roster(&["lead-1", "w1", "w2"])); assert!(payload.contains("Feature X is done")); assert!(payload.contains("`w1`")); @@ -405,18 +476,20 @@ fn wp2_wake_payload_includes_task_list() { updated_at: 0, }, ]; - let payload = build_wake_payload(&agent, &tasks, &[]); + let payload = build_wake_payload(&agent, &tasks, &[], &roster(&["lead-1", "w1", "w2"])); - assert!(payload.contains("Current Task Board")); + assert!(payload.contains("Current Task Board Summary")); + assert!(payload.contains("Showing 2 of 2 tasks.")); assert!(payload.contains("Implement auth")); assert!(payload.contains("in_progress")); assert!(payload.contains("Write tests")); assert!(payload.contains("pending")); assert!(payload.contains("w1")); assert!(payload.contains("w2")); + assert!(payload.contains("aaaaaaaa…")); assert!( - payload.contains("aaaaaaaa-1111-2222-3333-444444444444"), - "blocker task ID should appear in blocked_by column" + !payload.contains("aaaaaaaa-1111-2222-3333-444444444444"), + "summary blocked_by column should use short task IDs" ); } @@ -425,7 +498,7 @@ fn wp2_wake_payload_includes_task_list() { #[test] fn wp3_wake_payload_empty_builds_normally() { let agent = make_agent("w1", "Worker1", TeammateRole::Teammate); - let payload = build_wake_payload(&agent, &[], &[]); + let payload = build_wake_payload(&agent, &[], &[], &roster(&["w1"])); assert!(payload.contains("No new messages")); assert!(payload.contains("No tasks on the board")); diff --git a/crates/aionui-team/tests/session_service_integration.rs b/crates/aionui-team/tests/session_service_integration.rs index b1d62c53c..e8d8ba223 100644 --- a/crates/aionui-team/tests/session_service_integration.rs +++ b/crates/aionui-team/tests/session_service_integration.rs @@ -3,6 +3,7 @@ mod common; use std::collections::HashMap; use std::sync::Arc; use std::sync::Mutex; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use aionui_ai_agent::session_context::{ AcpSessionBuildContext, AgentSessionContext, AgentSessionKind, ConversationContext, WorkspaceContext, @@ -10,7 +11,10 @@ use aionui_ai_agent::session_context::{ use aionui_ai_agent::task_manager::AgentFactory; use aionui_ai_agent::types::BuildTaskOptions; use aionui_ai_agent::{ActiveLeaseRegistry, AgentError, IWorkerTaskManager, WorkerTaskManagerImpl}; -use aionui_api_types::{AcpBuildExtra, AddAgentRequest, CreateTeamRequest, TeamAgentInput, WebSocketMessage}; +use aionui_api_types::{ + AcpBuildExtra, AcpConfigOptionDto, AcpConfigSelectOptionDto, AddAgentRequest, CreateTeamRequest, + GetConfigOptionsResponse, TeamAgentInput, WebSocketMessage, +}; use aionui_common::{AgentKillReason, AgentType, PaginatedResult, ProviderWithModel}; use aionui_db::models::{ AgentMetadataRow, AssistantDefinitionRow, AssistantOverlayRow, ConversationRow, MessageRow, @@ -20,7 +24,7 @@ use aionui_db::models::{ use aionui_db::{ ConversationFilters, ConversationRowUpdate, DbError, IAgentMetadataRepository, IAssistantDefinitionRepository, IAssistantOverlayRepository, IConversationRepository, IProviderRepository, ITeamRepository, MessagePageParams, - MessagePageResult, MessageRowUpdate, MessageSearchRow, + MessagePageResult, MessageRowUpdate, MessageSearchRow, resolve_agent_binding_from_rows, }; use aionui_realtime::EventBroadcaster; @@ -43,12 +47,14 @@ use common::MockTeamRepo; struct MockConversationRepo { conversations: std::sync::Mutex>, + messages: std::sync::Mutex>, } impl MockConversationRepo { fn new() -> Self { Self { conversations: std::sync::Mutex::new(Vec::new()), + messages: std::sync::Mutex::new(Vec::new()), } } @@ -64,6 +70,16 @@ impl MockConversationRepo { self.conversations.lock().unwrap().len() } + fn messages_for(&self, conversation_id: &str) -> Vec { + self.messages + .lock() + .unwrap() + .iter() + .filter(|message| message.conversation_id == conversation_id) + .cloned() + .collect() + } + fn patch_extra(&self, id: &str, patch: serde_json::Value) -> Result<(), DbError> { let mut convs = self.conversations.lock().unwrap(); let conv = convs @@ -155,7 +171,8 @@ impl IConversationRepository for MockConversationRepo { has_more_after: false, }) } - async fn insert_message(&self, _message: &MessageRow) -> Result<(), DbError> { + async fn insert_message(&self, message: &MessageRow) -> Result<(), DbError> { + self.messages.lock().unwrap().push(message.clone()); Ok(()) } async fn update_message(&self, _id: &str, _updates: &MessageRowUpdate) -> Result<(), DbError> { @@ -166,11 +183,19 @@ impl IConversationRepository for MockConversationRepo { } async fn get_message_by_msg_id( &self, - _conv_id: &str, - _msg_id: &str, - _msg_type: &str, + conv_id: &str, + msg_id: &str, + msg_type: &str, ) -> Result, DbError> { - Ok(None) + Ok(self + .messages + .lock() + .unwrap() + .iter() + .find(|row| { + row.conversation_id == conv_id && row.msg_id.as_deref() == Some(msg_id) && row.r#type == msg_type + }) + .cloned()) } async fn search_messages( &self, @@ -203,7 +228,7 @@ impl AgentTurnExecutionPort for NoopTurnPort { async fn run_agent_turn(&self, request: AgentTurnRequest) -> Result { if let Some(on_started) = request.on_started.as_ref() { on_started(AgentTurnStarted { - team_run_id: request.team_run_id.clone().expect("team run id"), + team_run_id: request.team_run_id.clone(), slot_id: request.slot_id.clone(), role: request.role.clone(), conversation_id: request.conversation_id.clone(), @@ -230,7 +255,7 @@ impl AgentTurnExecutionPort for RecordingTurnPort { async fn run_agent_turn(&self, request: AgentTurnRequest) -> Result { if let Some(on_started) = request.on_started.as_ref() { on_started(AgentTurnStarted { - team_run_id: request.team_run_id.clone().expect("team run id"), + team_run_id: request.team_run_id.clone(), slot_id: request.slot_id.clone(), role: request.role.clone(), conversation_id: request.conversation_id.clone(), @@ -463,6 +488,38 @@ impl TeamConversationProvisioningPort for FakeConversationPorts { .await } + async fn get_config_options( + &self, + conversation_id: &str, + ) -> Result { + let extra = self + .repo + .get_extra(conversation_id) + .ok_or_else(|| aionui_team::TeamError::AgentNotFound(conversation_id.to_owned()))?; + let model = extra + .get("current_model_id") + .and_then(serde_json::Value::as_str) + .unwrap_or("mock-model") + .to_owned(); + Ok(GetConfigOptionsResponse { + config_options: vec![AcpConfigOptionDto { + id: "model".to_owned(), + name: None, + label: Some("Model".to_owned()), + description: None, + category: Some("model".to_owned()), + option_type: "select".to_owned(), + current_value: Some(model.clone()), + options: vec![AcpConfigSelectOptionDto { + value: model.clone(), + name: None, + label: Some(model), + description: None, + }], + }], + }) + } + async fn warmup_agent_process( &self, user_id: &str, @@ -611,6 +668,10 @@ impl RecordingBroadcaster { .cloned() .collect() } + + fn clear(&self) { + self.events.lock().unwrap().clear(); + } } impl EventBroadcaster for RecordingBroadcaster { @@ -887,6 +948,16 @@ impl CountingTaskManager { fn snapshot(&self) -> TaskManagerCalls { self.calls.lock().unwrap().clone() } + + fn reset_calls(&self) { + *self.calls.lock().unwrap() = TaskManagerCalls::default(); + } + + async fn remove_task_without_recording(&self, conversation_id: &str) { + self.inner + .kill_and_wait(conversation_id, Some(AgentKillReason::TeamMcpRebuild)) + .await; + } } #[async_trait::async_trait] @@ -915,8 +986,12 @@ impl IWorkerTaskManager for CountingTaskManager { conversation_id: &str, reason: Option, ) -> std::pin::Pin + Send>> { - let _ = self.kill(conversation_id, reason); - Box::pin(std::future::ready(())) + self.calls + .lock() + .unwrap() + .kill + .push((conversation_id.to_owned(), reason)); + self.inner.kill_and_wait(conversation_id, reason) } async fn clear(&self) { self.inner.clear().await @@ -1028,6 +1103,91 @@ fn success_factory() -> AgentFactory { }) } +fn blocking_first_build_factory(started: Arc, release: Arc) -> AgentFactory { + use futures_util::FutureExt; + + let calls = Arc::new(AtomicUsize::new(0)); + Arc::new(move |opts: BuildTaskOptions| { + let started = Arc::clone(&started); + let release = Arc::clone(&release); + let calls = Arc::clone(&calls); + async move { + if calls.fetch_add(1, Ordering::SeqCst) == 0 { + started.notify_one(); + release.notified().await; + } + + Ok(aionui_ai_agent::AgentInstance::Mock(Arc::new( + mock_agent::MockAgent::new(opts.context.conversation.conversation_id, opts.context.workspace.path), + ))) + } + .boxed() + }) +} + +struct GatedProvisioningFactory { + enabled: AtomicBool, + starts: Mutex>, + started: tokio::sync::Notify, + release: tokio::sync::Semaphore, +} + +impl Default for GatedProvisioningFactory { + fn default() -> Self { + Self { + enabled: AtomicBool::new(false), + starts: Mutex::new(Vec::new()), + started: tokio::sync::Notify::new(), + release: tokio::sync::Semaphore::new(0), + } + } +} + +impl GatedProvisioningFactory { + fn factory(self: &Arc) -> AgentFactory { + use futures_util::FutureExt; + + let gate = Arc::clone(self); + Arc::new(move |opts: BuildTaskOptions| { + let gate = Arc::clone(&gate); + async move { + let conversation_id = opts.context.conversation.conversation_id.clone(); + if gate.enabled.swap(false, Ordering::SeqCst) { + gate.starts.lock().unwrap().push(conversation_id.clone()); + gate.started.notify_waiters(); + gate.release.acquire().await.expect("provisioning gate closed").forget(); + } + Ok(aionui_ai_agent::AgentInstance::Mock(Arc::new( + mock_agent::MockAgent::new(conversation_id, opts.context.workspace.path), + ))) + } + .boxed() + }) + } + + fn enable(&self) { + self.enabled.store(true, Ordering::SeqCst); + } + + async fn wait_for_starts(&self, count: usize) { + loop { + let notified = self.started.notified(); + if self.starts.lock().unwrap().len() >= count { + return; + } + notified.await; + } + } + + fn starts(&self) -> Vec { + self.starts.lock().unwrap().clone() + } + + fn release(&self, count: usize) { + self.release.add_permits(count); + } +} + fn confirmations_factory(count: usize) -> AgentFactory { use aionui_common::Confirmation; use futures_util::FutureExt; @@ -1096,6 +1256,45 @@ impl TeamAssistantCatalogPort for EmptyTeamAssistantCatalog { } } +struct TestTeamAssistantCatalog { + agent_metadata_repo: Arc, + assistant_definition_repo: Arc, + assistant_overlay_repo: Arc, +} + +#[async_trait::async_trait] +impl TeamAssistantCatalogPort for TestTeamAssistantCatalog { + async fn list_team_selectable_assistants(&self) -> Result, TeamError> { + let agent_rows = self.agent_metadata_repo.list_all().await?; + let definitions = self.assistant_definition_repo.list().await?; + let mut result = Vec::new(); + + for definition in definitions { + let overlay = self.assistant_overlay_repo.get(&definition.id).await?; + if overlay.as_ref().is_some_and(|row| !row.enabled) { + continue; + } + let effective_agent_id = overlay + .as_ref() + .and_then(|row| row.agent_id_override.as_deref()) + .filter(|value| !value.trim().is_empty()) + .unwrap_or(definition.agent_id.as_str()); + let backend = resolve_agent_binding_from_rows(&agent_rows, effective_agent_id) + .map(|binding| binding.runtime_backend) + .unwrap_or_else(|| effective_agent_id.to_owned()); + result.push(TeamAssistantCatalogEntry { + assistant_id: definition.assistant_id, + name: definition.name, + backend, + description: definition.description.unwrap_or_default(), + skills: Vec::new(), + }); + } + + Ok(result) + } +} + struct EmptyProviderRepo; #[async_trait::async_trait] @@ -1321,10 +1520,15 @@ fn setup_with_factory_metadata_assistants_and_conversation_repo( let task_manager_dyn: Arc = task_manager.clone(); let backend_binary_path = Arc::new(std::path::PathBuf::from("/tmp/aioncore-test")); let provider_repo: Arc = Arc::new(EmptyProviderRepo); + let assistant_catalog: Arc = Arc::new(TestTeamAssistantCatalog { + agent_metadata_repo: agent_metadata_repo.clone(), + assistant_definition_repo: assistant_definition_repo.clone(), + assistant_overlay_repo: assistant_overlay_repo.clone(), + }); let svc = TeamSessionService::new( team_repo_dyn, agent_metadata_repo, - Arc::new(EmptyTeamAssistantCatalog), + assistant_catalog, assistant_definition_repo, assistant_overlay_repo, provider_repo, @@ -1377,10 +1581,15 @@ fn setup_with_ports_metadata_assistants_and_conversation_repo( let task_manager: Arc = Arc::new(CountingTaskManager::new(factory)); let backend_binary_path = Arc::new(std::path::PathBuf::from("/tmp/aioncore-test")); let provider_repo: Arc = Arc::new(EmptyProviderRepo); + let assistant_catalog: Arc = Arc::new(TestTeamAssistantCatalog { + agent_metadata_repo: agent_metadata_repo.clone(), + assistant_definition_repo: assistant_definition_repo.clone(), + assistant_overlay_repo: assistant_overlay_repo.clone(), + }); let svc = TeamSessionService::new( team_repo_dyn, agent_metadata_repo, - Arc::new(EmptyTeamAssistantCatalog), + assistant_catalog, assistant_definition_repo, assistant_overlay_repo, provider_repo, @@ -1435,7 +1644,7 @@ fn setup() -> Arc { } #[tokio::test] -async fn ensure_session_recovery_drain_runs_agent_turn_with_team_run_id() { +async fn recovery_creates_background_intents_without_restoring_old_memory_run() { let (svc, team_repo, turn_port, _conv_repo) = setup_with_recording_turn_port(); let created = svc .create_team( @@ -1451,7 +1660,7 @@ async fn ensure_session_recovery_drain_runs_agent_turn_with_team_run_id() { let lead_slot_id = created.leader_assistant_id.clone().expect("lead"); svc.stop_session("user1", &created.id) .await - .expect("stop auto-started session"); + .expect("clear existing session"); team_repo .write_message(&aionui_db::models::MailboxMessageRow { @@ -1485,7 +1694,10 @@ async fn ensure_session_recovery_drain_runs_agent_turn_with_team_run_id() { let requests = turn_port.requests.lock().unwrap(); assert_eq!(requests.len(), 1); assert_eq!(requests[0].slot_id, lead_slot_id); - assert!(requests[0].team_run_id.is_some(), "recovery turn must be TeamRun-owned"); + assert_eq!( + requests[0].team_run_id, None, + "recovery work must not synthesize a user-visible TeamRun" + ); } #[tokio::test] @@ -1505,7 +1717,7 @@ async fn teammate_first_wake_uses_canonical_prompt_at_service_boundary() { let worker_slot_id = created.assistants[1].slot_id.clone(); svc.stop_session("user1", &created.id) .await - .expect("stop auto-started session"); + .expect("clear existing session"); team_repo .write_message(&aionui_db::models::MailboxMessageRow { @@ -1575,7 +1787,7 @@ async fn ensure_session_does_not_run_self_message_only_recovery_turn() { let lead_slot_id = created.leader_assistant_id.clone().expect("lead"); svc.stop_session("user1", &created.id) .await - .expect("stop auto-started session"); + .expect("clear existing session"); team_repo .write_message(&aionui_db::models::MailboxMessageRow { @@ -1632,6 +1844,59 @@ fn setup_with_recording_broadcaster() -> (Arc, Arc ( + Arc, + Arc, + Arc, + Arc, +) { + let (service, team_repo, task_manager, recorder, _conversation_repo) = + setup_with_factory_recording_broadcaster_and_conversation_repo(factory); + (service, team_repo, task_manager, recorder) +} + +type RecordingServiceHarness = ( + Arc, + Arc, + Arc, + Arc, + Arc, +); + +fn setup_with_factory_recording_broadcaster_and_conversation_repo(factory: AgentFactory) -> RecordingServiceHarness { + let team_repo = Arc::new(FullMockTeamRepo::new()); + let team_repo_dyn: Arc = team_repo.clone(); + let conv_repo = Arc::new(MockConversationRepo::new()); + let recorder = Arc::new(RecordingBroadcaster::new()); + let broadcaster: Arc = recorder.clone(); + let agent_metadata_repo: Arc = Arc::new(StubAgentMetadataRepo::empty()); + let conversation_ports = Arc::new(FakeConversationPorts::new(conv_repo.clone())); + let conversation_port: Arc = conversation_ports.clone(); + let projection_store: Arc = conversation_ports.clone(); + let task_manager = Arc::new(CountingTaskManager::new(factory)); + let task_manager_dyn: Arc = task_manager.clone(); + let backend_binary_path = Arc::new(std::path::PathBuf::from("/tmp/aioncore-test")); + let provider_repo: Arc = Arc::new(EmptyProviderRepo); + let svc = TeamSessionService::new( + team_repo_dyn, + agent_metadata_repo, + Arc::new(EmptyTeamAssistantCatalog), + Arc::new(EmptyAssistantDefinitionRepo), + Arc::new(EmptyAssistantOverlayRepo), + provider_repo, + conversation_port, + projection_store, + broadcaster, + task_manager_dyn, + noop_turn_port(), + noop_cancellation_port(), + backend_binary_path, + ); + (svc, team_repo, task_manager, recorder, conv_repo) +} + fn make_agent_metadata_row(id: &str, backend: &str, icon: &str) -> AgentMetadataRow { AgentMetadataRow { id: id.to_owned(), @@ -1753,7 +2018,99 @@ fn two_agent_input() -> Vec { ] } -async fn reset_auto_started_session(svc: &Arc, tm: &Arc, team_id: &str) { +fn team_agent_input(name: &str, role: &str, model: &str) -> TeamAgentInput { + TeamAgentInput { + name: name.into(), + role: role.into(), + backend: Some("acp".into()), + model: model.into(), + assistant_id: None, + conversation_id: None, + } +} + +fn four_agent_input_leader_not_first() -> Vec { + vec![ + team_agent_input("Worker 1", "teammate", "worker-1"), + team_agent_input("Lead", "lead", "lead-model"), + team_agent_input("Worker 2", "teammate", "worker-2"), + team_agent_input("Worker 3", "teammate", "worker-3"), + ] +} + +fn five_agent_input_leader_not_first() -> Vec { + vec![ + team_agent_input("Worker 1", "teammate", "worker-1"), + team_agent_input("Worker 2", "teammate", "worker-2"), + team_agent_input("Lead", "lead", "lead-model"), + team_agent_input("Worker 3", "teammate", "worker-3"), + team_agent_input("Worker 4", "teammate", "worker-4"), + ] +} + +struct WarmupConcurrencyProbe { + active: AtomicUsize, + max_active: AtomicUsize, + starts: Mutex>, + start_times: Mutex>, + started_at: tokio::time::Instant, +} + +impl Default for WarmupConcurrencyProbe { + fn default() -> Self { + Self { + active: AtomicUsize::new(0), + max_active: AtomicUsize::new(0), + starts: Mutex::new(Vec::new()), + start_times: Mutex::new(Vec::new()), + started_at: tokio::time::Instant::now(), + } + } +} + +impl WarmupConcurrencyProbe { + fn factory(self: &Arc, delay: std::time::Duration) -> AgentFactory { + use futures_util::FutureExt; + + let probe = Arc::clone(self); + Arc::new(move |opts: BuildTaskOptions| { + let probe = Arc::clone(&probe); + async move { + let conversation_id = opts.context.conversation.conversation_id.clone(); + let elapsed = probe.started_at.elapsed(); + probe.starts.lock().unwrap().push(conversation_id.clone()); + probe + .start_times + .lock() + .unwrap() + .push((conversation_id.clone(), elapsed)); + let current = probe.active.fetch_add(1, Ordering::SeqCst) + 1; + probe.max_active.fetch_max(current, Ordering::SeqCst); + tokio::time::sleep(delay).await; + probe.active.fetch_sub(1, Ordering::SeqCst); + + Ok(aionui_ai_agent::AgentInstance::Mock(Arc::new( + mock_agent::MockAgent::new(conversation_id, opts.context.workspace.path), + ))) + } + .boxed() + }) + } + + fn max_active(&self) -> usize { + self.max_active.load(Ordering::SeqCst) + } + + fn starts(&self) -> Vec { + self.starts.lock().unwrap().clone() + } + + fn start_times(&self) -> Vec<(String, std::time::Duration)> { + self.start_times.lock().unwrap().clone() + } +} + +async fn reset_runtime_state(svc: &Arc, tm: &Arc, team_id: &str) { svc.stop_session("user1", team_id).await.unwrap(); tm.reset().await; } @@ -2417,7 +2774,6 @@ async fn spawned_preset_assistant_snapshot_is_frozen() { SpawnAgentRequest { name: "Writer".into(), assistant_id: Some("word-creator".into()), - model: Some("claude-sonnet-4".into()), }, ) .await @@ -2700,7 +3056,7 @@ async fn tc2_create_single_agent_team() { } #[tokio::test] -async fn tc4_first_agent_is_lead() { +async fn create_team_uses_explicit_leader_role_when_leader_is_not_first() { let svc = setup(); let resp = svc .create_team( @@ -2709,7 +3065,7 @@ async fn tc4_first_agent_is_lead() { name: "T".into(), agents: vec![ TeamAgentInput { - name: "A".into(), + name: "Worker".into(), role: "teammate".into(), backend: Some("acp".into()), model: "claude".into(), @@ -2717,8 +3073,8 @@ async fn tc4_first_agent_is_lead() { conversation_id: None, }, TeamAgentInput { - name: "B".into(), - role: "teammate".into(), + name: "Lead".into(), + role: "lead".into(), backend: Some("acp".into()), model: "claude".into(), assistant_id: None, @@ -2731,10 +3087,95 @@ async fn tc4_first_agent_is_lead() { .await .unwrap(); + assert_eq!(resp.assistants[0].name, "Lead"); assert_eq!(resp.assistants[0].role, "lead"); + assert_eq!(resp.assistants[1].name, "Worker"); + assert_eq!(resp.assistants[1].role, "teammate"); assert_eq!(resp.leader_assistant_id, Some(resp.assistants[0].slot_id.clone())); } +#[tokio::test] +async fn create_team_rejects_zero_leaders() { + let svc = setup(); + let result = svc + .create_team( + "user1", + CreateTeamRequest { + name: "T".into(), + agents: vec![TeamAgentInput { + name: "Worker".into(), + role: "teammate".into(), + backend: Some("acp".into()), + model: "claude".into(), + assistant_id: None, + conversation_id: None, + }], + workspace: None, + }, + ) + .await; + + assert!(matches!(result, Err(TeamError::InvalidRequest(_)))); +} + +#[tokio::test] +async fn create_team_rejects_multiple_leaders() { + let svc = setup(); + let result = svc + .create_team( + "user1", + CreateTeamRequest { + name: "T".into(), + agents: vec![ + TeamAgentInput { + name: "Lead A".into(), + role: "lead".into(), + backend: Some("acp".into()), + model: "claude".into(), + assistant_id: None, + conversation_id: None, + }, + TeamAgentInput { + name: "Lead B".into(), + role: "leader".into(), + backend: Some("acp".into()), + model: "claude".into(), + assistant_id: None, + conversation_id: None, + }, + ], + workspace: None, + }, + ) + .await; + + assert!(matches!(result, Err(TeamError::InvalidRequest(_)))); +} + +#[tokio::test] +async fn create_team_rejects_unknown_role() { + let svc = setup(); + let result = svc + .create_team( + "user1", + CreateTeamRequest { + name: "T".into(), + agents: vec![TeamAgentInput { + name: "Lead".into(), + role: "captain".into(), + backend: Some("acp".into()), + model: "claude".into(), + assistant_id: None, + conversation_id: None, + }], + workspace: None, + }, + ) + .await; + + assert!(matches!(result, Err(TeamError::InvalidRequest(_)))); +} + #[tokio::test] async fn tc5_empty_agents_returns_error() { let svc = setup(); @@ -3057,14 +3498,12 @@ async fn aa1_add_agent_to_team() { } #[tokio::test] -async fn aa_add_agent_inherits_team_workspace() { - let agent_metadata_repo: Arc = Arc::new(StubAgentMetadataRepo::empty()); - let (svc, _, conv_repo) = - setup_with_factory_and_metadata_and_conversation_repo(success_factory(), agent_metadata_repo); - let workspace = std::env::temp_dir().join(format!("aionui-team-workspace-{}", aionui_common::generate_id())); - std::fs::create_dir_all(&workspace).unwrap(); - let workspace = workspace.to_string_lossy().into_owned(); - let created = svc +async fn manual_add_without_active_run_queues_background_welcome_without_creating_run() { + let (svc, _, conv_repo) = setup_with_factory_and_metadata_and_conversation_repo( + success_factory(), + Arc::new(StubAgentMetadataRepo::empty()), + ); + let created = svc .create_team( "user1", CreateTeamRequest { @@ -3077,13 +3516,14 @@ async fn aa_add_agent_inherits_team_workspace() { assistant_id: None, conversation_id: None, }], - workspace: Some(workspace.clone()), + workspace: None, }, ) .await .unwrap(); + svc.ensure_session("user1", &created.id).await.unwrap(); - let agent = svc + let added = svc .add_agent( "user1", &created.id, @@ -3098,138 +3538,138 @@ async fn aa_add_agent_inherits_team_workspace() { .await .unwrap(); - let extra = conv_repo.get_extra(&agent.conversation_id).unwrap(); - assert_eq!( - extra.get("workspace").and_then(|v| v.as_str()), - Some(workspace.as_str()) + let leader_conversation_id = &created.assistants[0].conversation_id; + let leader_messages = conv_repo.messages_for(leader_conversation_id); + assert_eq!(leader_messages.len(), 1); + assert_eq!(leader_messages[0].position.as_deref(), Some("left")); + let leader_content: serde_json::Value = serde_json::from_str(&leader_messages[0].content).unwrap(); + assert_eq!(leader_content["sender_name"], "team_system"); + assert_eq!(leader_content["teammate_message"], true); + assert!( + leader_content["content"] + .as_str() + .unwrap() + .contains("manually added teammate") ); + assert!(leader_content["content"].as_str().unwrap().contains(&added.slot_id)); + + let added_messages = conv_repo.messages_for(&added.conversation_id); + assert_eq!(added_messages.len(), 1); + assert_eq!(added_messages[0].position.as_deref(), Some("left")); + let added_content: serde_json::Value = serde_json::from_str(&added_messages[0].content).unwrap(); + assert_eq!(added_content["sender_name"], "team_system"); + assert_eq!(added_content["teammate_message"], true); + assert!(added_content["content"].as_str().unwrap().contains("manually added")); + + let run_state = svc.get_run_state("user1", &created.id).await.unwrap(); + assert!(run_state.active_run.is_none()); + assert!(run_state.slot_work.iter().any(|slot| slot.slot_id == added.slot_id)); } #[tokio::test] -async fn add_agent_backfills_empty_team_workspace_from_leader_workspace() { - let agent_metadata_repo: Arc = Arc::new(StubAgentMetadataRepo::empty()); - let (svc, team_repo, _, conv_repo) = - setup_with_factory_metadata_team_repo_and_conversation_repo(success_factory(), agent_metadata_repo); +async fn add_agent_rejects_leader_role() { + let svc = setup(); let created = svc .create_team( "user1", CreateTeamRequest { - name: "Legacy".into(), - agents: vec![TeamAgentInput { - name: "Lead".into(), - role: "lead".into(), - backend: Some("acp".into()), - model: "claude".into(), - assistant_id: None, - conversation_id: None, - }], + name: "T".into(), + agents: two_agent_input(), workspace: None, }, ) .await .unwrap(); - let leader_workspace = conv_repo.get_extra(&created.assistants[0].conversation_id).unwrap()["workspace"] - .as_str() - .unwrap() - .to_owned(); - - force_team_workspace(&team_repo, &created.id, "").await; - let added = svc + let result = svc .add_agent( "user1", &created.id, AddAgentRequest { - name: "Worker".into(), - role: "teammate".into(), + name: "Second Leader".into(), + role: "lead".into(), backend: Some("acp".into()), model: "claude".into(), assistant_id: None, }, ) - .await - .unwrap(); + .await; - let got = svc.get_team("user1", &created.id).await.unwrap(); - assert_eq!(got.workspace, leader_workspace); - let added_extra = conv_repo.get_extra(&added.conversation_id).unwrap(); - assert_eq!( - added_extra.get("workspace").and_then(serde_json::Value::as_str), - Some(leader_workspace.as_str()) - ); + assert!(matches!(result, Err(TeamError::InvalidRequest(_)))); } #[tokio::test] -async fn add_agent_uses_team_temp_workspace_when_team_and_leader_workspaces_are_unusable() { - let agent_metadata_repo: Arc = Arc::new(StubAgentMetadataRepo::empty()); - let (svc, team_repo, _, conv_repo) = - setup_with_factory_metadata_team_repo_and_conversation_repo(success_factory(), agent_metadata_repo); +async fn add_agent_allows_same_assistant_id_multiple_times() { + let definition_repo: Arc = Arc::new(SingleAssistantDefinitionRepo { + row: word_creator_definition(), + }); + let (svc, _team_repo, _conversation_ports, _conv_repo) = setup_with_ports_metadata_assistants_and_conversation_repo( + success_factory(), + seeded_agent_metadata_repo(), + definition_repo, + Arc::new(EmptyAssistantOverlayRepo), + ); let created = svc .create_team( "user1", CreateTeamRequest { - name: "Legacy Empty".into(), - agents: vec![TeamAgentInput { - name: "Lead".into(), - role: "lead".into(), - backend: Some("acp".into()), - model: "claude".into(), - assistant_id: None, - conversation_id: None, - }], + name: "T".into(), + agents: two_agent_input(), workspace: None, }, ) .await .unwrap(); - force_team_workspace(&team_repo, &created.id, "").await; - conv_repo - .patch_extra( - &created.assistants[0].conversation_id, - serde_json::json!({ "workspace": "/tmp/aionui-team-missing-leader-workspace" }), - ) - .unwrap(); - - let added = svc - .add_agent( + for name in ["Worker Copy A", "Worker Copy B"] { + svc.add_agent( "user1", &created.id, AddAgentRequest { - name: "Worker".into(), + name: name.into(), role: "teammate".into(), backend: Some("acp".into()), model: "claude".into(), - assistant_id: None, + assistant_id: Some("word-creator".into()), }, ) .await .unwrap(); + } let got = svc.get_team("user1", &created.id).await.unwrap(); - assert!( - got.workspace - .contains(&format!("/conversations/team-temp-{}", created.id)), - "unexpected team temp workspace: {}", - got.workspace - ); - let added_extra = conv_repo.get_extra(&added.conversation_id).unwrap(); - assert_eq!( - added_extra.get("workspace").and_then(serde_json::Value::as_str), - Some(got.workspace.as_str()) - ); + let matching = got + .assistants + .iter() + .filter(|agent| agent.assistant_id.as_deref() == Some("word-creator")) + .count(); + assert_eq!(matching, 2); } #[tokio::test] -async fn add_agent_does_not_create_teammate_when_workspace_writeback_fails() { - let agent_metadata_repo: Arc = Arc::new(StubAgentMetadataRepo::empty()); - let (svc, team_repo, _, conv_repo) = - setup_with_factory_metadata_team_repo_and_conversation_repo(success_factory(), agent_metadata_repo); +async fn manual_add_agent_active_session_attaches_runtime_in_background_without_blocking_http() { + use futures_util::FutureExt; + + let build_count = Arc::new(AtomicUsize::new(0)); + let factory_count = build_count.clone(); + let factory: AgentFactory = Arc::new(move |opts: BuildTaskOptions| { + let build_index = factory_count.fetch_add(1, Ordering::SeqCst); + async move { + if build_index >= 1 { + tokio::time::sleep(std::time::Duration::from_millis(250)).await; + } + Ok(aionui_ai_agent::AgentInstance::Mock(Arc::new( + mock_agent::MockAgent::new(opts.context.conversation.conversation_id, opts.context.workspace.path), + ))) + } + .boxed() + }); + let (svc, _team_repo, task_manager, _recorder) = setup_with_factory_and_recording_broadcaster(factory); let created = svc .create_team( "user1", CreateTeamRequest { - name: "Writeback Failure".into(), + name: "T".into(), agents: vec![TeamAgentInput { name: "Lead".into(), role: "lead".into(), @@ -3243,12 +3683,12 @@ async fn add_agent_does_not_create_teammate_when_workspace_writeback_fails() { ) .await .unwrap(); + assert_eq!(task_manager.snapshot().build.len(), 0); + svc.ensure_session("user1", &created.id).await.unwrap(); + assert_eq!(task_manager.snapshot().build.len(), 1); - force_team_workspace(&team_repo, &created.id, "").await; - team_repo.fail_workspace_update(); - let before_count = conv_repo.conversation_count(); - - let err = svc + let started_at = std::time::Instant::now(); + let agent = svc .add_agent( "user1", &created.id, @@ -3261,25 +3701,53 @@ async fn add_agent_does_not_create_teammate_when_workspace_writeback_fails() { }, ) .await - .expect_err("workspace writeback failure must block teammate creation"); + .unwrap(); assert!( - err.to_string().contains("forced workspace writeback failure"), - "unexpected error: {err}" + started_at.elapsed() < std::time::Duration::from_millis(150), + "manual add HTTP path must not wait for delayed runtime attach" ); - assert_eq!(conv_repo.conversation_count(), before_count); + tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + if task_manager + .snapshot() + .build + .iter() + .any(|conversation_id| conversation_id == &agent.conversation_id) + { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } + }) + .await + .expect("manual add background attach should build the new runtime"); } #[tokio::test] -async fn add_agent_continues_when_team_temp_leader_patch_fails() { - let agent_metadata_repo: Arc = Arc::new(StubAgentMetadataRepo::empty()); - let (svc, team_repo, conversation_ports, conv_repo) = - setup_with_ports_team_repo_and_conversation_repo(success_factory(), agent_metadata_repo); +async fn manual_add_agent_attach_failure_marks_slot_error_and_notifies_leader() { + use futures_util::FutureExt; + + let fail_next = Arc::new(AtomicBool::new(false)); + let factory_fail_next = Arc::clone(&fail_next); + let factory: AgentFactory = Arc::new(move |opts: BuildTaskOptions| { + let should_fail = factory_fail_next.swap(false, Ordering::SeqCst); + async move { + if should_fail { + return Err(AgentError::internal("simulated manual add attach failure")); + } + Ok(aionui_ai_agent::AgentInstance::Mock(Arc::new( + mock_agent::MockAgent::new(opts.context.conversation.conversation_id, opts.context.workspace.path), + ))) + } + .boxed() + }); + let (svc, team_repo, task_manager, recorder) = setup_with_factory_and_recording_broadcaster(factory); let created = svc .create_team( "user1", CreateTeamRequest { - name: "Patch Failure".into(), + name: "T".into(), agents: vec![TeamAgentInput { name: "Lead".into(), role: "lead".into(), @@ -3293,19 +3761,13 @@ async fn add_agent_continues_when_team_temp_leader_patch_fails() { ) .await .unwrap(); + svc.ensure_session("user1", &created.id).await.unwrap(); + let original_scheduler = svc.get_session_scheduler(&created.id).unwrap(); + let lead_conversation_id = created.assistants[0].conversation_id.clone(); + task_manager.reset_calls(); + fail_next.store(true, Ordering::SeqCst); - force_team_workspace(&team_repo, &created.id, "").await; - conv_repo - .patch_extra( - &created.assistants[0].conversation_id, - serde_json::json!({ "workspace": "/tmp/aionui-team-missing-leader-workspace" }), - ) - .unwrap(); - conversation_ports - .fail_leader_workspace_patch - .store(true, std::sync::atomic::Ordering::SeqCst); - - let added = svc + let agent = svc .add_agent( "user1", &created.id, @@ -3320,50 +3782,598 @@ async fn add_agent_continues_when_team_temp_leader_patch_fails() { .await .unwrap(); - let got = svc.get_team("user1", &created.id).await.unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + let has_error_status = recorder.events_by_name("team.agentStatusChanged").iter().any(|event| { + event.data.get("slot_id").and_then(serde_json::Value::as_str) == Some(agent.slot_id.as_str()) + && event.data.get("status").and_then(serde_json::Value::as_str) == Some("error") + }); + if has_error_status { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } + }) + .await + .expect("manual add attach failure should mark the slot error"); + + tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + if recorder + .events_by_name("team.sessionStatusChanged") + .iter() + .any(|event| { + event.data.get("team_id").and_then(serde_json::Value::as_str) == Some(created.id.as_str()) + && event.data.get("status").and_then(serde_json::Value::as_str) == Some("failed") + && event.data.get("phase").and_then(serde_json::Value::as_str) == Some("attaching_agents") + }) + { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("dynamic attach failure must fail the team lifecycle"); + + assert!(Arc::ptr_eq( + &original_scheduler, + &svc.get_session_scheduler(&created.id) + .expect("published session remains") + )); assert!( - got.workspace - .contains(&format!("/conversations/team-temp-{}", created.id)) + task_manager.get_task(&lead_conversation_id).is_some(), + "healthy leader runtime must remain alive" ); - let added_extra = conv_repo.get_extra(&added.conversation_id).unwrap(); + assert!( + task_manager + .snapshot() + .build + .iter() + .all(|conversation_id| conversation_id == &agent.conversation_id), + "dynamic failure must not rebuild healthy members" + ); + + let lead_slot_id = created.leader_assistant_id.as_deref().expect("leader slot"); + let leader_messages = team_repo.get_history(&created.id, lead_slot_id, None).await.unwrap(); + assert!( + leader_messages + .iter() + .any(|message| message.content.contains("failed to attach its runtime")), + "leader should receive a persisted attach-failure notice" + ); + + svc.ensure_session("user1", &created.id) + .await + .expect("later ensure should retry only the failed member"); + assert!(Arc::ptr_eq( + &original_scheduler, + &svc.get_session_scheduler(&created.id) + .expect("same session after retry") + )); assert_eq!( - added_extra.get("workspace").and_then(serde_json::Value::as_str), - Some(got.workspace.as_str()) + task_manager + .snapshot() + .build + .iter() + .filter(|conversation_id| *conversation_id == &agent.conversation_id) + .count(), + 2, + "the failed member should have one failed attach and one later retry" + ); + assert!( + recorder + .events_by_name("team.sessionStatusChanged") + .iter() + .any(|event| { + event.data.get("team_id").and_then(serde_json::Value::as_str) == Some(created.id.as_str()) + && event.data.get("status").and_then(serde_json::Value::as_str) == Some("ready") + && event.data.get("server_count").and_then(serde_json::Value::as_u64) == Some(2) + }), + "single-member retry must restore team Ready" ); } #[tokio::test] -async fn provisioning_writes_typed_team_binding_for_create_and_add_agent() { - let agent_metadata_repo: Arc = Arc::new(StubAgentMetadataRepo::empty()); - let (svc, _, conv_repo) = - setup_with_factory_and_metadata_and_conversation_repo(success_factory(), agent_metadata_repo); +async fn failed_member_returns_conflict_and_removal_restores_ready() { + use futures_util::FutureExt; + + let build_count = Arc::new(AtomicUsize::new(0)); + let factory_count = Arc::clone(&build_count); + let factory: AgentFactory = Arc::new(move |opts: BuildTaskOptions| { + let build_index = factory_count.fetch_add(1, Ordering::SeqCst); + async move { + if build_index >= 1 { + return Err(AgentError::internal("provider-secret: dynamic attach failed")); + } + Ok(aionui_ai_agent::AgentInstance::Mock(Arc::new( + mock_agent::MockAgent::new(opts.context.conversation.conversation_id, opts.context.workspace.path), + ))) + } + .boxed() + }); + let (svc, _team_repo, task_manager, recorder) = setup_with_factory_and_recording_broadcaster(factory); let created = svc .create_team( "user1", CreateTeamRequest { - name: "Typed".into(), - agents: two_agent_input(), + name: "Failed member removal".into(), + agents: vec![TeamAgentInput { + name: "Lead".into(), + role: "lead".into(), + backend: Some("acp".into()), + model: "claude".into(), + assistant_id: None, + conversation_id: None, + }], workspace: None, }, ) .await .unwrap(); + svc.ensure_session("user1", &created.id).await.unwrap(); + let original_scheduler = svc.get_session_scheduler(&created.id).unwrap(); + let lead_conversation_id = created.assistants[0].conversation_id.clone(); + task_manager.reset_calls(); - for agent in &created.assistants { - let extra = conv_repo.get_extra(&agent.conversation_id).unwrap(); - assert_eq!(extra.get("teamId").and_then(|v| v.as_str()), Some(created.id.as_str())); - assert_eq!( - extra.get("slot_id").and_then(|v| v.as_str()), - Some(agent.slot_id.as_str()) - ); - assert_eq!(extra.get("role").and_then(|v| v.as_str()), Some(agent.role.as_str())); - assert_eq!( - extra.get("backend").and_then(|v| v.as_str()), - Some(agent.backend.as_str()) - ); - assert_eq!( - extra.get("session_mode").and_then(|v| v.as_str()), - Some("yolo"), + let failed = svc + .add_agent( + "user1", + &created.id, + AddAgentRequest { + name: "Broken".into(), + role: "teammate".into(), + backend: Some("acp".into()), + model: "claude".into(), + assistant_id: None, + }, + ) + .await + .unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + if recorder + .events_by_name("team.sessionStatusChanged") + .iter() + .any(|event| event.data.get("status").and_then(serde_json::Value::as_str) == Some("failed")) + { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("dynamic failure status"); + + let error = svc + .ensure_session("user1", &created.id) + .await + .expect_err("failed-member retry should report one deterministic failure"); + assert!(matches!( + error, + TeamError::MemberRuntimeFailed { + ref team_id, + ref slot_id, + ref conversation_id, + ref public_reason, + } if team_id == &created.id + && slot_id == &failed.slot_id + && conversation_id == &failed.conversation_id + && public_reason == "Agent runtime failed to start" + )); + + task_manager.reset_calls(); + recorder.clear(); + svc.remove_agent("user1", &created.id, &failed.slot_id).await.unwrap(); + + assert!(Arc::ptr_eq( + &original_scheduler, + &svc.get_session_scheduler(&created.id).expect("healthy session remains") + )); + assert!(task_manager.get_task(&lead_conversation_id).is_some()); + assert!( + task_manager + .snapshot() + .kill + .iter() + .all(|(conversation_id, _)| conversation_id != &lead_conversation_id) + ); + assert!( + recorder + .events_by_name("team.sessionStatusChanged") + .iter() + .any(|event| { + event.data.get("status").and_then(serde_json::Value::as_str) == Some("ready") + && event.data.get("server_count").and_then(serde_json::Value::as_u64) == Some(1) + }) + ); +} + +#[tokio::test] +async fn remove_during_attach_cancels_work_and_rejects_late_ready() { + let gate = Arc::new(GatedProvisioningFactory::default()); + let (svc, _team_repo, task_manager, recorder, conv_repo) = + setup_with_factory_recording_broadcaster_and_conversation_repo(gate.factory()); + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "Remove attaching member".into(), + agents: vec![TeamAgentInput { + name: "Lead".into(), + role: "lead".into(), + backend: Some("acp".into()), + model: "claude".into(), + assistant_id: None, + conversation_id: None, + }], + workspace: None, + }, + ) + .await + .unwrap(); + svc.ensure_session("user1", &created.id).await.unwrap(); + gate.enable(); + let added = svc + .add_agent( + "user1", + &created.id, + AddAgentRequest { + name: "Attaching".into(), + role: "teammate".into(), + backend: Some("acp".into()), + model: "claude".into(), + assistant_id: None, + }, + ) + .await + .unwrap(); + gate.wait_for_starts(1).await; + + let remove_service = Arc::clone(&svc); + let team_id = created.id.clone(); + let slot_id = added.slot_id.clone(); + let removal = tokio::spawn(async move { remove_service.remove_agent("user1", &team_id, &slot_id).await }); + tokio::task::yield_now().await; + assert!(!removal.is_finished(), "removal must join in-flight runtime cleanup"); + gate.release(1); + removal.await.unwrap().unwrap(); + + assert!(task_manager.get_task(&added.conversation_id).is_none()); + assert!( + conv_repo.get(&added.conversation_id).await.unwrap().is_none(), + "removed attaching member conversation must be deleted" + ); + assert!( + svc.get_team("user1", &created.id) + .await + .unwrap() + .assistants + .iter() + .all(|agent| agent.slot_id != added.slot_id) + ); + assert_eq!( + recorder + .events_by_name("team.agentRuntimeStatusChanged") + .into_iter() + .filter(|event| { + event.data.get("slot_id").and_then(serde_json::Value::as_str) == Some(added.slot_id.as_str()) + && event.data.get("status").and_then(serde_json::Value::as_str) == Some("ready") + }) + .count(), + 0, + "cancelled attach must not publish late Ready" + ); +} + +#[tokio::test] +async fn aa_add_agent_inherits_team_workspace() { + let agent_metadata_repo: Arc = Arc::new(StubAgentMetadataRepo::empty()); + let (svc, _, conv_repo) = + setup_with_factory_and_metadata_and_conversation_repo(success_factory(), agent_metadata_repo); + let workspace = std::env::temp_dir().join(format!("aionui-team-workspace-{}", aionui_common::generate_id())); + std::fs::create_dir_all(&workspace).unwrap(); + let workspace = workspace.to_string_lossy().into_owned(); + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "T".into(), + agents: vec![TeamAgentInput { + name: "Lead".into(), + role: "lead".into(), + backend: Some("acp".into()), + model: "claude".into(), + assistant_id: None, + conversation_id: None, + }], + workspace: Some(workspace.clone()), + }, + ) + .await + .unwrap(); + + let agent = svc + .add_agent( + "user1", + &created.id, + AddAgentRequest { + name: "Worker".into(), + role: "teammate".into(), + backend: Some("acp".into()), + model: "claude".into(), + assistant_id: None, + }, + ) + .await + .unwrap(); + + let extra = conv_repo.get_extra(&agent.conversation_id).unwrap(); + assert_eq!( + extra.get("workspace").and_then(|v| v.as_str()), + Some(workspace.as_str()) + ); +} + +#[tokio::test] +async fn add_agent_backfills_empty_team_workspace_from_leader_workspace() { + let agent_metadata_repo: Arc = Arc::new(StubAgentMetadataRepo::empty()); + let (svc, team_repo, _, conv_repo) = + setup_with_factory_metadata_team_repo_and_conversation_repo(success_factory(), agent_metadata_repo); + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "Legacy".into(), + agents: vec![TeamAgentInput { + name: "Lead".into(), + role: "lead".into(), + backend: Some("acp".into()), + model: "claude".into(), + assistant_id: None, + conversation_id: None, + }], + workspace: None, + }, + ) + .await + .unwrap(); + let leader_workspace = conv_repo.get_extra(&created.assistants[0].conversation_id).unwrap()["workspace"] + .as_str() + .unwrap() + .to_owned(); + + force_team_workspace(&team_repo, &created.id, "").await; + + let added = svc + .add_agent( + "user1", + &created.id, + AddAgentRequest { + name: "Worker".into(), + role: "teammate".into(), + backend: Some("acp".into()), + model: "claude".into(), + assistant_id: None, + }, + ) + .await + .unwrap(); + + let got = svc.get_team("user1", &created.id).await.unwrap(); + assert_eq!(got.workspace, leader_workspace); + let added_extra = conv_repo.get_extra(&added.conversation_id).unwrap(); + assert_eq!( + added_extra.get("workspace").and_then(serde_json::Value::as_str), + Some(leader_workspace.as_str()) + ); +} + +#[tokio::test] +async fn add_agent_uses_team_temp_workspace_when_team_and_leader_workspaces_are_unusable() { + let agent_metadata_repo: Arc = Arc::new(StubAgentMetadataRepo::empty()); + let (svc, team_repo, _, conv_repo) = + setup_with_factory_metadata_team_repo_and_conversation_repo(success_factory(), agent_metadata_repo); + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "Legacy Empty".into(), + agents: vec![TeamAgentInput { + name: "Lead".into(), + role: "lead".into(), + backend: Some("acp".into()), + model: "claude".into(), + assistant_id: None, + conversation_id: None, + }], + workspace: None, + }, + ) + .await + .unwrap(); + + force_team_workspace(&team_repo, &created.id, "").await; + conv_repo + .patch_extra( + &created.assistants[0].conversation_id, + serde_json::json!({ "workspace": "/tmp/aionui-team-missing-leader-workspace" }), + ) + .unwrap(); + + let added = svc + .add_agent( + "user1", + &created.id, + AddAgentRequest { + name: "Worker".into(), + role: "teammate".into(), + backend: Some("acp".into()), + model: "claude".into(), + assistant_id: None, + }, + ) + .await + .unwrap(); + + let got = svc.get_team("user1", &created.id).await.unwrap(); + assert!( + got.workspace + .contains(&format!("/conversations/team-temp-{}", created.id)), + "unexpected team temp workspace: {}", + got.workspace + ); + let added_extra = conv_repo.get_extra(&added.conversation_id).unwrap(); + assert_eq!( + added_extra.get("workspace").and_then(serde_json::Value::as_str), + Some(got.workspace.as_str()) + ); +} + +#[tokio::test] +async fn add_agent_does_not_create_teammate_when_workspace_writeback_fails() { + let agent_metadata_repo: Arc = Arc::new(StubAgentMetadataRepo::empty()); + let (svc, team_repo, _, conv_repo) = + setup_with_factory_metadata_team_repo_and_conversation_repo(success_factory(), agent_metadata_repo); + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "Writeback Failure".into(), + agents: vec![TeamAgentInput { + name: "Lead".into(), + role: "lead".into(), + backend: Some("acp".into()), + model: "claude".into(), + assistant_id: None, + conversation_id: None, + }], + workspace: None, + }, + ) + .await + .unwrap(); + + force_team_workspace(&team_repo, &created.id, "").await; + team_repo.fail_workspace_update(); + let before_count = conv_repo.conversation_count(); + + let err = svc + .add_agent( + "user1", + &created.id, + AddAgentRequest { + name: "Worker".into(), + role: "teammate".into(), + backend: Some("acp".into()), + model: "claude".into(), + assistant_id: None, + }, + ) + .await + .expect_err("workspace writeback failure must block teammate creation"); + + assert!( + err.to_string().contains("forced workspace writeback failure"), + "unexpected error: {err}" + ); + assert_eq!(conv_repo.conversation_count(), before_count); +} + +#[tokio::test] +async fn add_agent_continues_when_team_temp_leader_patch_fails() { + let agent_metadata_repo: Arc = Arc::new(StubAgentMetadataRepo::empty()); + let (svc, team_repo, conversation_ports, conv_repo) = + setup_with_ports_team_repo_and_conversation_repo(success_factory(), agent_metadata_repo); + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "Patch Failure".into(), + agents: vec![TeamAgentInput { + name: "Lead".into(), + role: "lead".into(), + backend: Some("acp".into()), + model: "claude".into(), + assistant_id: None, + conversation_id: None, + }], + workspace: None, + }, + ) + .await + .unwrap(); + + force_team_workspace(&team_repo, &created.id, "").await; + conv_repo + .patch_extra( + &created.assistants[0].conversation_id, + serde_json::json!({ "workspace": "/tmp/aionui-team-missing-leader-workspace" }), + ) + .unwrap(); + conversation_ports + .fail_leader_workspace_patch + .store(true, std::sync::atomic::Ordering::SeqCst); + + let added = svc + .add_agent( + "user1", + &created.id, + AddAgentRequest { + name: "Worker".into(), + role: "teammate".into(), + backend: Some("acp".into()), + model: "claude".into(), + assistant_id: None, + }, + ) + .await + .unwrap(); + + let got = svc.get_team("user1", &created.id).await.unwrap(); + assert!( + got.workspace + .contains(&format!("/conversations/team-temp-{}", created.id)) + ); + let added_extra = conv_repo.get_extra(&added.conversation_id).unwrap(); + assert_eq!( + added_extra.get("workspace").and_then(serde_json::Value::as_str), + Some(got.workspace.as_str()) + ); +} + +#[tokio::test] +async fn provisioning_writes_typed_team_binding_for_create_and_add_agent() { + let agent_metadata_repo: Arc = Arc::new(StubAgentMetadataRepo::empty()); + let (svc, _, conv_repo) = + setup_with_factory_and_metadata_and_conversation_repo(success_factory(), agent_metadata_repo); + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "Typed".into(), + agents: two_agent_input(), + workspace: None, + }, + ) + .await + .unwrap(); + + for agent in &created.assistants { + let extra = conv_repo.get_extra(&agent.conversation_id).unwrap(); + assert_eq!(extra.get("teamId").and_then(|v| v.as_str()), Some(created.id.as_str())); + assert_eq!( + extra.get("slot_id").and_then(|v| v.as_str()), + Some(agent.slot_id.as_str()) + ); + assert_eq!(extra.get("role").and_then(|v| v.as_str()), Some(agent.role.as_str())); + assert_eq!( + extra.get("backend").and_then(|v| v.as_str()), + Some(agent.backend.as_str()) + ); + assert_eq!( + extra.get("session_mode").and_then(|v| v.as_str()), + Some("yolo"), "Team provisioning should write the runtime seed for initial agents" ); } @@ -3495,6 +4505,153 @@ async fn ar1_remove_agent_from_team() { assert!(got.assistants.iter().all(|a| a.slot_id != worker_slot)); } +#[tokio::test] +async fn membership_persist_failure_does_not_delete_the_conversation() { + let (svc, team_repo, task_manager, conv_repo) = setup_with_factory_metadata_team_repo_and_conversation_repo( + success_factory(), + Arc::new(StubAgentMetadataRepo::empty()), + ); + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "Removal persistence failure".into(), + agents: two_agent_input(), + workspace: None, + }, + ) + .await + .unwrap(); + svc.ensure_session("user1", &created.id).await.unwrap(); + let worker = created.assistants[1].clone(); + team_repo.fail_agent_updates(); + + let error = svc + .remove_agent("user1", &created.id, &worker.slot_id) + .await + .expect_err("forced membership persistence failure must be returned"); + + assert!(error.to_string().contains("forced agent update failure")); + assert!( + conv_repo.get(&worker.conversation_id).await.unwrap().is_some(), + "conversation deletion must happen only after membership persistence" + ); + assert!( + svc.get_team("user1", &created.id) + .await + .unwrap() + .assistants + .iter() + .any(|agent| agent.slot_id == worker.slot_id), + "failed persistence must leave declared membership intact" + ); + assert!( + task_manager.get_task(&worker.conversation_id).is_none(), + "cancelled runtime must be cleaned before reporting persistence failure" + ); +} + +#[tokio::test] +async fn remove_tolerates_current_session_already_missing_the_slot() { + let (svc, _team_repo, _task_manager, conv_repo) = setup_with_factory_metadata_team_repo_and_conversation_repo( + success_factory(), + Arc::new(StubAgentMetadataRepo::empty()), + ); + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "Already absent runtime slot".into(), + agents: two_agent_input(), + workspace: None, + }, + ) + .await + .unwrap(); + svc.ensure_session("user1", &created.id).await.unwrap(); + let worker = created.assistants[1].clone(); + svc.get_session_scheduler(&created.id) + .unwrap() + .remove_agent(&worker.slot_id) + .await + .unwrap(); + + svc.remove_agent("user1", &created.id, &worker.slot_id) + .await + .expect("post-persistence runtime cleanup must be idempotent"); + + assert!(conv_repo.get(&worker.conversation_id).await.unwrap().is_none()); + assert!( + svc.get_team("user1", &created.id) + .await + .unwrap() + .assistants + .iter() + .all(|agent| agent.slot_id != worker.slot_id) + ); +} + +#[tokio::test] +async fn manual_remove_agent_projects_team_system_message_without_active_team_run() { + let (svc, _, conv_repo) = setup_with_factory_and_metadata_and_conversation_repo( + success_factory(), + Arc::new(StubAgentMetadataRepo::empty()), + ); + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "T".into(), + agents: two_agent_input(), + workspace: None, + }, + ) + .await + .unwrap(); + svc.ensure_session("user1", &created.id).await.unwrap(); + let leader_conversation_id = created.assistants[0].conversation_id.clone(); + let worker_slot = created.assistants[1].slot_id.clone(); + + svc.remove_agent("user1", &created.id, &worker_slot).await.unwrap(); + + let leader_messages = conv_repo.messages_for(&leader_conversation_id); + assert_eq!(leader_messages.len(), 1); + assert_eq!(leader_messages[0].position.as_deref(), Some("left")); + let content: serde_json::Value = serde_json::from_str(&leader_messages[0].content).unwrap(); + assert_eq!(content["sender_name"], "team_system"); + assert_eq!(content["teammate_message"], true); + assert!( + content["content"] + .as_str() + .unwrap() + .contains("was removed from the team") + ); + assert!(content["content"].as_str().unwrap().contains(&worker_slot)); +} + +#[tokio::test] +async fn remove_agent_rejects_leader() { + let svc = setup(); + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "T".into(), + agents: two_agent_input(), + workspace: None, + }, + ) + .await + .unwrap(); + + let lead_slot = created.assistants[0].slot_id.clone(); + let result = svc.remove_agent("user1", &created.id, &lead_slot).await; + + assert!(matches!(result, Err(TeamError::InvalidRequest(_)))); + let got = svc.get_team("user1", &created.id).await.unwrap(); + assert!(got.assistants.iter().any(|agent| agent.slot_id == lead_slot)); +} + #[tokio::test] async fn ar4_remove_nonexistent_agent() { let svc = setup(); @@ -3661,7 +4818,6 @@ async fn spawn_agent_in_session_succeeds_without_active_team_run() { let req = SpawnAgentRequest { name: "Helper".into(), assistant_id: Some("assistant-worker".into()), - model: Some("claude-sonnet-4".into()), }; let spawned = svc @@ -3686,6 +4842,74 @@ async fn spawn_agent_in_session_succeeds_without_active_team_run() { ); } +#[tokio::test] +async fn leader_spawn_then_immediate_ensure_joins_the_same_attach_operation() { + let gate = Arc::new(GatedProvisioningFactory::default()); + let (svc, _team_repo, task_manager, _conv_repo) = setup_with_factory_metadata_assistants_and_conversation_repo( + gate.factory(), + seeded_agent_metadata_repo(), + Arc::new(SingleAssistantDefinitionRepo { + row: word_creator_definition(), + }), + Arc::new(EmptyAssistantOverlayRepo), + ); + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "Leader spawn reconciliation".into(), + agents: two_agent_input(), + workspace: None, + }, + ) + .await + .unwrap(); + svc.ensure_session("user1", &created.id).await.unwrap(); + let original_scheduler = svc.get_session_scheduler(&created.id).unwrap(); + task_manager.reset_calls(); + gate.enable(); + + let spawned = svc + .spawn_agent_in_session( + &created.id, + created.leader_assistant_id.as_deref().unwrap(), + SpawnAgentRequest { + name: "Writer".into(), + assistant_id: Some("word-creator".into()), + }, + ) + .await + .unwrap(); + gate.wait_for_starts(1).await; + + let ensure_service = Arc::clone(&svc); + let team_id = created.id.clone(); + let ensure = tokio::spawn(async move { ensure_service.ensure_session("user1", &team_id).await }); + tokio::task::yield_now().await; + assert!( + !ensure.is_finished(), + "ensure must join the in-flight Leader spawn attach" + ); + assert_eq!(gate.starts(), vec![spawned.conversation_id.clone()]); + assert_eq!(task_manager.snapshot().build, vec![spawned.conversation_id.clone()]); + assert!( + task_manager + .snapshot() + .kill + .iter() + .all(|(conversation_id, _)| conversation_id == &spawned.conversation_id), + "Leader spawn reconciliation must not restart healthy members" + ); + assert!(Arc::ptr_eq( + &original_scheduler, + &svc.get_session_scheduler(&created.id).unwrap() + )); + + gate.release(1); + ensure.await.unwrap().unwrap(); + assert_eq!(gate.starts(), vec![spawned.conversation_id]); +} + #[tokio::test] async fn lead_send_agent_message_in_session_requires_active_team_run() { let svc = setup(); @@ -3754,7 +4978,6 @@ async fn spawn_agent_in_session_aborts_lease_when_persistence_fails() { let req = SpawnAgentRequest { name: "Helper".into(), assistant_id: Some("word-creator".into()), - model: Some("claude-sonnet-4".into()), }; let err = svc @@ -3801,7 +5024,6 @@ async fn spawn_agent_in_session_compensates_when_welcome_mailbox_write_fails() { let req = SpawnAgentRequest { name: "Helper".into(), assistant_id: Some("word-creator".into()), - model: Some("claude-sonnet-4".into()), }; let err = svc @@ -3863,33 +5085,104 @@ async fn es4_ensure_session_rejects_cross_user_access() { assert!(matches!(result, Err(aionui_team::TeamError::Forbidden(_)))); } -// -- W5-D31b-2: team.mcpStatus service-layer broadcasts --------------------- +// -- W5-D31b-2: team.sessionStatusChanged service-layer broadcasts ----------- // // The happy-path phase transitions are covered by focused service/session // assertions. This test keeps the load-failed broadcast covered end-to-end. -#[tokio::test] -async fn d31b2_ensure_session_broadcasts_load_failed_for_missing_team() { - let (svc, recorder) = setup_with_recording_broadcaster(); - let err = svc.ensure_session("user1", "nonexistent-team-xyz").await.unwrap_err(); - assert!(matches!(err, aionui_team::TeamError::TeamNotFound(_))); +#[tokio::test] +async fn d31b2_ensure_session_broadcasts_failed_loading_team_for_missing_team() { + let (svc, recorder) = setup_with_recording_broadcaster(); + let err = svc.ensure_session("user1", "nonexistent-team-xyz").await.unwrap_err(); + assert!(matches!(err, aionui_team::TeamError::TeamNotFound(_))); + + let failed = recorder + .events_by_name("team.sessionStatusChanged") + .into_iter() + .find(|e| { + e.data.get("status").and_then(|v| v.as_str()) == Some("failed") + && e.data.get("phase").and_then(|v| v.as_str()) == Some("loading_team") + }) + .expect("failed/loading_team broadcast expected"); + assert_eq!( + failed.data.get("team_id").and_then(|v| v.as_str()), + Some("nonexistent-team-xyz") + ); + assert!(failed.data.get("error").is_some()); +} + +#[tokio::test] +async fn ensure_session_broadcasts_starting_and_ready_session_status() { + let (svc, recorder) = setup_with_recording_broadcaster(); + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "T".into(), + agents: two_agent_input(), + workspace: None, + }, + ) + .await + .unwrap(); + + svc.ensure_session("user1", &created.id).await.unwrap(); + + let events = recorder.events_by_name("team.sessionStatusChanged"); + assert!( + events.iter().any(|event| { + event.data.get("team_id").and_then(|v| v.as_str()) == Some(created.id.as_str()) + && event.data.get("status").and_then(|v| v.as_str()) == Some("starting") + && event.data.get("phase").and_then(|v| v.as_str()) == Some("loading_team") + }), + "ensure_session must emit starting/loading_team" + ); + assert!( + events.iter().any(|event| { + event.data.get("team_id").and_then(|v| v.as_str()) == Some(created.id.as_str()) + && event.data.get("status").and_then(|v| v.as_str()) == Some("ready") + && event.data.get("server_count").and_then(|v| v.as_u64()) == Some(2) + }), + "ensure_session must emit ready with server_count" + ); +} + +#[tokio::test] +async fn ensure_session_existing_ready_session_broadcasts_ready_terminal_status() { + let (svc, recorder) = setup_with_recording_broadcaster(); + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "T".into(), + agents: two_agent_input(), + workspace: None, + }, + ) + .await + .unwrap(); + + svc.ensure_session("user1", &created.id).await.unwrap(); + recorder.clear(); + + svc.ensure_session("user1", &created.id).await.unwrap(); - let load_failed = recorder - .events_by_name("team.mcpStatus") - .into_iter() - .find(|e| { - e.data - .get("phase") - .and_then(|v| v.as_str()) - .map(|s| s == "load_failed") - .unwrap_or(false) - }) - .expect("load_failed broadcast expected"); - assert_eq!( - load_failed.data.get("team_id").and_then(|v| v.as_str()), - Some("nonexistent-team-xyz") + let events = recorder.events_by_name("team.sessionStatusChanged"); + assert!( + events.iter().all(|event| { + event.data.get("team_id").and_then(|v| v.as_str()) != Some(created.id.as_str()) + || event.data.get("status").and_then(|v| v.as_str()) != Some("starting") + }), + "existing ready session fast path must not emit starting because no lifecycle transition is happening" + ); + assert!( + events.iter().any(|event| { + event.data.get("team_id").and_then(|v| v.as_str()) == Some(created.id.as_str()) + && event.data.get("status").and_then(|v| v.as_str()) == Some("ready") + && event.data.get("server_count").and_then(|v| v.as_u64()) == Some(2) + }), + "existing ready session fast path must emit a ready terminal status" ); - assert!(load_failed.data.get("error").is_some()); } #[tokio::test] @@ -4090,57 +5383,498 @@ async fn dispose_all_cleans_up_sessions() { .create_team( "user1", CreateTeamRequest { - name: "B".into(), + name: "B".into(), + agents: two_agent_input(), + workspace: None, + }, + ) + .await + .unwrap(); + + svc.ensure_session("user1", &t1.id).await.unwrap(); + svc.ensure_session("user1", &t2.id).await.unwrap(); + + svc.dispose_all(); + + // After dispose, sessions are cleaned up. + assert!(svc.get_session_scheduler(&t1.id).is_none()); + assert!(svc.get_session_scheduler(&t2.id).is_none()); +} + +// =========================================================================== +// Test: Delete team stops active session (TD-2 + integration) +// =========================================================================== + +#[tokio::test] +async fn td_delete_team_stops_session() { + let svc = setup(); + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "T".into(), + agents: two_agent_input(), + workspace: None, + }, + ) + .await + .unwrap(); + + svc.ensure_session("user1", &created.id).await.unwrap(); + svc.remove_team("user1", &created.id).await.unwrap(); + + let result = svc.send_message("user1", &created.id, "Hello", None).await; + assert!(result.is_err()); +} + +// =========================================================================== +// Test: D9 ensure_session kill + rebuild closed loop +// =========================================================================== + +#[tokio::test] +async fn d9_create_team_persists_without_warming_initial_agents() { + let (svc, tm) = setup_with_factory(success_factory()); + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "T".into(), + agents: two_agent_input(), + workspace: None, + }, + ) + .await + .unwrap(); + + let calls = tm.snapshot(); + assert_eq!( + created.assistants.len(), + 2, + "create_team should still persist initial agents" + ); + assert!(calls.kill.is_empty(), "create_team must not rebuild initial agents"); + assert!(calls.build.is_empty(), "create_team must not warm initial agents"); + assert_eq!(tm.active_count(), 0, "create_team must not register live agent tasks"); + assert!( + svc.get_session_scheduler(&created.id).is_none(), + "create_team must not create a runtime session" + ); +} + +#[tokio::test] +async fn d9_ensure_session_kills_and_rebuilds_every_agent() { + let (svc, tm) = setup_with_factory(success_factory()); + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "T".into(), + agents: two_agent_input(), + workspace: None, + }, + ) + .await + .unwrap(); + + reset_runtime_state(&svc, &tm, &created.id).await; + svc.ensure_session("user1", &created.id).await.unwrap(); + + // Two agents → kill called 2x and get_or_build_task called 2x, each with + // the corresponding conversation_id. Order is agents-iteration order. + let calls = tm.snapshot(); + assert_eq!(calls.kill.len(), 2, "expected 2 kill calls"); + assert_eq!(calls.build.len(), 2, "expected 2 build calls"); + for (i, agent) in created.assistants.iter().enumerate() { + assert_eq!(calls.kill[i].0, agent.conversation_id); + assert_eq!(calls.kill[i].1, Some(AgentKillReason::TeamMcpRebuild)); + assert_eq!(calls.build[i], agent.conversation_id); + } +} + +#[tokio::test(start_paused = true)] +async fn d9_ensure_session_rebuilds_agents_with_staggered_bounded_parallelism() { + let probe = Arc::new(WarmupConcurrencyProbe::default()); + let (svc, _tm) = setup_with_factory(probe.factory(std::time::Duration::from_secs(20))); + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "T".into(), + agents: five_agent_input_leader_not_first(), + workspace: None, + }, + ) + .await + .unwrap(); + let mut expected_starts = Vec::new(); + expected_starts.extend( + created + .assistants + .iter() + .filter(|assistant| assistant.role == "lead") + .map(|assistant| assistant.conversation_id.clone()), + ); + expected_starts.extend( + created + .assistants + .iter() + .filter(|assistant| assistant.role != "lead") + .map(|assistant| assistant.conversation_id.clone()), + ); + + let svc_for_task = Arc::clone(&svc); + let team_id = created.id.clone(); + let handle = tokio::spawn(async move { svc_for_task.ensure_session("user1", &team_id).await }); + + tokio::time::advance(std::time::Duration::from_secs(120)).await; + handle.await.unwrap().unwrap(); + + let starts = probe.starts(); + assert_eq!( + starts, expected_starts, + "team rebuild warmup must start leader first and preserve teammate order" + ); + assert!( + probe.max_active() > 1, + "team rebuild warmup should overlap staggered agents when warmup takes longer than the launch interval" + ); + assert_eq!( + probe.max_active(), + 3, + "team rebuild warmup should cap concurrent agents at 3" + ); + let start_times = probe.start_times(); + assert_eq!(start_times.len(), expected_starts.len()); + for pair in start_times.windows(2).take(2) { + let delta = pair[1].1.saturating_sub(pair[0].1); + assert!( + delta >= std::time::Duration::from_secs(3), + "agent starts should be staggered by at least 3s; observed {delta:?}" + ); + assert!( + delta < std::time::Duration::from_secs(5), + "agent starts should use the configured 3s stagger, not the old 5s interval; observed {delta:?}" + ); + } +} + +#[tokio::test] +async fn d9_ensure_session_persists_team_mcp_stdio_config() { + // Each agent's conversation.extra must carry a `team_mcp_stdio_config` + // object by the time the factory is called — that is what the rebuilt + // typed Team context will expose to reach the MCP server. + use futures_util::FutureExt; + let (svc, _tm) = setup_with_factory(Arc::new(|opts: BuildTaskOptions| { + async move { + let context = opts.context; + let typed_has_cfg = context + .team + .as_ref() + .and_then(|team| team.mcp.as_ref()) + .is_some_and(|mcp| mcp.stdio.port > 0 && !mcp.stdio.slot_id.is_empty()); + let compat_has_cfg = match &context.kind { + AgentSessionKind::Acp(acp) => { + assert!(acp.belongs_to_team); + assert!(acp.team.is_some(), "ACP build context must carry typed team binding"); + acp.config + .team_mcp_stdio_config + .as_ref() + .is_some_and(|cfg| cfg.port > 0 && !cfg.slot_id.is_empty()) + } + _ => false, + }; + assert!( + typed_has_cfg && compat_has_cfg, + "factory called without typed team_mcp_stdio_config in context: {:?}", + context.team + ); + Ok(aionui_ai_agent::AgentInstance::Mock(Arc::new( + mock_agent::MockAgent::new(context.conversation.conversation_id, context.workspace.path), + ))) + } + .boxed() + })); + + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "T".into(), + agents: two_agent_input(), + workspace: None, + }, + ) + .await + .unwrap(); + + svc.ensure_session("user1", &created.id).await.unwrap(); +} + +#[tokio::test] +async fn d9_ensure_session_is_idempotent() { + let (svc, tm) = setup_with_factory(success_factory()); + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "T".into(), + agents: two_agent_input(), + workspace: None, + }, + ) + .await + .unwrap(); + + reset_runtime_state(&svc, &tm, &created.id).await; + svc.ensure_session("user1", &created.id).await.unwrap(); + svc.ensure_session("user1", &created.id).await.unwrap(); + + // Second call short-circuits — no additional kill/build calls. + let calls = tm.snapshot(); + assert_eq!(calls.kill.len(), 2, "second ensure_session must not re-kill"); + assert_eq!(calls.build.len(), 2, "second ensure_session must not re-build"); +} + +#[tokio::test] +async fn manual_add_then_immediate_ensure_joins_attach_without_rebuilding_session() { + let gate = Arc::new(GatedProvisioningFactory::default()); + let (svc, _repo, task_manager, _recorder) = setup_with_factory_and_recording_broadcaster(gate.factory()); + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "Join dynamic attach".into(), + agents: vec![team_agent_input("Lead", "lead", "claude")], + workspace: None, + }, + ) + .await + .unwrap(); + svc.ensure_session("user1", &created.id).await.unwrap(); + let original_scheduler = svc.get_session_scheduler(&created.id).expect("published session"); + let lead_conversation_id = created.assistants[0].conversation_id.clone(); + task_manager.reset_calls(); + gate.enable(); + + let added = svc + .add_agent( + "user1", + &created.id, + AddAgentRequest { + name: "Worker".into(), + role: "teammate".into(), + backend: Some("acp".into()), + model: "claude".into(), + assistant_id: None, + }, + ) + .await + .unwrap(); + gate.wait_for_starts(1).await; + + let svc_for_ensure = Arc::clone(&svc); + let team_id = created.id.clone(); + let ensure = tokio::spawn(async move { svc_for_ensure.ensure_session("user1", &team_id).await }); + tokio::task::yield_now().await; + assert!(!ensure.is_finished(), "ensure must join the pending dynamic attach"); + assert_eq!(gate.starts(), vec![added.conversation_id.clone()]); + assert!(Arc::ptr_eq( + &original_scheduler, + &svc.get_session_scheduler(&created.id).expect("same published session") + )); + assert!( + task_manager + .snapshot() + .kill + .iter() + .all(|(conversation_id, _)| conversation_id != &lead_conversation_id), + "joining a dynamic attach must not kill a healthy member" + ); + + gate.release(1); + ensure.await.unwrap().unwrap(); + assert_eq!( + task_manager + .snapshot() + .build + .iter() + .filter(|conversation_id| *conversation_id == &added.conversation_id) + .count(), + 1, + "the dynamic conversation must be built once" + ); +} + +#[tokio::test] +async fn concurrent_ensures_launch_one_dynamic_attach() { + let gate = Arc::new(GatedProvisioningFactory::default()); + let (svc, task_manager) = setup_with_factory(gate.factory()); + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "Concurrent repair".into(), agents: two_agent_input(), workspace: None, }, ) .await .unwrap(); + svc.ensure_session("user1", &created.id).await.unwrap(); + let original_scheduler = svc.get_session_scheduler(&created.id).expect("published session"); + let lead = created.assistants.iter().find(|agent| agent.role == "lead").unwrap(); + let worker = created + .assistants + .iter() + .find(|agent| agent.role == "teammate") + .unwrap(); + task_manager + .remove_task_without_recording(&worker.conversation_id) + .await; + task_manager.reset_calls(); + gate.enable(); + + let barrier = Arc::new(tokio::sync::Barrier::new(3)); + let mut handles = Vec::new(); + for _ in 0..2 { + let svc = Arc::clone(&svc); + let team_id = created.id.clone(); + let barrier = Arc::clone(&barrier); + handles.push(tokio::spawn(async move { + barrier.wait().await; + svc.ensure_session("user1", &team_id).await + })); + } + barrier.wait().await; + gate.wait_for_starts(1).await; + tokio::task::yield_now().await; + assert!(handles.iter().all(|handle| !handle.is_finished())); + assert_eq!(gate.starts(), vec![worker.conversation_id.clone()]); + assert_eq!(task_manager.snapshot().build, vec![worker.conversation_id.clone()]); + assert!( + task_manager + .snapshot() + .kill + .iter() + .all(|(conversation_id, _)| conversation_id != &lead.conversation_id), + "healthy members must not be killed during a one-slot repair" + ); + assert!(Arc::ptr_eq( + &original_scheduler, + &svc.get_session_scheduler(&created.id).expect("same published session") + )); - svc.ensure_session("user1", &t1.id).await.unwrap(); - svc.ensure_session("user1", &t2.id).await.unwrap(); - - svc.dispose_all(); - - // After dispose, sessions are cleaned up. - assert!(svc.get_session_scheduler(&t1.id).is_none()); - assert!(svc.get_session_scheduler(&t2.id).is_none()); + gate.release(1); + for handle in handles { + handle.await.unwrap().unwrap(); + } + assert_eq!(gate.starts(), vec![worker.conversation_id.clone()]); } -// =========================================================================== -// Test: Delete team stops active session (TD-2 + integration) -// =========================================================================== - #[tokio::test] -async fn td_delete_team_stops_session() { - let svc = setup(); +async fn stopped_session_rejects_late_attach_completion() { + let gate = Arc::new(GatedProvisioningFactory::default()); + let (svc, _repo, task_manager, recorder) = setup_with_factory_and_recording_broadcaster(gate.factory()); let created = svc .create_team( "user1", CreateTeamRequest { - name: "T".into(), - agents: two_agent_input(), + name: "Stopped late attach".into(), + agents: vec![team_agent_input("Lead", "lead", "claude")], workspace: None, }, ) .await .unwrap(); - svc.ensure_session("user1", &created.id).await.unwrap(); - svc.remove_team("user1", &created.id).await.unwrap(); + task_manager.reset_calls(); + gate.enable(); - let result = svc.send_message("user1", &created.id, "Hello", None).await; - assert!(result.is_err()); -} + let added = svc + .add_agent( + "user1", + &created.id, + AddAgentRequest { + name: "Worker".into(), + role: "teammate".into(), + backend: Some("acp".into()), + model: "claude".into(), + assistant_id: None, + }, + ) + .await + .unwrap(); + gate.wait_for_starts(1).await; + svc.stop_session("user1", &created.id).await.unwrap(); + let svc_for_replacement = Arc::clone(&svc); + let replacement_team_id = created.id.clone(); + let replacement = + tokio::spawn(async move { svc_for_replacement.ensure_session("user1", &replacement_team_id).await }); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + let replacement_kill_started = task_manager + .snapshot() + .kill + .iter() + .filter(|(conversation_id, _)| conversation_id == &added.conversation_id) + .count(); + if replacement_kill_started >= 2 { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("replacement bootstrap must begin replacing the same member before old attach release"); + assert!(!replacement.is_finished()); + gate.release(1); + replacement.await.unwrap().unwrap(); -// =========================================================================== -// Test: D9 ensure_session kill + rebuild closed loop -// =========================================================================== + tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + let cleanup_kills = task_manager + .snapshot() + .kill + .iter() + .filter(|(conversation_id, _)| conversation_id == &added.conversation_id) + .count(); + if cleanup_kills >= 2 { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("late completion must clean up its partial task"); + + let added_events = recorder + .events_by_name("team.agentRuntimeStatusChanged") + .into_iter() + .filter(|event| event.data.get("slot_id").and_then(serde_json::Value::as_str) == Some(added.slot_id.as_str())) + .collect::>(); + assert_eq!( + added_events + .iter() + .filter(|event| event.data.get("status").and_then(serde_json::Value::as_str) == Some("ready")) + .count(), + 1, + "only the replacement session may publish Ready for the member" + ); + assert!( + task_manager.get_task(&added.conversation_id).is_some(), + "stale cleanup must not kill the replacement session runtime" + ); + assert!(svc.get_session_scheduler(&created.id).is_some()); +} #[tokio::test] -async fn d9_create_team_auto_start_rebuilds_every_initial_agent() { - let (svc, tm) = setup_with_factory(success_factory()); +async fn d9_ensure_session_rollbacks_when_build_fails() { + // Factory always fails → ensure_session must propagate error and not + // insert into sessions, so send_message afterwards still errors. + use futures_util::FutureExt; + let failing_factory: AgentFactory = + Arc::new(|_opts: BuildTaskOptions| async move { Err(AgentError::internal("simulated build failure")) }.boxed()); + let (svc, tm) = setup_with_factory(failing_factory); let created = svc .create_team( "user1", @@ -4153,106 +5887,197 @@ async fn d9_create_team_auto_start_rebuilds_every_initial_agent() { .await .unwrap(); + reset_runtime_state(&svc, &tm, &created.id).await; + let result = svc.ensure_session("user1", &created.id).await; + assert!(result.is_err(), "ensure_session should propagate build error"); + + // Serial rebuild stops at the first failing agent, and no session is + // inserted after the failure. let calls = tm.snapshot(); assert_eq!( calls.kill.len(), - 2, - "auto-start should kill every initial agent before rebuild" + 3, + "failed bootstrap cleans the full two-member snapshot" + ); + assert_eq!(calls.build.len(), 1); + + let send_result = svc.send_message("user1", &created.id, "Hello", None).await; + assert!( + send_result.is_err(), + "session must not be registered after build failure" ); - assert_eq!(calls.build.len(), 2, "auto-start should warm every initial agent"); - for (i, agent) in created.assistants.iter().enumerate() { - assert_eq!(calls.kill[i].0, agent.conversation_id); - assert_eq!(calls.kill[i].1, Some(AgentKillReason::TeamMcpRebuild)); - assert_eq!(calls.build[i], agent.conversation_id); - } - assert_eq!(tm.active_count(), 2, "auto-start must register every initial agent"); } #[tokio::test] -async fn d9_ensure_session_kills_and_rebuilds_every_agent() { - let (svc, tm) = setup_with_factory(success_factory()); +async fn cold_bootstrap_failure_stops_session_and_cleans_all_successful_runtimes() { + use futures_util::FutureExt; + + let fail_conversation_id: Arc>> = Arc::new(Mutex::new(None)); + let fail_target = Arc::clone(&fail_conversation_id); + let factory: AgentFactory = Arc::new(move |opts: BuildTaskOptions| { + let fail_target = Arc::clone(&fail_target); + async move { + let conversation_id = opts.context.conversation.conversation_id.clone(); + if fail_target.lock().unwrap().as_deref() == Some(conversation_id.as_str()) { + return Err(AgentError::internal("simulated build failure")); + } + + Ok(aionui_ai_agent::AgentInstance::Mock(Arc::new( + mock_agent::MockAgent::new(conversation_id, opts.context.workspace.path), + ))) + } + .boxed() + }); + let (svc, _team_repo, tm, recorder) = setup_with_factory_and_recording_broadcaster(factory); let created = svc .create_team( "user1", CreateTeamRequest { name: "T".into(), - agents: two_agent_input(), + agents: four_agent_input_leader_not_first(), workspace: None, }, ) .await .unwrap(); + let failed_agent = created + .assistants + .iter() + .find(|assistant| assistant.name == "Worker 3") + .expect("failed agent") + .conversation_id + .clone(); + *fail_conversation_id.lock().unwrap() = Some(failed_agent); - reset_auto_started_session(&svc, &tm, &created.id).await; - svc.ensure_session("user1", &created.id).await.unwrap(); + let result = svc.ensure_session("user1", &created.id).await; + assert!(result.is_err(), "ensure_session should propagate build error"); + let error = result.unwrap_err().to_string(); + assert!( + error.contains("Worker 3") + && error.contains("backend=acp") + && error.contains("model=worker-3") + && error.contains("role=teammate"), + "rebuild error should identify the failing agent by name, backend, model, and role: {error}" + ); - // Two agents → kill called 2x and get_or_build_task called 2x, each with - // the corresponding conversation_id. Order is agents-iteration order. let calls = tm.snapshot(); - assert_eq!(calls.kill.len(), 2, "expected 2 kill calls"); - assert_eq!(calls.build.len(), 2, "expected 2 build calls"); - for (i, agent) in created.assistants.iter().enumerate() { - assert_eq!(calls.kill[i].0, agent.conversation_id); - assert_eq!(calls.kill[i].1, Some(AgentKillReason::TeamMcpRebuild)); - assert_eq!(calls.build[i], agent.conversation_id); + assert_eq!( + calls.build.len(), + 4, + "serial rebuild should stop only after the failing attempted agent" + ); + assert_eq!( + calls.kill.len(), + 11, + "cleanup is idempotent after partial-success cleanup" + ); + for agent in &created.assistants { + assert!( + calls + .kill + .iter() + .filter(|(conversation_id, _)| conversation_id == &agent.conversation_id) + .count() + >= 2, + "bootstrap failure must issue final cleanup for {}", + agent.conversation_id + ); } + assert_eq!(tm.active_count(), 0); + assert!( + svc.get_session_scheduler(&created.id).is_none(), + "session must not be registered after partial rebuild failure" + ); + + let team_session_failed = recorder + .events_by_name("team.sessionStatusChanged") + .into_iter() + .find(|event| { + event.data.get("team_id").and_then(|v| v.as_str()) == Some(created.id.as_str()) + && event.data.get("status").and_then(|v| v.as_str()) == Some("failed") + && event.data.get("phase").and_then(|v| v.as_str()) == Some("attaching_agents") + }); + assert!( + team_session_failed.is_some(), + "partial rebuild failure must emit a team-level failed/attaching_agents terminal event" + ); } #[tokio::test] -async fn d9_ensure_session_persists_team_mcp_stdio_config() { - // Each agent's conversation.extra must carry a `team_mcp_stdio_config` - // object by the time the factory is called — that is what the rebuilt - // typed Team context will expose to reach the MCP server. - use futures_util::FutureExt; - let (svc, _tm) = setup_with_factory(Arc::new(|opts: BuildTaskOptions| { - async move { - let context = opts.context; - let typed_has_cfg = context - .team - .as_ref() - .and_then(|team| team.mcp.as_ref()) - .is_some_and(|mcp| mcp.stdio.port > 0 && !mcp.stdio.slot_id.is_empty()); - let compat_has_cfg = match &context.kind { - AgentSessionKind::Acp(acp) => { - assert!(acp.belongs_to_team); - assert!(acp.team.is_some(), "ACP build context must carry typed team binding"); - acp.config - .team_mcp_stdio_config - .as_ref() - .is_some_and(|cfg| cfg.port > 0 && !cfg.slot_id.is_empty()) - } - _ => false, - }; - assert!( - typed_has_cfg && compat_has_cfg, - "factory called without typed team_mcp_stdio_config in context: {:?}", - context.team - ); - Ok(aionui_ai_agent::AgentInstance::Mock(Arc::new( - mock_agent::MockAgent::new(context.conversation.conversation_id, context.workspace.path), - ))) - } - .boxed() - })); - +async fn ensure_session_serializes_manual_add_until_rebuild_completes() { + let build_started = Arc::new(tokio::sync::Notify::new()); + let release_build = Arc::new(tokio::sync::Notify::new()); + let (svc, _tm) = setup_with_factory(blocking_first_build_factory( + Arc::clone(&build_started), + Arc::clone(&release_build), + )); let created = svc .create_team( "user1", CreateTeamRequest { name: "T".into(), - agents: two_agent_input(), + agents: vec![TeamAgentInput { + name: "Lead".into(), + role: "lead".into(), + backend: Some("acp".into()), + model: "claude".into(), + assistant_id: None, + conversation_id: None, + }], workspace: None, }, ) .await .unwrap(); - svc.ensure_session("user1", &created.id).await.unwrap(); + let svc_for_ensure = Arc::clone(&svc); + let ensure_team_id = created.id.clone(); + let ensure_handle = tokio::spawn(async move { svc_for_ensure.ensure_session("user1", &ensure_team_id).await }); + tokio::time::timeout(std::time::Duration::from_secs(2), build_started.notified()) + .await + .expect("ensure_session should start rebuilding before mutation attempt"); + + let svc_for_add = Arc::clone(&svc); + let add_team_id = created.id.clone(); + let mut add_handle = tokio::spawn(async move { + svc_for_add + .add_agent( + "user1", + &add_team_id, + AddAgentRequest { + name: "Worker".into(), + role: "teammate".into(), + backend: Some("acp".into()), + model: "claude".into(), + assistant_id: None, + }, + ) + .await + }); + + tokio::select! { + result = &mut add_handle => panic!("add_agent completed while ensure_session was rebuilding: {result:?}"), + _ = tokio::time::sleep(std::time::Duration::from_millis(200)) => {} + } + + release_build.notify_waiters(); + ensure_handle.await.unwrap().unwrap(); + add_handle.await.unwrap().unwrap(); + + let scheduler = svc + .get_session_scheduler(&created.id) + .expect("session should be registered after ensure_session"); + assert_eq!(scheduler.list_agents().await.len(), 2); } #[tokio::test] -async fn d9_ensure_session_is_idempotent() { - let (svc, tm) = setup_with_factory(success_factory()); +async fn ensure_session_serializes_manual_remove_until_rebuild_completes() { + let build_started = Arc::new(tokio::sync::Notify::new()); + let release_build = Arc::new(tokio::sync::Notify::new()); + let (svc, _tm) = setup_with_factory(blocking_first_build_factory( + Arc::clone(&build_started), + Arc::clone(&release_build), + )); let created = svc .create_team( "user1", @@ -4264,25 +6089,49 @@ async fn d9_ensure_session_is_idempotent() { ) .await .unwrap(); + let worker_slot = created.assistants[1].slot_id.clone(); - reset_auto_started_session(&svc, &tm, &created.id).await; - svc.ensure_session("user1", &created.id).await.unwrap(); - svc.ensure_session("user1", &created.id).await.unwrap(); + let svc_for_ensure = Arc::clone(&svc); + let ensure_team_id = created.id.clone(); + let ensure_handle = tokio::spawn(async move { svc_for_ensure.ensure_session("user1", &ensure_team_id).await }); + tokio::time::timeout(std::time::Duration::from_secs(2), build_started.notified()) + .await + .expect("ensure_session should start rebuilding before mutation attempt"); + + let svc_for_remove = Arc::clone(&svc); + let remove_team_id = created.id.clone(); + let remove_slot = worker_slot.clone(); + let mut remove_handle = tokio::spawn(async move { + svc_for_remove + .remove_agent("user1", &remove_team_id, &remove_slot) + .await + }); - // Second call short-circuits — no additional kill/build calls. - let calls = tm.snapshot(); - assert_eq!(calls.kill.len(), 2, "second ensure_session must not re-kill"); - assert_eq!(calls.build.len(), 2, "second ensure_session must not re-build"); + tokio::select! { + result = &mut remove_handle => panic!("remove_agent completed while ensure_session was rebuilding: {result:?}"), + _ = tokio::time::sleep(std::time::Duration::from_millis(200)) => {} + } + + release_build.notify_waiters(); + ensure_handle.await.unwrap().unwrap(); + remove_handle.await.unwrap().unwrap(); + + let scheduler = svc + .get_session_scheduler(&created.id) + .expect("session should be registered after ensure_session"); + let agents = scheduler.list_agents().await; + assert_eq!(agents.len(), 1); + assert!(agents.iter().all(|agent| agent.slot_id != worker_slot)); } #[tokio::test] -async fn d9_ensure_session_rollbacks_when_build_fails() { - // Factory always fails → ensure_session must propagate error and not - // insert into sessions, so send_message afterwards still errors. - use futures_util::FutureExt; - let failing_factory: AgentFactory = - Arc::new(|_opts: BuildTaskOptions| async move { Err(AgentError::internal("simulated build failure")) }.boxed()); - let (svc, tm) = setup_with_factory(failing_factory); +async fn ensure_session_serializes_manual_rename_until_rebuild_completes() { + let build_started = Arc::new(tokio::sync::Notify::new()); + let release_build = Arc::new(tokio::sync::Notify::new()); + let (svc, _tm) = setup_with_factory(blocking_first_build_factory( + Arc::clone(&build_started), + Arc::clone(&release_build), + )); let created = svc .create_team( "user1", @@ -4294,22 +6143,39 @@ async fn d9_ensure_session_rollbacks_when_build_fails() { ) .await .unwrap(); + let worker_slot = created.assistants[1].slot_id.clone(); - reset_auto_started_session(&svc, &tm, &created.id).await; - let result = svc.ensure_session("user1", &created.id).await; - assert!(result.is_err(), "ensure_session should propagate build error"); + let svc_for_ensure = Arc::clone(&svc); + let ensure_team_id = created.id.clone(); + let ensure_handle = tokio::spawn(async move { svc_for_ensure.ensure_session("user1", &ensure_team_id).await }); + tokio::time::timeout(std::time::Duration::from_secs(2), build_started.notified()) + .await + .expect("ensure_session should start rebuilding before mutation attempt"); + + let svc_for_rename = Arc::clone(&svc); + let rename_team_id = created.id.clone(); + let rename_slot = worker_slot.clone(); + let mut rename_handle = tokio::spawn(async move { + svc_for_rename + .rename_agent("user1", &rename_team_id, &rename_slot, "Senior Worker") + .await + }); - // Rebuild aborts on the first warmup failure, so only the first agent - // is killed/built. No session is inserted, so send_message still errors. - let calls = tm.snapshot(); - assert_eq!(calls.kill.len(), 1); - assert_eq!(calls.build.len(), 1); + tokio::select! { + result = &mut rename_handle => panic!("rename_agent completed while ensure_session was rebuilding: {result:?}"), + _ = tokio::time::sleep(std::time::Duration::from_millis(200)) => {} + } - let send_result = svc.send_message("user1", &created.id, "Hello", None).await; - assert!( - send_result.is_err(), - "session must not be registered after build failure" - ); + release_build.notify_waiters(); + ensure_handle.await.unwrap().unwrap(); + rename_handle.await.unwrap().unwrap(); + + let scheduler = svc + .get_session_scheduler(&created.id) + .expect("session should be registered after ensure_session"); + let agents = scheduler.list_agents().await; + let renamed = agents.iter().find(|agent| agent.slot_id == worker_slot).unwrap(); + assert_eq!(renamed.name, "Senior Worker"); } // =========================================================================== @@ -4413,7 +6279,7 @@ async fn d115_remove_team_kills_every_agent_process() { .await .unwrap(); - reset_auto_started_session(&svc, &tm, &created.id).await; + reset_runtime_state(&svc, &tm, &created.id).await; // Bring two agents online — after ensure_session, active_count == 2. svc.ensure_session("user1", &created.id).await.unwrap(); assert_eq!(tm.active_count(), 2, "ensure_session must register 2 live agents");