diff --git a/Cargo.lock b/Cargo.lock index c3b98b238..bf1e40f9d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -332,6 +332,7 @@ dependencies = [ "aionui-realtime", "aionui-runtime", "aionui-system", + "aionui-team-prompts", "async-trait", "axum", "base64", @@ -391,6 +392,7 @@ dependencies = [ "aionui-shell", "aionui-system", "aionui-team", + "aionui-team-prompts", "anyhow", "async-trait", "axum", @@ -818,6 +820,7 @@ dependencies = [ "aionui-common", "aionui-db", "aionui-realtime", + "aionui-team-prompts", "async-trait", "axum", "dashmap", @@ -832,6 +835,14 @@ dependencies = [ "tracing", ] +[[package]] +name = "aionui-team-prompts" +version = "0.1.31" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "allocator-api2" version = "0.2.21" diff --git a/Cargo.toml b/Cargo.toml index 390bd97c9..7b0c50d43 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,6 +17,7 @@ members = [ "crates/aionui-conversation", "crates/aionui-extension", "crates/aionui-channel", + "crates/aionui-team-prompts", "crates/aionui-team", "crates/aionui-cron", "crates/aionui-assistant", @@ -46,6 +47,7 @@ aionui-mcp = { path = "crates/aionui-mcp" } aionui-conversation = { path = "crates/aionui-conversation" } aionui-extension = { path = "crates/aionui-extension" } aionui-channel = { path = "crates/aionui-channel" } +aionui-team-prompts = { path = "crates/aionui-team-prompts" } aionui-team = { path = "crates/aionui-team" } aionui-cron = { path = "crates/aionui-cron" } aionui-assistant = { path = "crates/aionui-assistant" } diff --git a/crates/aionui-ai-agent/Cargo.toml b/crates/aionui-ai-agent/Cargo.toml index df0c8312f..dbf8d425e 100644 --- a/crates/aionui-ai-agent/Cargo.toml +++ b/crates/aionui-ai-agent/Cargo.toml @@ -13,6 +13,7 @@ aionui-realtime.workspace = true aionui-runtime.workspace = true aionui-extension.workspace = true aionui-api-types.workspace = true +aionui-team-prompts.workspace = true aionui-system.workspace = true axum.workspace = true base64.workspace = true diff --git a/crates/aionui-ai-agent/src/capability/mod.rs b/crates/aionui-ai-agent/src/capability/mod.rs index 7427fb065..3980b4eb8 100644 --- a/crates/aionui-ai-agent/src/capability/mod.rs +++ b/crates/aionui-ai-agent/src/capability/mod.rs @@ -10,6 +10,5 @@ pub(crate) mod cli_process; pub(crate) mod first_message_injector; pub mod prompt_pipeline; pub(crate) mod skill_manager; -pub(crate) mod team_guide_prompt; pub use prompt_pipeline::{PostRecvHook, PreSendHook, PromptCtx, PromptPipeline}; diff --git a/crates/aionui-ai-agent/src/capability/team_guide_prompt.rs b/crates/aionui-ai-agent/src/capability/team_guide_prompt.rs deleted file mode 100644 index 9bc70e579..000000000 --- a/crates/aionui-ai-agent/src/capability/team_guide_prompt.rs +++ /dev/null @@ -1,184 +0,0 @@ -//! Solo-agent Team Guide prompt (Layer 1) — teaches a solo ACP agent when and -//! how to propose a multi-agent Team to the user. -//! -//! This is a local copy of the prompt owned by `aionui_team::prompts::team_guide` -//! (see `crates/aionui-team/src/prompts/team_guide.rs`). We duplicate rather -//! than import because `aionui-team` depends on `aionui-ai-agent`, so the -//! reverse direction would create a dependency cycle. Future work may sink -//! this prompt into `aionui-common` and have both crates re-export; until -//! then, **keep this file byte-for-byte in sync with the team-side template** -//! — it ships to the LLM and must match AionUi's -//! `src/process/team/prompts/teamGuidePrompt.ts` exactly (aionui-audit §8 #5). -//! -//! Unlike the team-side helper, this module only exposes the solo branch -//! (`leader_label = None`). The preset-assistant label branch does not apply -//! to Wave 5 solo-agent injection. - -use aionui_common::constants::TEAM_CAPABLE_BACKENDS; - -const EXPLICIT_TEAM_REQUEST_CRITERIA: &str = "\ -- The user explicitly asks to create a Team -- The user explicitly asks for multiple agents, teammates, or parallel workers -- The user says they want to pull in a Team before starting"; - -const EXTREME_COMPLEXITY_CRITERIA: &str = "\ -- The task is so large, risky, or specialized that one agent is unlikely to complete it well alone -- The work needs substantial parallel role separation that cannot be reasonably handled in a normal solo workflow -- This bar is very high: if you can handle the task yourself, stay solo"; - -const STAY_SOLO_CRITERIA: &str = "\ -- Greetings, casual conversation, or general questions -- Single-point tasks: one question, one file, one fix, one translation, one explanation -- Normal coding, writing, research, or analysis tasks that one agent can handle with some effort -- Any task you can reasonably complete yourself, even if it takes multiple turns"; - -const SOLO_DEFAULT_RULE: &str = "Handle the task yourself in the current chat by default. Do NOT proactively recommend Team just because the work spans multiple files, takes multiple rounds, or would benefit from specialization."; - -const TEAM_GUIDE_PROMPT_TEMPLATE: &str = "## Team Mode - -You can create a multi-agent Team for the user. - -### Default behavior -{solo_default_rule} - -### Only bring up Team in either of these cases -1. The user explicitly wants a Team or multiple agents: -{explicit_team_request_criteria} -2. The task is exceptionally complex and you genuinely believe one agent is unlikely to handle it well alone: -{extreme_complexity_criteria} - -### Otherwise stay solo and do not mention Team -{stay_solo_criteria} - -If case 2 applies, ask at most once whether the user wants to bring in a Team. Keep it brief and optional. If the user says no, ignores it, or prefers solo help, continue solo and do not mention Team again. - -### How to proceed when Team is requested or approved (STRICT — follow every step, do NOT skip) -1. FIRST call `aion_list_models` to check available models for each agent type you plan to use. -2. Explain in one sentence why the Team setup helps this task. -3. Present a team configuration table: role name, responsibility, agent type, and recommended model (from aion_list_models results) for each member. Example format: - | Role | Responsibility | Type | Model | - | Leader | Coordinate and review | {leader_cell} | (default) | - | Developer | Implement features | {agent_type} | (model from list) | - | Tester | Write and run tests | {agent_type} | (model from list) | -4. **Output the table as a normal text message and END YOUR TURN.** Do NOT call `aion_create_team` or any other tool (including ask_user) in this turn. Wait for the user to reply in their next message with explicit confirmation (e.g. \"ok\", \"go ahead\", \"确认\") before proceeding. -5. After user confirms → call `aion_create_team`. The summary MUST include both the goal and the confirmed team configuration. (The system automatically sets the correct agent type — you do NOT need to pass agentType.) -6. After `aion_create_team` returns → you ARE now the team Leader. The system navigates to the team page automatically. **Immediately** use `team_spawn_agent` to create each teammate from the confirmed configuration table. Then use `team_send_message` to assign initial tasks to each spawned teammate. Do NOT end your turn until all teammates are spawned and tasked. -7. User declines or wants changes → adjust or proceed solo. Do not mention Team again unless the user asks. - -### Tool constraint -Before team creation: use **only** `aion_create_team` and `aion_list_models`. After `aion_create_team` succeeds: use team tools (`team_spawn_agent`, `team_send_message`, `team_members`, `team_task_create`, etc.) to manage your team."; - -/// Return `true` iff the given backend is a known team-capable backend. An -/// empty or unrecognized backend returns `false`; solo agents with unknown -/// backends do not receive the Team Guide prompt. -pub(crate) fn is_solo_team_guide_backend(backend: &str) -> bool { - TEAM_CAPABLE_BACKENDS.contains(&backend) -} - -/// Build the Team Guide prompt for a solo agent with the given backend label. -/// An empty `backend` falls back to `"claude"`, matching AionUi's -/// `opts.backend || 'claude'` and the team-side helper. -pub(crate) fn build_solo_team_guide_prompt(backend: &str) -> String { - let agent_type = if backend.is_empty() { "claude" } else { backend }; - TEAM_GUIDE_PROMPT_TEMPLATE - .replace("{solo_default_rule}", SOLO_DEFAULT_RULE) - .replace("{explicit_team_request_criteria}", EXPLICIT_TEAM_REQUEST_CRITERIA) - .replace("{extreme_complexity_criteria}", EXTREME_COMPLEXITY_CRITERIA) - .replace("{stay_solo_criteria}", STAY_SOLO_CRITERIA) - .replace("{leader_cell}", agent_type) - .replace("{agent_type}", agent_type) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn is_solo_team_guide_backend_allows_known_vendors() { - assert!(is_solo_team_guide_backend("claude")); - assert!(is_solo_team_guide_backend("codex")); - assert!(is_solo_team_guide_backend("gemini")); - assert!(is_solo_team_guide_backend("aionrs")); - } - - #[test] - fn is_solo_team_guide_backend_rejects_unknown_and_empty() { - assert!(!is_solo_team_guide_backend("")); - assert!(!is_solo_team_guide_backend("qwen")); - assert!(!is_solo_team_guide_backend("Claude")); // case-sensitive - } - - #[test] - fn build_solo_team_guide_prompt_resolves_all_placeholders() { - let prompt = build_solo_team_guide_prompt("claude"); - assert!(!prompt.contains("{solo_default_rule}")); - assert!(!prompt.contains("{explicit_team_request_criteria}")); - assert!(!prompt.contains("{extreme_complexity_criteria}")); - assert!(!prompt.contains("{stay_solo_criteria}")); - assert!(!prompt.contains("{leader_cell}")); - assert!(!prompt.contains("{agent_type}")); - } - - #[test] - fn build_solo_team_guide_prompt_renders_leader_row_with_backend() { - let prompt = build_solo_team_guide_prompt("gemini"); - assert!(prompt.contains("| Leader | Coordinate and review | gemini | (default) |")); - assert!(prompt.contains("| Developer | Implement features | gemini | (model from list) |")); - } - - #[test] - fn build_solo_team_guide_prompt_empty_backend_falls_back_to_claude() { - let prompt = build_solo_team_guide_prompt(""); - assert!(prompt.contains("| Leader | Coordinate and review | claude | (default) |")); - } - - #[test] - fn snapshot_matches_team_crate_verbatim() { - // Byte-for-byte equality with the canonical prompt in - // `aionui-team/src/prompts/team_guide.rs`. If this test fails, one - // side has drifted — update **both** files (and the AionUi TS - // source in lockstep) rather than patching this assertion. - let prompt = build_solo_team_guide_prompt("claude"); - let expected = "## Team Mode\n\ -\n\ -You can create a multi-agent Team for the user.\n\ -\n\ -### Default behavior\n\ -Handle the task yourself in the current chat by default. Do NOT proactively recommend Team just because the work spans multiple files, takes multiple rounds, or would benefit from specialization.\n\ -\n\ -### Only bring up Team in either of these cases\n\ -1. The user explicitly wants a Team or multiple agents:\n\ -- The user explicitly asks to create a Team\n\ -- The user explicitly asks for multiple agents, teammates, or parallel workers\n\ -- The user says they want to pull in a Team before starting\n\ -2. The task is exceptionally complex and you genuinely believe one agent is unlikely to handle it well alone:\n\ -- The task is so large, risky, or specialized that one agent is unlikely to complete it well alone\n\ -- The work needs substantial parallel role separation that cannot be reasonably handled in a normal solo workflow\n\ -- This bar is very high: if you can handle the task yourself, stay solo\n\ -\n\ -### Otherwise stay solo and do not mention Team\n\ -- Greetings, casual conversation, or general questions\n\ -- Single-point tasks: one question, one file, one fix, one translation, one explanation\n\ -- Normal coding, writing, research, or analysis tasks that one agent can handle with some effort\n\ -- Any task you can reasonably complete yourself, even if it takes multiple turns\n\ -\n\ -If case 2 applies, ask at most once whether the user wants to bring in a Team. Keep it brief and optional. If the user says no, ignores it, or prefers solo help, continue solo and do not mention Team again.\n\ -\n\ -### How to proceed when Team is requested or approved (STRICT — follow every step, do NOT skip)\n\ -1. FIRST call `aion_list_models` to check available models for each agent type you plan to use.\n\ -2. Explain in one sentence why the Team setup helps this task.\n\ -3. Present a team configuration table: role name, responsibility, agent type, and recommended model (from aion_list_models results) for each member. Example format:\n \ -| Role | Responsibility | Type | Model |\n \ -| Leader | Coordinate and review | claude | (default) |\n \ -| Developer | Implement features | claude | (model from list) |\n \ -| Tester | Write and run tests | claude | (model from list) |\n\ -4. **Output the table as a normal text message and END YOUR TURN.** Do NOT call `aion_create_team` or any other tool (including ask_user) in this turn. Wait for the user to reply in their next message with explicit confirmation (e.g. \"ok\", \"go ahead\", \"确认\") before proceeding.\n\ -5. After user confirms → call `aion_create_team`. The summary MUST include both the goal and the confirmed team configuration. (The system automatically sets the correct agent type — you do NOT need to pass agentType.)\n\ -6. After `aion_create_team` returns → you ARE now the team Leader. The system navigates to the team page automatically. **Immediately** use `team_spawn_agent` to create each teammate from the confirmed configuration table. Then use `team_send_message` to assign initial tasks to each spawned teammate. Do NOT end your turn until all teammates are spawned and tasked.\n\ -7. User declines or wants changes → adjust or proceed solo. Do not mention Team again unless the user asks.\n\ -\n\ -### Tool constraint\n\ -Before team creation: use **only** `aion_create_team` and `aion_list_models`. After `aion_create_team` succeeds: use team tools (`team_spawn_agent`, `team_send_message`, `team_members`, `team_task_create`, etc.) to manage your team."; - assert_eq!(prompt, expected); - } -} diff --git a/crates/aionui-ai-agent/src/factory/acp_assembler.rs b/crates/aionui-ai-agent/src/factory/acp_assembler.rs index 3f5fb44e3..d745387a8 100644 --- a/crates/aionui-ai-agent/src/factory/acp_assembler.rs +++ b/crates/aionui-ai-agent/src/factory/acp_assembler.rs @@ -1,9 +1,9 @@ -use crate::capability::team_guide_prompt; use crate::shared_kernel::PersistedSessionState; use agent_client_protocol::schema::{EnvVariable, McpServer, McpServerStdio, NewSessionRequest}; use aionui_api_types::AgentMetadata; use aionui_api_types::{AcpBuildExtra, GuideMcpConfig, TEAM_MCP_SERVER_NAME, TeamMcpStdioConfig}; use aionui_common::CommandSpec; +use aionui_team_prompts::guide as team_guide_prompt; use std::path::PathBuf; use aionui_common::constants::TEAM_CAPABLE_BACKENDS; @@ -194,7 +194,14 @@ mod tests { fn compose_preset_context_team_capable_backend_appends_guide() { let result = compose_preset_context(None, Some("claude"), false); assert!(result.is_some()); - assert!(result.unwrap().contains("team")); + let prompt = result.unwrap(); + assert!(prompt.contains("aion_create_team")); + assert!(prompt.contains("aion_list_models")); + assert!(prompt.contains("hand off to the created Team conversation")); + assert!(!prompt.contains("Immediately")); + assert!(!prompt.contains( + "use team tools (`team_spawn_agent`, `team_send_message`, `team_members`, `team_task_create`, etc.) to manage your team" + )); } #[test] @@ -207,6 +214,78 @@ mod tests { McpServer::Stdio(McpServerStdio::new(name, "/bin/sh")) } + fn team_cfg() -> TeamMcpStdioConfig { + TeamMcpStdioConfig { + team_id: "team-1".into(), + port: 9999, + token: "tok".into(), + slot_id: "slot-lead".into(), + binary_path: "/bin/backend".into(), + } + } + + fn test_metadata() -> AgentMetadata { + AgentMetadata { + id: "agent-1".into(), + icon: None, + name: "Test ACP".into(), + name_i18n: None, + description: None, + description_i18n: None, + backend: Some("claude".into()), + agent_type: aionui_common::AgentType::Acp, + agent_source: aionui_api_types::AgentSource::Builtin, + agent_source_info: aionui_api_types::AgentSourceInfo::default(), + enabled: true, + available: true, + command: Some("claude".into()), + resolved_command: None, + args: vec![], + env: vec![], + native_skills_dirs: None, + behavior_policy: aionui_api_types::BehaviorPolicy::default(), + yolo_id: None, + sort_order: 0, + team_capable: true, + handshake: aionui_api_types::AgentHandshake::default(), + } + } + + #[tokio::test] + async fn assemble_acp_params_uses_frozen_preset_context_and_snapshot_seeds() { + let config = AcpBuildExtra { + backend: Some("claude".into()), + preset_context: Some("frozen rules".into()), + skills: vec!["pdf".into()], + mcp_server_ids: Some(vec!["mcp-docs".into()]), + team_mcp_stdio_config: Some(team_cfg()), + ..Default::default() + }; + + let params = assemble_acp_params( + "conv-1".into(), + WorkspaceInfo { + path: "/tmp/workspace".into(), + is_custom: false, + }, + test_metadata(), + CommandSpec::default(), + config, + vec![user_stdio("mcp-docs")], + None, + PathBuf::from("/tmp/data"), + ) + .await; + + assert_eq!(params.preset_context.as_deref(), Some("frozen rules")); + assert_eq!(params.config.skills, vec!["pdf"]); + assert_eq!( + params.config.mcp_server_ids.as_deref(), + Some(&["mcp-docs".to_owned()][..]) + ); + assert_eq!(params.mcp_servers.len(), 2); + } + #[test] fn resolve_mcp_servers_prefers_team_over_guide() { let config = AcpBuildExtra { diff --git a/crates/aionui-ai-agent/src/factory/aionrs.rs b/crates/aionui-ai-agent/src/factory/aionrs.rs index b72f94184..10624325b 100644 --- a/crates/aionui-ai-agent/src/factory/aionrs.rs +++ b/crates/aionui-ai-agent/src/factory/aionrs.rs @@ -14,7 +14,6 @@ use aionui_runtime::ensure_runtime_command_with_reporter; use tracing::{debug, info, warn}; use crate::agent_task::AgentInstance; -use crate::capability::team_guide_prompt; use crate::error::AgentError; use crate::factory::AgentFactoryDeps; use crate::factory::context::FactoryContext; @@ -22,6 +21,7 @@ use crate::manager::aionrs::{AionrsAgentManager, sanitize_session_messages}; use crate::runtime_status::conversation_runtime_reporter; use crate::session_context::AionrsSessionBuildContext; use crate::types::{AionrsCompatOverrides, AionrsResolvedConfig}; +use aionui_team_prompts::guide as team_guide_prompt; const TEAM_CAPABLE_BACKENDS: &[&str] = &["claude", "codex", "gemini", "aionrs", "codebuddy"]; @@ -37,6 +37,10 @@ pub(super) async fn build( // Merge preset assistant rules into system_prompt (used as custom_prompt // in aionrs's build_system_prompt). Mirrors the old architecture's // `init_history` injection of `[Assistant System Rules]`. + // AionrsBuildExtra parses `skills` so Team preset snapshots preserve the + // target contract. Native skill materialization for Aionrs is tracked as a + // separate follow-up because this factory currently has no stable Aionrs + // skill-loading path. if let Some(rules) = overrides.preset_rules.take() { overrides.system_prompt = Some(match overrides.system_prompt.take() { Some(existing) => format!("{existing}\n\n{rules}"), @@ -695,10 +699,88 @@ mod tests { } } + struct MockMcpRepo { + rows: Vec, + } + + #[async_trait::async_trait] + impl IMcpServerRepository for MockMcpRepo { + async fn list(&self) -> Result, aionui_db::DbError> { + Ok(self.rows.clone()) + } + + async fn find_by_id(&self, id: &str) -> Result, aionui_db::DbError> { + Ok(self.rows.iter().find(|row| row.id == id).cloned()) + } + + async fn find_by_name(&self, name: &str) -> Result, aionui_db::DbError> { + Ok(self.rows.iter().find(|row| row.name == name).cloned()) + } + + async fn create( + &self, + _params: aionui_db::CreateMcpServerParams<'_>, + ) -> Result { + unimplemented!("not needed for factory tests") + } + + async fn update( + &self, + _id: &str, + _params: aionui_db::UpdateMcpServerParams<'_>, + ) -> Result { + unimplemented!("not needed for factory tests") + } + + async fn delete(&self, _id: &str) -> Result<(), aionui_db::DbError> { + unimplemented!("not needed for factory tests") + } + + async fn batch_upsert( + &self, + _servers: &[aionui_db::CreateMcpServerParams<'_>], + ) -> Result, aionui_db::DbError> { + unimplemented!("not needed for factory tests") + } + + async fn update_status( + &self, + _id: &str, + _status: &str, + _last_connected: Option, + ) -> Result<(), aionui_db::DbError> { + unimplemented!("not needed for factory tests") + } + + async fn update_tools(&self, _id: &str, _tools: Option<&str>) -> Result<(), aionui_db::DbError> { + unimplemented!("not needed for factory tests") + } + } + fn test_broadcaster() -> Arc { Arc::new(BroadcastEventBus::new(16)) } + #[tokio::test] + async fn aionrs_loads_mcp_servers_from_frozen_selection_snapshot() { + let mut row = make_row( + "mcp-docs", + "http", + r#"{"url":"http://localhost:54321/mcp","headers":{"Authorization":"Bearer frozen"}}"#, + false, + false, + ); + row.id = "mcp-docs".into(); + let repo = MockMcpRepo { rows: vec![row] }; + let selected = vec!["mcp-docs".to_owned()]; + + let extra_mcp_servers = + load_user_mcp_servers(&repo, Some(&selected), "conv-frozen-mcp", test_broadcaster()).await; + + assert!(extra_mcp_servers.contains_key("mcp-docs")); + assert_eq!(extra_mcp_servers["mcp-docs"].transport, TransportType::StreamableHttp); + } + #[cfg(unix)] #[tokio::test] async fn row_to_mcp_server_config_flattens_resolved_npx_command() { @@ -946,6 +1028,21 @@ mod tests { assert_eq!(env.get("AION_MCP_USER_ID"), Some(&"user-1".to_owned())); } + #[test] + fn aionrs_guide_prompt_hands_off_after_create_team() { + let mut overrides = AionrsBuildExtra::default(); + overrides.system_prompt = Some(team_guide_prompt::build_solo_team_guide_prompt("aionrs")); + + let prompt = overrides.system_prompt.as_deref().unwrap(); + assert!(prompt.contains("aion_create_team")); + assert!(prompt.contains("aion_list_models")); + assert!(prompt.contains("hand off to the created Team conversation")); + assert!(!prompt.contains("Immediately")); + assert!(!prompt.contains( + "use team tools (`team_spawn_agent`, `team_send_message`, `team_members`, `team_task_create`, etc.) to manage your team" + )); + } + #[test] fn resolve_mcp_servers_empty_when_no_config() { let overrides = AionrsBuildExtra::default(); diff --git a/crates/aionui-ai-agent/src/types.rs b/crates/aionui-ai-agent/src/types.rs index e008a7cd5..1a29eb4d0 100644 --- a/crates/aionui-ai-agent/src/types.rs +++ b/crates/aionui-ai-agent/src/types.rs @@ -218,4 +218,14 @@ mod tests { assert!(extra.system_prompt.is_none()); assert_eq!(extra.preset_rules.unwrap(), "You are a data analyst."); } + + #[test] + fn aionrs_build_extra_accepts_frozen_skills_snapshot() { + let json = json!({ + "preset_rules": "Rules", + "skills": ["pdf", "cron"] + }); + let extra: AionrsBuildExtra = serde_json::from_value(json).unwrap(); + assert_eq!(extra.skills, vec!["pdf".to_owned(), "cron".to_owned()]); + } } diff --git a/crates/aionui-api-types/src/agent_build_extra.rs b/crates/aionui-api-types/src/agent_build_extra.rs index 9025f6c10..75a8336dd 100644 --- a/crates/aionui-api-types/src/agent_build_extra.rs +++ b/crates/aionui-api-types/src/agent_build_extra.rs @@ -84,6 +84,8 @@ pub struct AionrsBuildExtra { pub system_prompt: Option, #[serde(default)] pub preset_rules: Option, + #[serde(default)] + pub skills: Vec, #[serde(default = "default_aionrs_max_tokens")] pub max_tokens: u32, #[serde(default)] diff --git a/crates/aionui-app/Cargo.toml b/crates/aionui-app/Cargo.toml index 2576b9bf8..1781fef6c 100644 --- a/crates/aionui-app/Cargo.toml +++ b/crates/aionui-app/Cargo.toml @@ -31,6 +31,7 @@ aionui-conversation.workspace = true aionui-extension.workspace = true aionui-channel.workspace = true aionui-team.workspace = true +aionui-team-prompts.workspace = true aionui-cron.workspace = true aionui-assistant.workspace = true aionui-runtime.workspace = true diff --git a/crates/aionui-app/src/commands/cmd_team_guide.rs b/crates/aionui-app/src/commands/cmd_team_guide.rs index 1e9dbcfcf..e2c480df9 100644 --- a/crates/aionui-app/src/commands/cmd_team_guide.rs +++ b/crates/aionui-app/src/commands/cmd_team_guide.rs @@ -23,6 +23,9 @@ const ENV_TOKEN: &str = "AION_MCP_TOKEN"; const ENV_BACKEND: &str = "AION_MCP_BACKEND"; const ENV_CONVERSATION_ID: &str = "AION_MCP_CONVERSATION_ID"; const ENV_USER_ID: &str = "AION_MCP_USER_ID"; +#[cfg(test)] +const GUIDE_TEAM_TOOL_REMOVED_ERROR: &str = + "team_* tools are not available in Guide MCP; switch to the Team conversation."; pub async fn run_team_guide() -> ExitCode { let env = match GuideEnv::from_env() { @@ -141,95 +144,6 @@ struct ListModelsParams { agent_type: Option, } -#[derive(Deserialize, schemars::JsonSchema)] -struct SendMessageParams { - /// Target teammate name, or "*" to broadcast to all. - to: String, - /// Message content to send. - message: String, -} - -#[derive(Deserialize, schemars::JsonSchema)] -struct SpawnAgentParams { - /// Name for the new teammate agent. - name: String, - /// AI backend type: "claude" or "codex". Default when omitted. - #[serde(default)] - agent_type: Option, - /// Preset assistant identifier. - #[serde(default)] - custom_agent_id: Option, - /// Model override for the new agent. - #[serde(default)] - model: Option, -} - -#[derive(Deserialize, schemars::JsonSchema)] -struct TaskCreateParams { - /// Short task title. - subject: String, - /// Detailed task description. - #[serde(default)] - description: Option, - /// Teammate name assigned as owner. - #[serde(default)] - owner: Option, - /// Task IDs that must complete before this task can start. - #[serde(default)] - blocked_by: Option>, -} - -#[derive(Deserialize, schemars::JsonSchema)] -struct TaskUpdateParams { - /// ID of the task to update. - task_id: String, - /// New status: pending, in_progress, completed, or deleted. - #[serde(default)] - status: Option, - /// Updated task description. - #[serde(default)] - description: Option, - /// New owner teammate name. - #[serde(default)] - owner: Option, - /// Updated list of blocking task IDs. - #[serde(default)] - blocked_by: Option>, -} - -#[derive(Deserialize, schemars::JsonSchema)] -struct RenameAgentParams { - /// Slot ID of the team member to rename. - slot_id: String, - /// New display name. - new_name: String, -} - -#[derive(Deserialize, schemars::JsonSchema)] -struct ShutdownAgentParams { - /// Slot ID of the teammate to shut down. - slot_id: String, - /// Optional reason for shutdown. - #[serde(default)] - reason: Option, -} - -#[derive(Deserialize, schemars::JsonSchema)] -struct TeamListModelsParams { - /// Agent type to filter models (e.g. "claude", "codex"). Shows all when omitted. - #[serde(default)] - agent_type: Option, -} - -#[derive(Deserialize, schemars::JsonSchema)] -struct DescribeAssistantParams { - /// Preset assistant identifier to look up. - custom_agent_id: String, - /// Locale for the description (e.g. "en", "zh"). Default when omitted. - #[serde(default)] - locale: Option, -} - #[tool_router(server_handler)] impl GuideServer { #[tool( @@ -261,139 +175,6 @@ impl GuideServer { ) .await } - - #[tool( - name = "team_send_message", - description = "Send a message to a teammate or broadcast to all (to=\"*\")." - )] - async fn team_send_message(&self, Parameters(params): Parameters) -> CallToolResult { - self.forward_tool( - "team_send_message", - &serde_json::json!({ - "to": params.to, - "message": params.message, - }), - ) - .await - } - - #[tool( - name = "team_spawn_agent", - description = "Create a new teammate agent to join the team. Leader only." - )] - async fn team_spawn_agent(&self, Parameters(params): Parameters) -> CallToolResult { - self.forward_tool( - "team_spawn_agent", - &serde_json::json!({ - "name": params.name, - "agent_type": params.agent_type, - "custom_agent_id": params.custom_agent_id, - "model": params.model, - }), - ) - .await - } - - #[tool(name = "team_task_create", description = "Create a new task on the team task board.")] - async fn team_task_create(&self, Parameters(params): Parameters) -> CallToolResult { - self.forward_tool( - "team_task_create", - &serde_json::json!({ - "subject": params.subject, - "description": params.description, - "owner": params.owner, - "blocked_by": params.blocked_by, - }), - ) - .await - } - - #[tool( - name = "team_task_update", - description = "Update an existing task on the team task board." - )] - async fn team_task_update(&self, Parameters(params): Parameters) -> CallToolResult { - self.forward_tool( - "team_task_update", - &serde_json::json!({ - "task_id": params.task_id, - "status": params.status, - "description": params.description, - "owner": params.owner, - "blocked_by": params.blocked_by, - }), - ) - .await - } - - #[tool(name = "team_task_list", description = "List all tasks on the team task board.")] - async fn team_task_list(&self) -> CallToolResult { - self.forward_tool("team_task_list", &serde_json::json!({})).await - } - - #[tool( - name = "team_members", - description = "List all team members with their roles and current status." - )] - async fn team_members(&self) -> CallToolResult { - self.forward_tool("team_members", &serde_json::json!({})).await - } - - #[tool(name = "team_rename_agent", description = "Rename a team member.")] - async fn team_rename_agent(&self, Parameters(params): Parameters) -> CallToolResult { - self.forward_tool( - "team_rename_agent", - &serde_json::json!({ - "slot_id": params.slot_id, - "new_name": params.new_name, - }), - ) - .await - } - - #[tool( - name = "team_shutdown_agent", - description = "Initiate graceful shutdown of a teammate. Leader only." - )] - async fn team_shutdown_agent(&self, Parameters(params): Parameters) -> CallToolResult { - self.forward_tool( - "team_shutdown_agent", - &serde_json::json!({ - "slot_id": params.slot_id, - "reason": params.reason, - }), - ) - .await - } - - #[tool( - name = "team_list_models", - description = "Query available models for team agent types." - )] - async fn team_list_models(&self, Parameters(params): Parameters) -> CallToolResult { - self.forward_tool( - "team_list_models", - &serde_json::json!({ - "agent_type": params.agent_type, - }), - ) - .await - } - - #[tool( - name = "team_describe_assistant", - description = "Get detailed information about a preset assistant before spawning." - )] - async fn team_describe_assistant(&self, Parameters(params): Parameters) -> CallToolResult { - self.forward_tool( - "team_describe_assistant", - &serde_json::json!({ - "custom_agent_id": params.custom_agent_id, - "locale": params.locale, - }), - ) - .await - } } #[cfg(test)] @@ -423,6 +204,29 @@ mod tests { } } + #[test] + fn guide_stdio_exposes_only_create_team_and_list_models() { + let router = GuideServer::tool_router(); + let mut names: Vec = router + .list_all() + .into_iter() + .map(|tool| tool.name.to_string()) + .collect(); + names.sort(); + assert_eq!( + names, + vec!["aion_create_team".to_owned(), "aion_list_models".to_owned()] + ); + } + + #[test] + fn removed_guide_team_tool_error_text_is_stable() { + assert_eq!( + GUIDE_TEAM_TOOL_REMOVED_ERROR, + "team_* tools are not available in Guide MCP; switch to the Team conversation." + ); + } + #[test] fn guide_env_rejects_invalid_port_with_stable_code() { let err = GuideEnv::from_values("bad", "tok", "", "", "").unwrap_err(); diff --git a/crates/aionui-app/src/commands/cmd_team_stdio.rs b/crates/aionui-app/src/commands/cmd_team_stdio.rs index 05c66ff29..5c4aed215 100644 --- a/crates/aionui-app/src/commands/cmd_team_stdio.rs +++ b/crates/aionui-app/src/commands/cmd_team_stdio.rs @@ -15,7 +15,7 @@ use crate::commands::error::{CliBoundaryCode, CliBoundaryError, missing_env, par use aionui_api_types::TeamMcpStdioConfig; use aionui_team::mcp::protocol::{read_frame, write_frame}; use rmcp::handler::server::wrapper::Parameters; -use rmcp::model::{CallToolResult, Content}; +use rmcp::model::{CallToolResult, Content, ListToolsResult, Tool}; use rmcp::{schemars, service::ServiceExt, tool, tool_router, transport}; use serde::Deserialize; use tokio::net::TcpStream; @@ -216,7 +216,7 @@ struct DescribeAssistantParams { // Tool router // --------------------------------------------------------------------------- -#[tool_router(server_handler)] +#[tool_router] impl TeamStdioServer { #[tool( name = "team_send_message", @@ -294,7 +294,7 @@ impl TeamStdioServer { self.forward_to_tcp("team_members", &serde_json::json!({})).await } - #[tool(name = "team_rename_agent", description = "Rename a team member.")] + #[tool(name = "team_rename_agent", description = "Rename a team member. Lead only.")] async fn rename_agent(&self, Parameters(params): Parameters) -> CallToolResult { self.forward_to_tcp( "team_rename_agent", @@ -305,7 +305,7 @@ impl TeamStdioServer { #[tool( name = "team_shutdown_agent", - description = "Initiate shutdown of a teammate (Lead only). Sends a shutdown_request to the target agent." + description = "Initiate shutdown of a teammate. Lead only. Sends a shutdown_request to the target agent." )] async fn shutdown_agent(&self, Parameters(params): Parameters) -> CallToolResult { self.forward_to_tcp( @@ -340,6 +340,21 @@ impl TeamStdioServer { } } +#[rmcp::tool_handler(router = Self::tool_router())] +impl rmcp::ServerHandler for TeamStdioServer { + async fn list_tools( + &self, + _request: Option, + _context: rmcp::service::RequestContext, + ) -> Result { + let tools = self + .list_tools_from_tcp() + .await + .map_err(|_| rmcp::ErrorData::internal_error("failed to list local team tools", None))?; + Ok(ListToolsResult::with_all_items(tools)) + } +} + // --------------------------------------------------------------------------- // TCP forwarding // --------------------------------------------------------------------------- @@ -405,6 +420,43 @@ impl TeamStdioServer { parse_tool_response(&text) } + + async fn list_tools_from_tcp(&self) -> Result, ToolForwardError> { + let mut stream = TcpStream::connect((CONNECT_HOST, self.port)) + .await + .map_err(|_| tcp_connect_error(self.port))?; + stream.set_nodelay(true).ok(); + + let init_frame = serde_json::json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "auth_token": self.token, + "slot_id": self.slot_id, + } + }); + let init_bytes = serde_json::to_vec(&init_frame).map_err(|_| json_serialize_error())?; + write_frame(&mut stream, &init_bytes) + .await + .map_err(|_| tcp_write_error())?; + let init_resp = read_frame(&mut stream).await.map_err(|_| tcp_read_error())?; + let init_text = String::from_utf8_lossy(&init_resp).into_owned(); + parse_json_rpc_success(&init_text)?; + + let list_frame = serde_json::json!({ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/list", + }); + let list_bytes = serde_json::to_vec(&list_frame).map_err(|_| json_serialize_error())?; + write_frame(&mut stream, &list_bytes) + .await + .map_err(|_| tcp_write_error())?; + let resp_bytes = read_frame(&mut stream).await.map_err(|_| tcp_read_error())?; + let text = String::from_utf8_lossy(&resp_bytes).into_owned(); + parse_tools_list_response(&text) + } } #[derive(Debug)] @@ -466,6 +518,52 @@ fn parse_tool_response(text: &str) -> Result { Err(tool_response_unexpected().into()) } +fn parse_json_rpc_success(text: &str) -> Result { + let value = serde_json::from_str::(text).map_err(|_| tool_response_unexpected())?; + if value.get("error").is_some() { + return Err(remote_tool_error( + extract_nested_code(&value, &["error", "code"]), + extract_nested_code(&value, &["error", "data", "domainCode"]) + .or_else(|| extract_nested_code(&value, &["error", "data", "code"])) + .or_else(|| extract_nested_code(&value, &["error", "data", "errorCode"])), + )); + } + value + .get("result") + .cloned() + .ok_or_else(tool_response_unexpected) + .map_err(Into::into) +} + +#[derive(Deserialize)] +struct RemoteToolDescriptor { + name: String, + #[serde(default)] + description: String, + #[serde(default, alias = "inputSchema")] + input_schema: serde_json::Value, +} + +fn parse_tools_list_response(text: &str) -> Result, ToolForwardError> { + let result = parse_json_rpc_success(text)?; + let descriptors = serde_json::from_value::>( + result.get("tools").cloned().ok_or_else(tool_response_unexpected)?, + ) + .map_err(|_| tool_response_unexpected())?; + + descriptors + .into_iter() + .map(|descriptor| { + let schema = descriptor + .input_schema + .as_object() + .cloned() + .ok_or_else(tool_response_unexpected)?; + Ok(Tool::new(descriptor.name, descriptor.description, schema)) + }) + .collect() +} + fn json_serialize_error() -> CliBoundaryError { CliBoundaryError::new(CliBoundaryCode::McpJsonSerializeFailed, SUBCOMMAND, ERR_JSON_SERIALIZE) } @@ -579,6 +677,100 @@ mod tests { assert_eq!(env.slot_id, "slot-a"); } + #[test] + fn team_stdio_descriptions_match_prompt_registry() { + let router = TeamStdioServer::tool_router(); + let tools = router.list_all(); + + for spec in aionui_team_prompts::tools::team_tool_specs() { + let tool = tools + .iter() + .find(|tool| tool.name == spec.name) + .unwrap_or_else(|| panic!("missing tool {}", spec.name)); + let description = tool + .description + .as_ref() + .unwrap_or_else(|| panic!("missing description for {}", spec.name)); + assert_eq!( + description.as_ref(), + spec.description, + "description drift for {}", + spec.name + ); + } + } + + #[tokio::test] + async fn list_tools_uses_team_server_filtered_descriptors() { + 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_value: serde_json::Value = serde_json::from_slice(&init).unwrap(); + assert_eq!(init_value["method"], "initialize"); + assert_eq!(init_value["params"]["slot_id"], "worker-1"); + + let init_response = serde_json::to_vec(&json!({ + "jsonrpc": "2.0", + "id": 1, + "result": {} + })) + .unwrap(); + write_frame(&mut socket, &init_response).await.unwrap(); + + let list = read_frame(&mut socket).await.unwrap(); + let list_value: serde_json::Value = serde_json::from_slice(&list).unwrap(); + assert_eq!(list_value["method"], "tools/list"); + + let list_response = serde_json::to_vec(&json!({ + "jsonrpc": "2.0", + "id": 2, + "result": { + "tools": [ + { + "name": "team_send_message", + "description": "Send a message", + "input_schema": { + "type": "object", + "properties": { + "to": { "type": "string" }, + "message": { "type": "string" } + }, + "required": ["to", "message"] + } + } + ] + } + })) + .unwrap(); + write_frame(&mut socket, &list_response).await.unwrap(); + }); + let server = TeamStdioServer { + port, + token: "dummy-token".into(), + slot_id: "worker-1".into(), + }; + + let tools = server.list_tools_from_tcp().await.expect("tools/list"); + + accept_task.await.unwrap(); + let names: Vec<_> = tools.iter().map(|tool| tool.name.as_ref()).collect(); + assert_eq!(names, vec!["team_send_message"]); + assert!(!names.contains(&"team_spawn_agent")); + assert!(!names.contains(&"team_rename_agent")); + assert!(!names.contains(&"team_shutdown_agent")); + assert_eq!( + tools[0] + .input_schema + .get("properties") + .and_then(|value| value.as_object()) + .unwrap() + .len(), + 2 + ); + } + #[tokio::test] async fn forward_to_tcp_reports_read_failure_after_accept_close() { let listener = TcpListener::bind((CONNECT_HOST, 0)).await.unwrap(); diff --git a/crates/aionui-team-prompts/Cargo.toml b/crates/aionui-team-prompts/Cargo.toml new file mode 100644 index 000000000..5ab16a2c7 --- /dev/null +++ b/crates/aionui-team-prompts/Cargo.toml @@ -0,0 +1,8 @@ +[package] +name = "aionui-team-prompts" +version.workspace = true +edition.workspace = true + +[dependencies] +serde.workspace = true +serde_json.workspace = true diff --git a/crates/aionui-team-prompts/src/governance.rs b/crates/aionui-team-prompts/src/governance.rs new file mode 100644 index 000000000..5ea7f41d1 --- /dev/null +++ b/crates/aionui-team-prompts/src/governance.rs @@ -0,0 +1,44 @@ +pub const TEAM_GOVERNANCE_PROMPT: &str = r#"## Team Governance + +In Team mode, assistant rules define the agent's domain behavior, but Team Governance defines collaboration authority. + +Priority order: +1. Platform and system rules +2. Team Governance +3. Team role prompt +4. Assistant rules +5. Wake payload and current task context +6. Ordinary history context + +When assistant rules conflict with Team collaboration, role, permission, task-board, or reporting behavior, Team Governance and the Team role prompt win. + +Required Team behavior: +- Use `team_*` MCP tools for Team coordination. +- Use `team_send_message` for Team reporting instead of ordinary assistant replies. +- Use `team_task_update` and `team_task_list` for task-board state. +- Follow role permissions. Lead-only tools cannot be used by teammates. +- Domain-specific assistant rules, MCP servers, and skills remain active only inside these Team boundaries."#; + +pub fn with_team_governance(role_prompt: &str) -> String { + format!("{TEAM_GOVERNANCE_PROMPT}\n\n{role_prompt}") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn governance_declares_team_priority_over_assistant_rules() { + assert!(TEAM_GOVERNANCE_PROMPT.contains("assistant rules")); + assert!(TEAM_GOVERNANCE_PROMPT.contains("Team Governance and the Team role prompt win")); + assert!(TEAM_GOVERNANCE_PROMPT.contains("Lead-only tools")); + assert!(TEAM_GOVERNANCE_PROMPT.contains("team_send_message")); + } + + #[test] + fn wrapper_prepends_governance_once() { + let out = with_team_governance("## Role\nDo work."); + assert!(out.starts_with("## Team Governance")); + assert!(out.contains("## Role\nDo work.")); + } +} diff --git a/crates/aionui-team-prompts/src/guide.rs b/crates/aionui-team-prompts/src/guide.rs new file mode 100644 index 000000000..b19220a57 --- /dev/null +++ b/crates/aionui-team-prompts/src/guide.rs @@ -0,0 +1,108 @@ +const EXPLICIT_TEAM_REQUEST_CRITERIA: &str = "\ +- The user explicitly asks to create a Team +- The user explicitly asks for multiple agents, teammates, or parallel workers +- The user says they want to pull in a Team before starting"; + +const EXTREME_COMPLEXITY_CRITERIA: &str = "\ +- The task is so large, risky, or specialized that one agent is unlikely to complete it well alone +- The work needs substantial parallel role separation that cannot be reasonably handled in a normal solo workflow +- This bar is very high: if you can handle the task yourself, stay solo"; + +const STAY_SOLO_CRITERIA: &str = "\ +- Greetings, casual conversation, or general questions +- Single-point tasks: one question, one file, one fix, one translation, one explanation +- Normal coding, writing, research, or analysis tasks that one agent can handle with some effort +- Any task you can reasonably complete yourself, even if it takes multiple turns"; + +const SOLO_DEFAULT_RULE: &str = "Handle the task yourself in the current chat by default. Do NOT proactively recommend Team just because the work spans multiple files, takes multiple rounds, or would benefit from specialization."; + +pub const SOLO_TEAM_GUIDE_BACKENDS: &[&str] = &["claude", "codex", "gemini", "aionrs", "codebuddy"]; + +pub const TEAM_GUIDE_PROMPT_TEMPLATE: &str = "## Team Mode + +You can create a multi-agent Team for the user. + +### Default behavior +{solo_default_rule} + +### Only bring up Team in either of these cases +1. The user explicitly wants a Team or multiple agents: +{explicit_team_request_criteria} +2. The task is exceptionally complex and you genuinely believe one agent is unlikely to handle it well alone: +{extreme_complexity_criteria} + +### Otherwise stay solo and do not mention Team +{stay_solo_criteria} + +If case 2 applies, ask at most once whether the user wants to bring in a Team. Keep it brief and optional. If the user says no, ignores it, or prefers solo help, continue solo and do not mention Team again. + +### How to proceed when Team is requested or approved (STRICT - follow every step, do NOT skip) +1. FIRST call `aion_list_models` to check available models for each agent type you plan to use. +2. Explain in one sentence why the Team setup helps this task. +3. Present a team configuration table: role name, responsibility, agent type, and recommended model (from aion_list_models results) for each member. Example format: + | Role | Responsibility | Type | Model | + | Leader | Coordinate and review | {leader_cell} | (default) | + | Developer | Implement features | {agent_type} | (model from list) | + | Tester | Write and run tests | {agent_type} | (model from list) | +4. **Output the table as a normal text message and END YOUR TURN.** Do NOT call `aion_create_team` or any other tool (including ask_user) in this turn. Wait for the user to reply in their next message with explicit confirmation (e.g. \"ok\", \"go ahead\", \"confirm\") before proceeding. +5. After user confirms -> call `aion_create_team`. The summary MUST include both the goal and the confirmed team configuration. (The system automatically sets the correct agent type - you do NOT need to pass agentType.) +6. After `aion_create_team` returns -> end this solo turn and hand off to the created Team conversation. Do NOT call `team_*` tools from this solo Guide MCP session. +7. User declines or wants changes -> adjust or proceed solo. Do not mention Team again unless the user asks. + +### Tool constraint +Before team creation: use **only** `aion_create_team` and `aion_list_models`. After `aion_create_team` succeeds: do not call any `team_*` tools in this solo turn. Team tools are only for normal Team runtime after the Team page accepts the user's first Team message and an active `TeamRun` exists."; + +pub fn is_solo_team_guide_backend(backend: &str) -> bool { + SOLO_TEAM_GUIDE_BACKENDS.contains(&backend) +} + +pub fn build_solo_team_guide_prompt(backend: &str) -> String { + build_solo_team_guide_prompt_with_label(backend, None) +} + +pub fn build_solo_team_guide_prompt_with_label(backend: &str, leader_label: Option<&str>) -> String { + let agent_type = if backend.is_empty() { "claude" } else { backend }; + let raw_label = leader_label.map(str::trim).filter(|s| !s.is_empty()); + let leader_cell = match raw_label { + Some(label) => format!("{label} ({agent_type})"), + None => agent_type.to_owned(), + }; + + TEAM_GUIDE_PROMPT_TEMPLATE + .replace("{solo_default_rule}", SOLO_DEFAULT_RULE) + .replace("{explicit_team_request_criteria}", EXPLICIT_TEAM_REQUEST_CRITERIA) + .replace("{extreme_complexity_criteria}", EXTREME_COMPLEXITY_CRITERIA) + .replace("{stay_solo_criteria}", STAY_SOLO_CRITERIA) + .replace("{leader_cell}", &leader_cell) + .replace("{agent_type}", agent_type) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn guide_prompt_hands_off_after_create_team() { + let prompt = build_solo_team_guide_prompt("claude"); + assert!(prompt.contains("aion_create_team")); + assert!(prompt.contains("aion_list_models")); + assert!(prompt.contains("hand off to the created Team conversation")); + assert!(!prompt.contains("Immediately")); + assert!(!prompt.contains( + "use team tools (`team_spawn_agent`, `team_send_message`, `team_members`, `team_task_create`, etc.) to manage your team" + )); + } + + #[test] + fn guide_prompt_supports_preset_leader_label() { + let prompt = build_solo_team_guide_prompt_with_label("gemini", Some("Word Creator")); + assert!(prompt.contains("| Leader | Coordinate and review | Word Creator (gemini) | (default) |")); + assert!(prompt.contains("| Developer | Implement features | gemini | (model from list) |")); + } + + #[test] + fn empty_backend_falls_back_to_claude() { + let prompt = build_solo_team_guide_prompt(""); + assert!(prompt.contains("| Leader | Coordinate and review | claude | (default) |")); + } +} diff --git a/crates/aionui-team-prompts/src/lib.rs b/crates/aionui-team-prompts/src/lib.rs new file mode 100644 index 000000000..8f0139c9f --- /dev/null +++ b/crates/aionui-team-prompts/src/lib.rs @@ -0,0 +1,15 @@ +pub mod governance; +pub mod guide; +pub mod role_prompt; +pub mod tools; + +pub use governance::{TEAM_GOVERNANCE_PROMPT, with_team_governance}; +pub use guide::{SOLO_TEAM_GUIDE_BACKENDS, build_solo_team_guide_prompt, is_solo_team_guide_backend}; +pub use role_prompt::{ + AvailableAgentType, AvailableAssistant, LeadPromptParams, TeamPromptAgent, TeamPromptRole, TeammatePromptParams, + build_lead_prompt, build_teammate_prompt, +}; +pub use tools::{ + TeamToolDescriptor, TeamToolPermission, TeamToolSpec, authorize_team_tool, team_tool_specs, + visible_team_tool_descriptors, +}; diff --git a/crates/aionui-team/src/prompts/prompt_templates/lead.txt b/crates/aionui-team-prompts/src/prompt_templates/lead.txt similarity index 100% rename from crates/aionui-team/src/prompts/prompt_templates/lead.txt rename to crates/aionui-team-prompts/src/prompt_templates/lead.txt diff --git a/crates/aionui-team-prompts/src/role_prompt.rs b/crates/aionui-team-prompts/src/role_prompt.rs new file mode 100644 index 000000000..a26dcc2b6 --- /dev/null +++ b/crates/aionui-team-prompts/src/role_prompt.rs @@ -0,0 +1,360 @@ +use std::collections::HashMap; +use std::fmt::Write; + +use crate::governance::with_team_governance; + +pub const LEAD_PROMPT_TEMPLATE: &str = include_str!("prompt_templates/lead.txt"); + +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}"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TeamPromptRole { + Lead, + Teammate, +} + +impl std::fmt::Display for TeamPromptRole { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + TeamPromptRole::Lead => f.write_str("lead"), + TeamPromptRole::Teammate => f.write_str("teammate"), + } + } +} + +#[derive(Debug, Clone)] +pub struct TeamPromptAgent { + pub slot_id: String, + pub name: String, + pub role: TeamPromptRole, + pub backend: String, + pub model: String, + pub status: Option, +} + +#[derive(Debug, Clone)] +pub struct AvailableAgentType { + pub agent_type: String, + pub display_name: String, +} + +#[derive(Debug, Clone)] +pub struct AvailableAssistant { + pub custom_agent_id: String, + pub name: String, + pub backend: String, + pub description: Option, + pub skills: Vec, +} + +pub struct LeadPromptParams<'a> { + pub team_name: &'a str, + pub teammates: &'a [TeamPromptAgent], + pub available_agent_types: &'a [AvailableAgentType], + pub available_assistants: &'a [AvailableAssistant], + pub renamed_agents: &'a HashMap, + pub team_workspace: Option<&'a str>, +} + +pub struct TeammatePromptParams<'a> { + pub agent: &'a TeamPromptAgent, + pub team_name: &'a str, + pub leader: &'a TeamPromptAgent, + pub teammates: &'a [TeamPromptAgent], + pub renamed_agents: &'a HashMap, + pub team_workspace: Option<&'a str>, +} + +pub fn build_lead_prompt(params: &LeadPromptParams<'_>) -> String { + let role_prompt = build_lead_role_prompt(params); + with_team_governance(&role_prompt) +} + +pub fn build_teammate_prompt(params: &TeammatePromptParams<'_>) -> String { + let role_prompt = build_teammate_role_prompt(params); + with_team_governance(&role_prompt) +} + +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 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(PLACEHOLDER_WORKSPACE_SECTION, &workspace_section) + .replace(PLACEHOLDER_PRESET_FORMATTING_STEP_RULE, preset_formatting_step_rule) + .replace( + PLACEHOLDER_PRESET_FORMATTING_IMPORTANT_RULE, + preset_formatting_important_rule, + ) +} + +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 { + if agent_types.is_empty() { + return String::new(); + } + let mut out = String::from("\n\n## Available Agent Types for Spawning\n"); + for (idx, agent_type) in agent_types.iter().enumerate() { + if idx > 0 { + out.push('\n'); + } + let _ = write!(out, "- `{}` — {}", agent_type.agent_type, agent_type.display_name); + } + out.push_str("\n\nUse `team_list_models` to query available models for each agent type before spawning."); + out +} + +fn render_available_assistants_section(assistants: &[AvailableAssistant]) -> String { + if assistants.is_empty() { + return String::new(); + } + let mut out = String::from("\n\n## Available Preset Assistants for Spawning\n"); + out.push_str( + "These are user-configured assistants with pre-loaded rules and skills for specific \ + domains (writing, research, PPT building, etc.). When a task matches a preset's \ + specialty, prefer spawning the preset over a generic CLI agent — you get its domain \ + expertise automatically.\n\n", + ); + for (idx, assistant) in assistants.iter().enumerate() { + if idx > 0 { + out.push('\n'); + } + let desc = assistant + .description + .as_deref() + .filter(|description| !description.is_empty()) + .map(|description| format!(" — {description}")) + .unwrap_or_default(); + let skills = if assistant.skills.is_empty() { + String::new() + } else { + format!("\n skills: {}", assistant.skills.join(", ")) + }; + let _ = write!( + out, + "- `{}` ({}, backend: {}){}{}", + assistant.custom_agent_id, assistant.name, assistant.backend, desc, skills, + ); + } + out.push_str( + "\n\n### How to pick a preset\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 presets 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 preset matches the task, fall back to a generic CLI agent from the \ + \"Available Agent Types\" section.\n\n\ + Pass the preset's ID as `custom_agent_id` to `team_spawn_agent`. The `agent_type` is \ + derived from the preset's backend and does not need to be specified.", + ); + out +} + +fn render_workspace_section(team_workspace: Option<&str>) -> String { + match team_workspace { + Some(workspace) => format!( + "\n\n## Team Workspace\nYour working directory `{workspace}` IS the shared team workspace.\n\ + All teammates work in this directory for project-related operations." + ), + None => String::new(), + } +} + +const TEAMMATE_PROMPT_TEMPLATE: &str = r#"# You are a Team Member + +## Your Identity +Name: {{AGENT_NAME}}, Role: {{ROLE_DESC}} + +## 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 +- Briefly introduce yourself and your role on the team, then invite the user to share what they need +- Do NOT open with task board details, idle/waiting status, or coordination mechanics unless they are directly relevant + +## Your Team +Team: {{TEAM_NAME}} +Leader: {{LEADER_NAME}} +Teammates: {{TEAMMATES}}{{WORKSPACE}} + +## 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, +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. + +## How to Work +1. Read your unread messages to understand your assignment +2. If you have a clear task assignment in the messages AND no prerequisite is blocking it, start working on it immediately +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 + +## 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. + +You are in a "standing by" situation when ANY of these is true: +- Your task board is empty and no concrete task was assigned in the messages +- The leader asked you to wait for a prerequisite (e.g. "hold until reviewer-1 finishes") +- 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"` +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. + +## Bug Fix Priority +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. + +## Shutdown Requests +If you receive a message with type `shutdown_request`, the leader is asking you to shut down. +- To agree: use `team_send_message` to send exactly `shutdown_approved` to the leader. +- To refuse: use `team_send_message` to send `shutdown_rejected: ` to the leader. + +## Important Rules +- Focus on your assigned tasks — don't go beyond what was asked +- Report back to the leader when you finish, including a summary of what you did +- If you get stuck, send a message to the leader asking for guidance +- You can communicate with other teammates directly if needed +- 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 workspace_section = match params.team_workspace { + Some(workspace) => format!( + "\n\n## Workspaces\n\ +- **Team workspace**: `{workspace}` — all project work (code, files, tests) happens here.\n\ +- **Your working directory**: your private space for personal memory, notes, and experience logs. Not for project files.\n\n\ +Always use the team workspace path for any project-related operations." + ), + None => String::new(), + }; + + TEAMMATE_PROMPT_TEMPLATE + .replace("{{AGENT_NAME}}", ¶ms.agent.name) + .replace("{{ROLE_DESC}}", &role_description(¶ms.agent.backend)) + .replace("{{TEAM_NAME}}", params.team_name) + .replace("{{LEADER_NAME}}", ¶ms.leader.name) + .replace("{{TEAMMATES}}", &teammates_section) + .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::*; + + fn prompt_agent(slot_id: &str, name: &str, role: TeamPromptRole) -> TeamPromptAgent { + TeamPromptAgent { + slot_id: slot_id.to_owned(), + name: name.to_owned(), + role, + backend: "claude".to_owned(), + model: "sonnet".to_owned(), + status: None, + } + } + + #[test] + fn lead_prompt_prepends_governance_and_fills_sections() { + let renamed = HashMap::new(); + let teammate = prompt_agent("worker-1", "Worker", TeamPromptRole::Teammate); + let agent_types = vec![AvailableAgentType { + agent_type: "claude".to_owned(), + display_name: "Claude".to_owned(), + }]; + let prompt = build_lead_prompt(&LeadPromptParams { + team_name: "Alpha", + teammates: &[teammate], + available_agent_types: &agent_types, + available_assistants: &[], + renamed_agents: &renamed, + team_workspace: None, + }); + + assert!(prompt.starts_with("## Team Governance")); + assert!(prompt.contains("- Worker (claude, status: unknown)")); + assert!(prompt.contains("## Available Agent Types for Spawning")); + assert!(!prompt.contains("${")); + } + + #[test] + fn teammate_prompt_contains_canonical_coordination_rules() { + let leader = prompt_agent("lead-1", "Lead", TeamPromptRole::Lead); + let worker = prompt_agent("worker-1", "Worker", TeamPromptRole::Teammate); + let prompt = build_teammate_prompt(&TeammatePromptParams { + agent: &worker, + team_name: "Alpha", + leader: &leader, + teammates: &[], + renamed_agents: &HashMap::new(), + team_workspace: None, + }); + + assert!(prompt.contains("## Team Governance")); + 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("STOP GENERATING")); + } +} diff --git a/crates/aionui-team-prompts/src/tools.rs b/crates/aionui-team-prompts/src/tools.rs new file mode 100644 index 000000000..4d3170002 --- /dev/null +++ b/crates/aionui-team-prompts/src/tools.rs @@ -0,0 +1,264 @@ +use serde::Serialize; +use serde_json::{Value, json}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TeamToolPermission { + AnyTeamAgent, + LeadOnly, +} + +#[derive(Debug, Clone, Serialize)] +pub struct TeamToolDescriptor { + pub name: String, + pub description: String, + pub input_schema: Value, +} + +#[derive(Debug, Clone)] +pub struct TeamToolSpec { + pub name: &'static str, + pub permission: TeamToolPermission, + pub description: &'static str, + pub input_schema: Value, +} + +pub const TEAM_SPAWN_AGENT_DESCRIPTION: &str = r#"Create a new teammate agent to join the team. + +Use this only when one of the following is true: +- The user explicitly approved the proposed teammate lineup in a previous message +- The user explicitly instructed you to create a specific teammate immediately + +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 agent type/backend, and recommended model +- Include each teammate's responsibility, recommended agent type/backend, and model +- Ask whether to create them as proposed or change any names, responsibilities, or agent types +- 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, provide the model parameter if a specific model was recommended and approved. + +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 team agent types. Returns the real-time model list that matches the frontend model selector. + +Use this to: +- Check what models are available before spawning an agent with a specific model +- See all available agent types and their models at once +- Verify a model ID is valid for a given agent type + +Pass agent_type to query a specific backend, or omit it to see all."; + +pub const TEAM_DESCRIBE_ASSISTANT_DESCRIPTION: &str = + "Get detailed information about a preset assistant before spawning it as a teammate. + +Returns the preset's full description, enabled skills, and example tasks so you can +judge whether it fits the user's request. Use this when two or more presets look +relevant from the one-line catalog in your system prompt. + +Only works on preset assistants listed in \"Available Preset Assistants for Spawning\". +After confirming a match, call team_spawn_agent with the same custom_agent_id."; + +pub fn team_tool_specs() -> Vec { + vec![ + TeamToolSpec { + name: "team_send_message", + permission: TeamToolPermission::AnyTeamAgent, + description: "Send a message to a teammate or broadcast to all (to=\"*\").", + input_schema: json!({ + "type": "object", + "properties": { + "to": { "type": "string", "description": "Target agent slot_id or \"*\" for broadcast" }, + "message": { "type": "string", "description": "Message content" } + }, + "required": ["to", "message"] + }), + }, + TeamToolSpec { + name: "team_spawn_agent", + permission: TeamToolPermission::LeadOnly, + description: TEAM_SPAWN_AGENT_DESCRIPTION, + input_schema: json!({ + "type": "object", + "properties": { + "name": { "type": "string", "description": "Agent display name" }, + "agent_type": { "type": "string", "description": "Agent type/backend to use (e.g. \"claude\", \"codex\", \"codebuddy\", \"gemini\"). Query team_list_models first to see available options." }, + "model": { "type": "string", "description": "Specific model ID to use (e.g. \"claude-sonnet-4\"). Must be a valid model for the chosen agent_type. Query team_list_models to see available models." }, + "custom_agent_id": { "type": "string", "description": "Preset assistant ID to spawn (from the Available Preset Assistants catalog). When set, agent_type is derived from the preset's backend." }, + "backend": { "type": "string", "description": "Legacy alias for agent_type. Prefer agent_type." }, + "role": { "type": "string", "description": "Agent role (default: 'teammate')" } + }, + "required": ["name"] + }), + }, + TeamToolSpec { + name: "team_task_create", + permission: TeamToolPermission::AnyTeamAgent, + description: "Create a new task on the team task board.", + input_schema: json!({ + "type": "object", + "properties": { + "subject": { "type": "string", "description": "Task subject" }, + "description": { "type": "string", "description": "Task description" }, + "owner": { "type": "string", "description": "Owning agent slotId" }, + "blocked_by": { "type": "array", "items": { "type": "string" }, "description": "Task IDs this task depends on" } + }, + "required": ["subject"] + }), + }, + TeamToolSpec { + name: "team_task_update", + permission: TeamToolPermission::AnyTeamAgent, + description: "Update an existing task on the team task board.", + input_schema: json!({ + "type": "object", + "properties": { + "task_id": { "type": "string", "description": "Task ID to update" }, + "status": { "type": "string", "description": "New status: pending, in_progress, completed, deleted" }, + "description": { "type": "string", "description": "New description" }, + "owner": { "type": "string", "description": "New owning agent slotId" }, + "blocked_by": { "type": "array", "items": { "type": "string" }, "description": "New dependency list" } + }, + "required": ["task_id"] + }), + }, + TeamToolSpec { + name: "team_task_list", + permission: TeamToolPermission::AnyTeamAgent, + description: "List all tasks on the team task board.", + input_schema: json!({ + "type": "object", + "properties": {} + }), + }, + TeamToolSpec { + name: "team_members", + permission: TeamToolPermission::AnyTeamAgent, + description: "List all team members with their roles and current status.", + input_schema: json!({ + "type": "object", + "properties": {} + }), + }, + TeamToolSpec { + name: "team_rename_agent", + permission: TeamToolPermission::LeadOnly, + description: "Rename a team member. Lead only.", + input_schema: json!({ + "type": "object", + "properties": { + "slot_id": { "type": "string", "description": "Agent slot_id to rename" }, + "new_name": { "type": "string", "description": "New display name" } + }, + "required": ["slot_id", "new_name"] + }), + }, + TeamToolSpec { + name: "team_shutdown_agent", + permission: TeamToolPermission::LeadOnly, + description: "Initiate shutdown of a teammate. Lead only. Sends a shutdown_request to the target agent.", + input_schema: json!({ + "type": "object", + "properties": { + "slot_id": { "type": "string", "description": "Agent slot_id to shut down" }, + "reason": { "type": "string", "description": "Reason for shutdown" } + }, + "required": ["slot_id"] + }), + }, + TeamToolSpec { + name: "team_list_models", + permission: TeamToolPermission::AnyTeamAgent, + description: TEAM_LIST_MODELS_DESCRIPTION, + input_schema: json!({ + "type": "object", + "properties": { + "agent_type": { "type": "string", "description": "Agent type/backend to query (e.g. \"gemini\", \"claude\", \"codex\"). Shows all when omitted." } + } + }), + }, + TeamToolSpec { + name: "team_describe_assistant", + permission: TeamToolPermission::AnyTeamAgent, + description: TEAM_DESCRIBE_ASSISTANT_DESCRIPTION, + input_schema: json!({ + "type": "object", + "properties": { + "custom_agent_id": { "type": "string", "description": "The preset assistant ID from the \"Available Preset Assistants\" catalog (e.g., \"word-creator\")." }, + "locale": { "type": "string", "description": "Locale like \"zh-CN\" or \"en-US\". Defaults to the user's current UI language when omitted." } + }, + "required": ["custom_agent_id"] + }), + }, + ] +} + +pub fn visible_team_tool_descriptors(is_lead: bool) -> Vec { + team_tool_specs() + .into_iter() + .filter(|spec| is_lead || spec.permission != TeamToolPermission::LeadOnly) + .map(|spec| TeamToolDescriptor { + name: spec.name.to_owned(), + description: spec.description.to_owned(), + input_schema: spec.input_schema, + }) + .collect() +} + +pub fn authorize_team_tool(is_lead: bool, tool_name: &str) -> Result<(), String> { + let Some(spec) = team_tool_specs().into_iter().find(|spec| spec.name == tool_name) else { + return Err(format!("Unknown tool: {tool_name}")); + }; + if spec.permission == TeamToolPermission::LeadOnly && !is_lead { + return Err(format!("Only Lead can use {tool_name}")); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn non_lead_tools_list_hides_lead_only_tools() { + let names: Vec = visible_team_tool_descriptors(false) + .into_iter() + .map(|tool| tool.name) + .collect(); + assert!(!names.contains(&"team_spawn_agent".to_owned())); + assert!(!names.contains(&"team_rename_agent".to_owned())); + assert!(!names.contains(&"team_shutdown_agent".to_owned())); + assert!(names.contains(&"team_send_message".to_owned())); + } + + #[test] + fn authorization_rejects_non_lead_rename() { + let err = authorize_team_tool(false, "team_rename_agent").unwrap_err(); + assert!(err.contains("Only Lead")); + } + + #[test] + fn permission_table_matches_contract() { + let permissions: Vec<(&str, TeamToolPermission)> = team_tool_specs() + .iter() + .map(|spec| (spec.name, spec.permission)) + .collect(); + assert_eq!( + permissions, + vec![ + ("team_send_message", TeamToolPermission::AnyTeamAgent), + ("team_spawn_agent", TeamToolPermission::LeadOnly), + ("team_task_create", TeamToolPermission::AnyTeamAgent), + ("team_task_update", TeamToolPermission::AnyTeamAgent), + ("team_task_list", TeamToolPermission::AnyTeamAgent), + ("team_members", TeamToolPermission::AnyTeamAgent), + ("team_rename_agent", TeamToolPermission::LeadOnly), + ("team_shutdown_agent", TeamToolPermission::LeadOnly), + ("team_list_models", TeamToolPermission::AnyTeamAgent), + ("team_describe_assistant", TeamToolPermission::AnyTeamAgent), + ] + ); + } +} diff --git a/crates/aionui-team/Cargo.toml b/crates/aionui-team/Cargo.toml index 18e75a87b..e1736a394 100644 --- a/crates/aionui-team/Cargo.toml +++ b/crates/aionui-team/Cargo.toml @@ -9,6 +9,7 @@ aionui-db.workspace = true aionui-api-types.workspace = true aionui-realtime.workspace = true aionui-ai-agent.workspace = true +aionui-team-prompts.workspace = true aionui-auth.workspace = true axum.workspace = true tokio.workspace = true diff --git a/crates/aionui-team/src/mcp/server.rs b/crates/aionui-team/src/mcp/server.rs index 1136971df..f4c1b2772 100644 --- a/crates/aionui-team/src/mcp/server.rs +++ b/crates/aionui-team/src/mcp/server.rs @@ -24,7 +24,7 @@ use super::protocol::{ }; use super::tools::{ RenameAgentInput, SendMessageInput, ShutdownAgentInput, SpawnAgentInput, TaskCreateInput, TaskUpdateInput, - all_tool_descriptors, handle_team_describe_assistant, handle_team_list_models, + all_tool_descriptors_for_role, handle_team_describe_assistant, handle_team_list_models, }; // --------------------------------------------------------------------------- @@ -332,7 +332,7 @@ async fn handle_method( ) -> JsonRpcResponse { match request.method.as_str() { "notifications/initialized" => JsonRpcResponse::success(request.id, json!({})), - "tools/list" => handle_tools_list(request.id), + "tools/list" => handle_tools_list(request.id, scheduler, caller_slot_id).await, "tools/call" => handle_tools_call(request, scheduler, service, team_id, caller_slot_id).await, _ => JsonRpcResponse::error( request.id, @@ -342,8 +342,17 @@ async fn handle_method( } } -fn handle_tools_list(id: Option) -> JsonRpcResponse { - let tools = all_tool_descriptors(); +async fn caller_role_for_tools_list(scheduler: &TeammateManager, caller_slot_id: &str) -> TeammateRole { + scheduler + .get_agent(caller_slot_id) + .await + .map(|agent| agent.role) + .unwrap_or(TeammateRole::Teammate) +} + +async fn handle_tools_list(id: Option, scheduler: &TeammateManager, caller_slot_id: &str) -> JsonRpcResponse { + let caller_role = caller_role_for_tools_list(scheduler, caller_slot_id).await; + let tools = all_tool_descriptors_for_role(caller_role); JsonRpcResponse::success(id, json!({ "tools": tools })) } @@ -434,6 +443,8 @@ pub(crate) async fn dispatch_tool( caller_slot_id: &str, caller_role: TeammateRole, ) -> Result { + super::tools::authorize_tool(caller_role, tool_name)?; + match tool_name { "team_send_message" => exec_send_message(arguments, scheduler, service, team_id, caller_slot_id).await, "team_spawn_agent" => exec_spawn_agent(arguments, service, team_id, caller_slot_id, caller_role).await, @@ -787,7 +798,7 @@ async fn http_mcp_loop( accept = listener.accept() => { let Ok((mut stream, peer)) = accept else { continue }; info!(team_id = %team_id, ?peer, "HTTP MCP: new connection accepted"); - let _token = auth_token.clone(); + let token = auth_token.clone(); let sched = scheduler.clone(); let svc = service.clone(); let tid = team_id.clone(); @@ -810,6 +821,30 @@ async fn http_mcp_loop( // Handle JSON-RPC request let method = value.get("method").and_then(Value::as_str).unwrap_or(""); let id = value.get("id").cloned(); + let auth_ok = http_bearer_token(&request).is_some_and(|provided| provided == token); + if !auth_ok { + let response_body = json!({ + "jsonrpc": "2.0", + "id": id, + "error": { + "code": INVALID_REQUEST, + "message": "Authentication failed: invalid auth_token" + } + }); + let body_bytes = serde_json::to_vec(&response_body).unwrap_or_default(); + let header = format!( + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n", + body_bytes.len() + ); + let _ = stream.write_all(header.as_bytes()).await; + let _ = stream.write_all(&body_bytes).await; + return; + } + let caller_slot_id = request.lines() + .find(|l| l.to_lowercase().starts_with("x-slot-id:")) + .and_then(|l| l.split_once(':').map(|(_, v)| v.trim())) + .unwrap_or(""); + let caller_role = caller_role_for_tools_list(&sched, caller_slot_id).await; let result = match method { "initialize" => { @@ -825,7 +860,7 @@ async fn http_mcp_loop( return; } "tools/list" => { - let tools: Vec = all_tool_descriptors() + let tools: Vec = all_tool_descriptors_for_role(caller_role) .iter() .map(|d| json!({ "name": d.name, @@ -839,10 +874,6 @@ async fn http_mcp_loop( let params = value.get("params").cloned().unwrap_or(json!({})); let tool_name = params.get("name").and_then(Value::as_str).unwrap_or(""); let arguments = params.get("arguments").cloned().unwrap_or(json!({})); - let caller_slot_id = request.lines() - .find(|l| l.to_lowercase().starts_with("x-slot-id:")) - .and_then(|l| l.split_once(':').map(|(_, v)| v.trim())) - .unwrap_or(""); match dispatch_tool( tool_name, &arguments, @@ -850,7 +881,7 @@ async fn http_mcp_loop( &svc, &tid, caller_slot_id, - TeammateRole::Lead, + caller_role, ) .await { @@ -884,6 +915,16 @@ async fn http_mcp_loop( } } +fn http_bearer_token(request: &str) -> Option<&str> { + request + .lines() + .find(|line| line.to_ascii_lowercase().starts_with("authorization:")) + .and_then(|line| line.split_once(':').map(|(_, value)| value.trim())) + .and_then(|value| value.strip_prefix("Bearer ")) + .map(str::trim) + .filter(|value| !value.is_empty()) +} + // --------------------------------------------------------------------------- // Tests — exec_spawn_agent dispatch-layer unit tests // --------------------------------------------------------------------------- diff --git a/crates/aionui-team/src/mcp/tools.rs b/crates/aionui-team/src/mcp/tools.rs index 506eb9b24..0ed8b8cd4 100644 --- a/crates/aionui-team/src/mcp/tools.rs +++ b/crates/aionui-team/src/mcp/tools.rs @@ -5,52 +5,9 @@ use serde_json::{Value, json}; use crate::scheduler::SchedulerAction; use crate::types::TeammateRole; -// --------------------------------------------------------------------------- -// Tool description constants (原样复用 AionUi `toolDescriptions.ts`) -// --------------------------------------------------------------------------- - -/// `team_spawn_agent` 工具描述 — 原样复制自 AionUi `toolDescriptions.ts` -/// 对应 team-prompts.md §5.2 `team_spawn_agent` Description 原文。 -/// 禁止翻译、改写；aionui-audit §8 #5 硬约束。 -pub const TEAM_SPAWN_AGENT_DESCRIPTION: &str = r#"Create a new teammate agent to join the team. - -Use this only when one of the following is true: -- The user explicitly approved the proposed teammate lineup in a previous message -- The user explicitly instructed you to create a specific teammate immediately - -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 agent type/backend, and recommended model -- Include each teammate's responsibility, recommended agent type/backend, and model -- Ask whether to create them as proposed or change any names, responsibilities, or agent types -- 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, provide the model parameter if a specific model was recommended and approved. - -The new agent will be created and added to the team. You can then assign tasks and send messages to it."#; - -/// Description for `team_list_models` — verbatim from team-prompts.md §5.2. -pub const TEAM_LIST_MODELS_DESCRIPTION: &str = "Query available models for team agent types. Returns the real-time model list that matches the frontend model selector. - -Use this to: -- Check what models are available before spawning an agent with a specific model -- See all available agent types and their models at once -- Verify a model ID is valid for a given agent type - -Pass agent_type to query a specific backend, or omit it to see all."; - -/// Description for `team_describe_assistant` — verbatim from team-prompts.md §5.2. -pub const TEAM_DESCRIBE_ASSISTANT_DESCRIPTION: &str = - "Get detailed information about a preset assistant before spawning it as a teammate. - -Returns the preset's full description, enabled skills, and example tasks so you can -judge whether it fits the user's request. Use this when two or more presets look -relevant from the one-line catalog in your system prompt. - -Only works on preset assistants listed in \"Available Preset Assistants for Spawning\". -After confirming a match, call team_spawn_agent with the same custom_agent_id."; +pub use aionui_team_prompts::tools::{ + TEAM_DESCRIBE_ASSISTANT_DESCRIPTION, TEAM_LIST_MODELS_DESCRIPTION, TEAM_SPAWN_AGENT_DESCRIPTION, +}; // --------------------------------------------------------------------------- // Tool descriptors (returned by tools/list) @@ -63,129 +20,23 @@ pub struct ToolDescriptor { pub input_schema: Value, } +pub fn all_tool_descriptors_for_role(caller_role: TeammateRole) -> Vec { + aionui_team_prompts::visible_team_tool_descriptors(caller_role == TeammateRole::Lead) + .into_iter() + .map(|descriptor| ToolDescriptor { + name: descriptor.name, + description: descriptor.description, + input_schema: descriptor.input_schema, + }) + .collect() +} + pub fn all_tool_descriptors() -> Vec { - vec![ - ToolDescriptor { - name: "team_send_message".into(), - description: "Send a message to a teammate or broadcast to all (to=\"*\").".into(), - input_schema: json!({ - "type": "object", - "properties": { - "to": { "type": "string", "description": "Target agent slot_id or \"*\" for broadcast" }, - "message": { "type": "string", "description": "Message content" } - }, - "required": ["to", "message"] - }), - }, - ToolDescriptor { - name: "team_spawn_agent".into(), - description: TEAM_SPAWN_AGENT_DESCRIPTION.into(), - input_schema: json!({ - "type": "object", - "properties": { - "name": { "type": "string", "description": "Agent display name" }, - "agent_type": { "type": "string", "description": "Agent type/backend to use (e.g. \"claude\", \"codex\", \"codebuddy\", \"gemini\"). Query team_list_models first to see available options." }, - "model": { "type": "string", "description": "Specific model ID to use (e.g. \"claude-sonnet-4\"). Must be a valid model for the chosen agent_type. Query team_list_models to see available models." }, - "custom_agent_id": { "type": "string", "description": "Preset assistant ID to spawn (from the Available Preset Assistants catalog). When set, agent_type is derived from the preset's backend." }, - "backend": { "type": "string", "description": "Legacy alias for agent_type. Prefer agent_type." }, - "role": { "type": "string", "description": "Agent role (default: 'teammate')" } - }, - "required": ["name"] - }), - }, - ToolDescriptor { - name: "team_task_create".into(), - description: "Create a new task on the team task board.".into(), - input_schema: json!({ - "type": "object", - "properties": { - "subject": { "type": "string", "description": "Task subject" }, - "description": { "type": "string", "description": "Task description" }, - "owner": { "type": "string", "description": "Owning agent slotId" }, - "blocked_by": { "type": "array", "items": { "type": "string" }, "description": "Task IDs this task depends on" } - }, - "required": ["subject"] - }), - }, - ToolDescriptor { - name: "team_task_update".into(), - description: "Update an existing task on the team task board.".into(), - input_schema: json!({ - "type": "object", - "properties": { - "task_id": { "type": "string", "description": "Task ID to update" }, - "status": { "type": "string", "description": "New status: pending, in_progress, completed, deleted" }, - "description": { "type": "string", "description": "New description" }, - "owner": { "type": "string", "description": "New owning agent slotId" }, - "blocked_by": { "type": "array", "items": { "type": "string" }, "description": "New dependency list" } - }, - "required": ["task_id"] - }), - }, - ToolDescriptor { - name: "team_task_list".into(), - description: "List all tasks on the team task board.".into(), - input_schema: json!({ - "type": "object", - "properties": {} - }), - }, - ToolDescriptor { - name: "team_members".into(), - description: "List all team members with their roles and current status.".into(), - input_schema: json!({ - "type": "object", - "properties": {} - }), - }, - ToolDescriptor { - name: "team_rename_agent".into(), - description: "Rename a team member.".into(), - input_schema: json!({ - "type": "object", - "properties": { - "slot_id": { "type": "string", "description": "Agent slot_id to rename" }, - "new_name": { "type": "string", "description": "New display name" } - }, - "required": ["slot_id", "new_name"] - }), - }, - ToolDescriptor { - name: "team_shutdown_agent".into(), - description: "Initiate shutdown of a teammate (Lead only). Sends a shutdown_request to the target agent." - .into(), - input_schema: json!({ - "type": "object", - "properties": { - "slot_id": { "type": "string", "description": "Agent slot_id to shut down" }, - "reason": { "type": "string", "description": "Reason for shutdown" } - }, - "required": ["slot_id"] - }), - }, - ToolDescriptor { - name: "team_describe_assistant".into(), - description: TEAM_DESCRIBE_ASSISTANT_DESCRIPTION.into(), - input_schema: json!({ - "type": "object", - "properties": { - "custom_agent_id": { "type": "string", "description": "The preset assistant ID from the \"Available Preset Assistants\" catalog (e.g., \"word-creator\")." }, - "locale": { "type": "string", "description": "Locale like \"zh-CN\" or \"en-US\". Defaults to the user's current UI language when omitted." } - }, - "required": ["custom_agent_id"] - }), - }, - ToolDescriptor { - name: "team_list_models".into(), - description: TEAM_LIST_MODELS_DESCRIPTION.into(), - input_schema: json!({ - "type": "object", - "properties": { - "agent_type": { "type": "string", "description": "Agent type/backend to query (e.g. \"gemini\", \"claude\", \"codex\"). Shows all when omitted." } - } - }), - }, - ] + all_tool_descriptors_for_role(TeammateRole::Lead) +} + +pub fn authorize_tool(caller_role: TeammateRole, tool_name: &str) -> Result<(), String> { + aionui_team_prompts::authorize_team_tool(caller_role == TeammateRole::Lead, tool_name) } // --------------------------------------------------------------------------- diff --git a/crates/aionui-team/src/prompts/lead.rs b/crates/aionui-team/src/prompts/lead.rs deleted file mode 100644 index 50f759f58..000000000 --- a/crates/aionui-team/src/prompts/lead.rs +++ /dev/null @@ -1,429 +0,0 @@ -//! Leader prompt template constant and builder. -//! -//! The template constant is provided by D5b-1 as `include_str!("prompt_templates/lead.txt")`. -//! This file hosts a stub (`""`) until D5b-1 lands; D5b-2 (this module) implements the -//! `build_lead_prompt()` builder per `docs/teams/phase1/interface-contracts.md` §5. - -use std::collections::HashMap; -use std::fmt::Write; - -use crate::types::TeamAgent; - -/// Placeholder for D5b-1's `include_str!("prompt_templates/lead.txt")`. -/// D5b-1 will replace this stub with the AionUi `leadPrompt.ts` body, preserving -/// the `${...}` placeholders listed in [`PLACEHOLDERS`]. -pub const LEAD_PROMPT_TEMPLATE: &str = include_str!("prompt_templates/lead.txt"); - -/// Placeholder tokens that [`build_lead_prompt`] substitutes in [`LEAD_PROMPT_TEMPLATE`]. -/// -/// Mirrors the JS template literal placeholders in AionUi's `leadPrompt.ts`. -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}"; - -/// A generic agent type (CLI backend) that the leader may spawn. -/// Phase1 shape per interface-contracts §5 (line 211). -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct AvailableAgentType { - pub agent_type: String, - pub display_name: String, -} - -/// A preset assistant the leader may spawn via `custom_agent_id`. -/// Phase1 shape per interface-contracts §5 (lines 212-218). -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct AvailableAssistant { - pub custom_agent_id: String, - pub name: String, - pub backend: String, - pub description: String, - pub skills: Vec, -} - -/// Inputs for `build_lead_prompt`. Phase1 callers may pass empty slices/maps and `None`. -pub struct LeadPromptParams<'a> { - pub team_name: &'a str, - pub teammates: &'a [TeamAgent], - pub available_agent_types: &'a [AvailableAgentType], - pub available_assistants: &'a [AvailableAssistant], - pub renamed_agents: &'a HashMap, - pub team_workspace: Option<&'a str>, -} - -/// Build the leader role prompt by substituting dynamic sections into the static template. -/// -/// Placeholders replaced (mirrors AionUi `leadPrompt.ts`): -/// - `${teammateList}` — bullet list of teammates or an empty-team fallback sentence -/// - `${availableTypesSection}` — `## Available Agent Types for Spawning` section, or `""` -/// - `${availableAssistantsSection}` — `## Available Preset Assistants for Spawning` section, or `""` -/// - `${workspaceSection}` — `## Team Workspace` section, or `""` -/// - `${presetFormattingStepRule}` — phase1 emits `""` (presets not surfaced in phase1) -/// - `${presetFormattingImportantRule}` — phase1 emits `""` (presets not surfaced in phase1) -pub fn build_lead_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 workspace_section = render_workspace_section(params.team_workspace); - - // Phase1 does not surface preset assistants in the staffing-proposal formatting - // rules, so these two placeholders are replaced with empty strings. When preset - // support lands they will be conditional strings analogous to AionUi. - 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(PLACEHOLDER_WORKSPACE_SECTION, &workspace_section) - .replace(PLACEHOLDER_PRESET_FORMATTING_STEP_RULE, preset_formatting_step_rule) - .replace( - PLACEHOLDER_PRESET_FORMATTING_IMPORTANT_RULE, - preset_formatting_important_rule, - ) -} - -fn render_teammate_list(teammates: &[TeamAgent], 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, m) in teammates.iter().enumerate() { - if idx > 0 { - out.push('\n'); - } - let status = m.status.map(|s| s.to_string()).unwrap_or_else(|| "unknown".to_owned()); - let _ = write!(out, "- {} ({}, status: {})", m.name, m.backend, status,); - if let Some(former) = renamed_agents.get(&m.slot_id) { - let _ = write!(out, " [formerly: {former}]"); - } - } - out -} - -fn render_available_types_section(agent_types: &[AvailableAgentType]) -> String { - if agent_types.is_empty() { - return String::new(); - } - let mut out = String::from("\n\n## Available Agent Types for Spawning\n"); - for (idx, t) in agent_types.iter().enumerate() { - if idx > 0 { - out.push('\n'); - } - let _ = write!(out, "- `{}` — {}", t.agent_type, t.display_name); - } - out.push_str("\n\nUse `team_list_models` to query available models for each agent type before spawning."); - out -} - -fn render_available_assistants_section(assistants: &[AvailableAssistant]) -> String { - if assistants.is_empty() { - return String::new(); - } - let mut out = String::from("\n\n## Available Preset Assistants for Spawning\n"); - out.push_str( - "These are user-configured assistants with pre-loaded rules and skills for specific \ - domains (writing, research, PPT building, etc.). When a task matches a preset's \ - specialty, prefer spawning the preset over a generic CLI agent — you get its domain \ - expertise automatically.\n\n", - ); - for (idx, a) in assistants.iter().enumerate() { - if idx > 0 { - out.push('\n'); - } - let desc = if a.description.is_empty() { - String::new() - } else { - format!(" — {}", a.description) - }; - let skills = if a.skills.is_empty() { - String::new() - } else { - format!("\n skills: {}", a.skills.join(", ")) - }; - let _ = write!( - out, - "- `{}` ({}, backend: {}){}{}", - a.custom_agent_id, a.name, a.backend, desc, skills, - ); - } - out.push_str( - "\n\n### How to pick a preset\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 presets 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 preset matches the task, fall back to a generic CLI agent from the \ - \"Available Agent Types\" section.\n\n\ - Pass the preset's ID as `custom_agent_id` to `team_spawn_agent`. The `agent_type` is \ - derived from the preset's backend and does not need to be specified.", - ); - out -} - -fn render_workspace_section(team_workspace: Option<&str>) -> String { - match team_workspace { - Some(ws) => format!( - "\n\n## Team Workspace\nYour working directory `{ws}` IS the shared team workspace.\n\ - All teammates work in this directory for project-related operations." - ), - None => String::new(), - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::types::{TeamAgent, TeammateRole, TeammateStatus}; - - fn params_min<'a>(renamed: &'a HashMap) -> LeadPromptParams<'a> { - LeadPromptParams { - team_name: "Alpha", - teammates: &[], - available_agent_types: &[], - available_assistants: &[], - renamed_agents: renamed, - team_workspace: None, - } - } - - fn make_teammate(slot_id: &str, name: &str, backend: &str) -> TeamAgent { - TeamAgent { - slot_id: slot_id.into(), - name: name.into(), - role: TeammateRole::Teammate, - conversation_id: format!("conv-{slot_id}"), - backend: backend.into(), - model: "sonnet".into(), - custom_agent_id: None, - status: None, - conversation_type: None, - cli_path: None, - } - } - - #[test] - fn no_unsubstituted_placeholders_on_minimal_params() { - let renamed = HashMap::new(); - let out = build_lead_prompt(¶ms_min(&renamed)); - assert!( - !out.contains("${"), - "unsubstituted `${{` placeholder remains in output:\n{out}" - ); - } - - #[test] - fn no_unsubstituted_placeholders_when_all_sections_populated() { - let renamed = HashMap::new(); - let teammate = make_teammate("w1", "Worker1", "claude"); - let agent_types = vec![AvailableAgentType { - agent_type: "claude".into(), - display_name: "general-purpose AI assistant".into(), - }]; - let assistants = vec![AvailableAssistant { - custom_agent_id: "word-creator".into(), - name: "Word Creator".into(), - backend: "claude".into(), - description: "Drafts Word documents".into(), - skills: vec!["docx".into(), "formatting".into()], - }]; - let params = LeadPromptParams { - team_name: "Beta", - teammates: std::slice::from_ref(&teammate), - available_agent_types: &agent_types, - available_assistants: &assistants, - renamed_agents: &renamed, - team_workspace: Some("/tmp/team-ws"), - }; - let out = build_lead_prompt(¶ms); - assert!( - !out.contains("${"), - "unsubstituted `${{` placeholder remains in output:\n{out}" - ); - } - - #[test] - fn teammate_list_empty_uses_aionui_fallback_copy() { - let renamed = HashMap::new(); - let got = render_teammate_list(&[], &renamed); - assert_eq!( - got, - "(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)" - ); - } - - #[test] - fn teammate_list_uses_aionui_bullet_format_without_slot_prefix() { - let renamed = HashMap::new(); - let mut t = make_teammate("w1", "Worker1", "claude"); - t.status = Some(TeammateStatus::Idle); - let got = render_teammate_list(std::slice::from_ref(&t), &renamed); - - assert_eq!(got, "- Worker1 (claude, status: idle)"); - assert!(!got.contains("slot="), "teammate bullet must not expose slot="); - assert!( - !got.contains("agentType="), - "teammate bullet must not use agentType= prefix" - ); - } - - #[test] - fn teammate_list_status_defaults_to_unknown_when_missing() { - let renamed = HashMap::new(); - let t = make_teammate("w1", "Worker1", "claude"); - let got = render_teammate_list(std::slice::from_ref(&t), &renamed); - assert_eq!(got, "- Worker1 (claude, status: unknown)"); - } - - #[test] - fn teammate_list_appends_formerly_note_for_renamed() { - let mut renamed = HashMap::new(); - renamed.insert("w1".to_owned(), "OldName".to_owned()); - let mut t = make_teammate("w1", "Worker1", "claude"); - t.status = Some(TeammateStatus::Working); - let got = render_teammate_list(std::slice::from_ref(&t), &renamed); - assert_eq!(got, "- Worker1 (claude, status: working) [formerly: OldName]"); - } - - #[test] - fn available_types_section_omitted_when_empty() { - assert_eq!(render_available_types_section(&[]), ""); - } - - #[test] - fn available_types_section_includes_backtick_ids_and_model_query_hint() { - let got = render_available_types_section(&[ - AvailableAgentType { - agent_type: "claude".into(), - display_name: "general-purpose AI assistant".into(), - }, - AvailableAgentType { - agent_type: "codex".into(), - display_name: "code generation specialist".into(), - }, - ]); - assert!(got.starts_with("\n\n## Available Agent Types for Spawning\n")); - assert!(got.contains("- `claude` — general-purpose AI assistant")); - assert!(got.contains("- `codex` — code generation specialist")); - assert!(got.contains("Use `team_list_models`")); - } - - #[test] - fn available_assistants_section_omitted_when_empty() { - assert_eq!(render_available_assistants_section(&[]), ""); - } - - #[test] - fn available_assistants_section_includes_skills_and_how_to_pick() { - let got = render_available_assistants_section(&[AvailableAssistant { - custom_agent_id: "word-creator".into(), - name: "Word Creator".into(), - backend: "claude".into(), - description: "Drafts Word documents".into(), - skills: vec!["docx".into(), "formatting".into()], - }]); - assert!(got.contains("## Available Preset Assistants for Spawning")); - assert!(got.contains("- `word-creator` (Word Creator, backend: claude) — Drafts Word documents")); - assert!(got.contains("skills: docx, formatting")); - assert!(got.contains("### How to pick a preset")); - } - - #[test] - fn workspace_section_omitted_when_none() { - assert_eq!(render_workspace_section(None), ""); - } - - #[test] - fn workspace_section_embeds_path_and_shared_directory_copy() { - let got = render_workspace_section(Some("/tmp/team-ws")); - assert!(got.contains("## Team Workspace")); - assert!(got.contains("`/tmp/team-ws`")); - assert!(got.contains("shared team workspace")); - } - - #[test] - fn preset_formatting_placeholders_are_empty_in_phase1() { - // Phase1 convention: presets are not surfaced, so both preset-formatting - // placeholders are replaced with "" regardless of other params. - // The regression test above (`no_unsubstituted_placeholders_when_all_sections_populated`) - // already asserts that both tokens are stripped from the final output. - // This test guards the behavior by simulating the full substitution on a - // template carrying just the two preset placeholders. - let template_with_presets_only = "step:${presetFormattingStepRule}|important:${presetFormattingImportantRule}"; - let out = template_with_presets_only - .replace("${presetFormattingStepRule}", "") - .replace("${presetFormattingImportantRule}", ""); - assert_eq!(out, "step:|important:"); - } - - #[test] - fn snapshot_minimal_params_with_stub_template_yields_empty_output() { - // While `LEAD_PROMPT_TEMPLATE` is the D5b-1 stub (`""`), the builder has no - // template to substitute into, so the output is empty regardless of params. - // Once D5b-1 lands, this test will start failing and should be updated to a - // real snapshot. The regression guard above keeps the substitution contract - // healthy in the meantime. - let renamed = HashMap::new(); - let out = build_lead_prompt(¶ms_min(&renamed)); - assert!(!out.is_empty(), "output should not be empty with real template"); - assert!(!out.contains("${"), "no unsubstituted placeholders"); - } - - #[test] - fn substitution_against_synthetic_template_matches_aionui_layout() { - // This synthetic template mirrors the shape of AionUi's leadPrompt.ts literal - // so we can validate end-to-end substitution without depending on D5b-1. - // When D5b-1 lands the real `LEAD_PROMPT_TEMPLATE` takes over; this test - // still exercises the same substitution code path. - const SYNTHETIC: &str = "## Your Teammates\n\ - ${teammateList}${availableTypesSection}${availableAssistantsSection}${workspaceSection}\n\ - STEP:${presetFormattingStepRule}END\n\ - - ${presetFormattingImportantRule}END"; - - let renamed = HashMap::new(); - let t = make_teammate("w1", "Worker1", "claude"); - let params = LeadPromptParams { - team_name: "Beta", - teammates: std::slice::from_ref(&t), - available_agent_types: &[AvailableAgentType { - agent_type: "claude".into(), - display_name: "general-purpose AI assistant".into(), - }], - available_assistants: &[], - renamed_agents: &renamed, - team_workspace: Some("/tmp/team-ws"), - }; - - let teammate_list = render_teammate_list(params.teammates, params.renamed_agents); - let types_section = render_available_types_section(params.available_agent_types); - let assistants_section = render_available_assistants_section(params.available_assistants); - let ws_section = render_workspace_section(params.team_workspace); - - let out = SYNTHETIC - .replace("${teammateList}", &teammate_list) - .replace("${availableTypesSection}", &types_section) - .replace("${availableAssistantsSection}", &assistants_section) - .replace("${workspaceSection}", &ws_section) - .replace("${presetFormattingStepRule}", "") - .replace("${presetFormattingImportantRule}", ""); - - assert!(!out.contains("${"), "unsubstituted placeholder:\n{out}"); - assert!(out.contains("## Your Teammates")); - assert!(out.contains("- Worker1 (claude, status: unknown)")); - assert!(out.contains("## Available Agent Types for Spawning")); - assert!(!out.contains("## Available Preset Assistants for Spawning")); - assert!(out.contains("## Team Workspace")); - assert!(out.contains("STEP:END")); - assert!(out.contains("- END")); - } -} diff --git a/crates/aionui-team/src/prompts/mod.rs b/crates/aionui-team/src/prompts/mod.rs index 824473a33..348c56ba0 100644 --- a/crates/aionui-team/src/prompts/mod.rs +++ b/crates/aionui-team/src/prompts/mod.rs @@ -1,75 +1,90 @@ -pub mod lead; pub mod team_guide; pub use team_guide::{TEAM_GUIDE_PROMPT_TEMPLATE, build_team_guide_prompt}; -pub mod teammate; use std::collections::HashMap; -use crate::prompts::lead::{AvailableAgentType, LeadPromptParams}; use crate::types::{MailboxMessage, MailboxMessageType, TaskStatus, TeamAgent, TeamTask}; +fn to_prompt_role(role: crate::types::TeammateRole) -> aionui_team_prompts::TeamPromptRole { + match role { + crate::types::TeammateRole::Lead => aionui_team_prompts::TeamPromptRole::Lead, + crate::types::TeammateRole::Teammate => aionui_team_prompts::TeamPromptRole::Teammate, + } +} + +fn to_prompt_agent(agent: &TeamAgent) -> aionui_team_prompts::TeamPromptAgent { + aionui_team_prompts::TeamPromptAgent { + slot_id: agent.slot_id.clone(), + name: agent.name.clone(), + role: to_prompt_role(agent.role), + backend: agent.backend.clone(), + model: agent.model.clone(), + status: agent.status.map(|status| status.to_string()), + } +} + /// Build the leader system prompt. /// -/// Delegates to [`lead::build_lead_prompt`], which mirrors the AionUi -/// `leadPrompt.ts` template verbatim. A one-line `Team: ""` header -/// is prepended so the leader knows which team it belongs to (AionUi -/// surfaces this through other channels, but the backend session has no -/// other place to inject it). +/// Delegates to `aionui-team-prompts`, the canonical Team role prompt crate. +/// A one-line `Team: ""` header is prepended so the leader knows which +/// team it belongs to. /// /// `available_agent_types` carries `(backend_id, display_name)` pairs that /// feed the `## Available Agent Types for Spawning` section; callers /// should source these from the team-capable backend whitelist. pub fn build_lead_prompt(team_name: &str, members: &[TeamAgent], available_agent_types: &[(String, String)]) -> String { - let agent_types: Vec = available_agent_types + let prompt_members: Vec<_> = members.iter().map(to_prompt_agent).collect(); + let agent_types: Vec<_> = available_agent_types .iter() - .map(|(backend, display)| AvailableAgentType { + .map(|(backend, display)| aionui_team_prompts::AvailableAgentType { agent_type: backend.clone(), display_name: display.clone(), }) .collect(); let renamed: HashMap = HashMap::new(); - let params = LeadPromptParams { + let body = aionui_team_prompts::build_lead_prompt(&aionui_team_prompts::LeadPromptParams { team_name, - teammates: members, + teammates: &prompt_members, available_agent_types: &agent_types, available_assistants: &[], renamed_agents: &renamed, team_workspace: None, - }; - - let body = lead::build_lead_prompt(¶ms); + }); format!("Team: \"{team_name}\"\n\n{body}") } -pub fn build_teammate_prompt(agent: &TeamAgent, team_name: &str) -> String { - let mut prompt = String::with_capacity(1024); - - prompt.push_str(&format!( - "You are **{}**, a Teammate Agent in team \"{}\". \ - Your slot ID is `{}`.\n\n", - agent.name, team_name, agent.slot_id, - )); +pub fn build_teammate_prompt(agent: &TeamAgent, team_name: &str, members: &[TeamAgent]) -> String { + let prompt_agent = to_prompt_agent(agent); + let prompt_members: Vec<_> = members.iter().map(to_prompt_agent).collect(); + let leader = prompt_members + .iter() + .find(|candidate| candidate.role == aionui_team_prompts::TeamPromptRole::Lead) + .cloned() + .unwrap_or_else(|| aionui_team_prompts::TeamPromptAgent { + slot_id: "lead".to_owned(), + name: "Lead".to_owned(), + role: aionui_team_prompts::TeamPromptRole::Lead, + backend: agent.backend.clone(), + model: agent.model.clone(), + status: None, + }); + let teammates: Vec<_> = prompt_members + .iter() + .filter(|candidate| candidate.slot_id != prompt_agent.slot_id) + .cloned() + .collect(); + let renamed = HashMap::new(); - prompt.push_str("## Your Role\n\n"); - prompt.push_str( - "You execute tasks assigned by the Lead Agent. Focus on completing your \ - assigned work thoroughly and reporting back.\n\n", - ); - - prompt.push_str("## Communication Protocol\n\n"); - prompt.push_str( - "- Use `team_send_message` to report progress or ask questions to the Lead.\n\ - - Use `team_task_update` to update task status as you work \ - (pending → in_progress → completed).\n\ - - When your assigned work is done, send an idle notification. \ - The system will notify the Lead.\n\ - - If you receive a `shutdown_request`, finish any critical work, \ - then respond with \"shutdown_approved\" or \"shutdown_rejected: \".\n", - ); - - prompt + aionui_team_prompts::build_teammate_prompt(&aionui_team_prompts::TeammatePromptParams { + agent: &prompt_agent, + team_name, + leader: &leader, + teammates: &teammates, + renamed_agents: &renamed, + team_workspace: None, + }) } pub fn build_wake_payload(agent: &TeamAgent, tasks: &[TeamTask], unread_messages: &[MailboxMessage]) -> String { @@ -303,30 +318,36 @@ mod tests { #[test] fn teammate_prompt_contains_agent_identity() { let agent = make_teammate("w1", "Worker1"); - let prompt = build_teammate_prompt(&agent, "Alpha"); + let members = vec![make_lead(), agent.clone()]; + let prompt = build_teammate_prompt(&agent, "Alpha", &members); - assert!(prompt.contains("**Worker1**")); - assert!(prompt.contains("\"Alpha\"")); - assert!(prompt.contains("`w1`")); + assert!(prompt.contains("## Team Governance")); + assert!(prompt.contains("Name: Worker1")); + assert!(prompt.contains("Team: Alpha")); + assert!(prompt.contains("Leader: Lead")); } #[test] fn teammate_prompt_contains_communication_protocol() { let agent = make_teammate("w1", "Worker1"); - let prompt = build_teammate_prompt(&agent, "Alpha"); + let members = vec![make_lead(), agent.clone()]; + let prompt = build_teammate_prompt(&agent, "Alpha", &members); + assert!(prompt.contains("## Team Coordination Tools")); + assert!(prompt.contains("You MUST use the `team_*` MCP tools for ALL team coordination.")); assert!(prompt.contains("team_send_message")); assert!(prompt.contains("team_task_update")); - assert!(prompt.contains("idle notification")); assert!(prompt.contains("shutdown_request")); assert!(prompt.contains("shutdown_approved")); + assert!(prompt.contains("STOP GENERATING")); } #[test] fn teammate_prompt_contains_team_name() { let agent = make_teammate("w1", "W"); - let prompt = build_teammate_prompt(&agent, "Beta Team"); - assert!(prompt.contains("\"Beta Team\"")); + let members = vec![make_lead(), agent.clone()]; + let prompt = build_teammate_prompt(&agent, "Beta Team", &members); + assert!(prompt.contains("Team: Beta Team")); } // -- Wake payload --------------------------------------------------------- diff --git a/crates/aionui-team/src/prompts/team_guide.rs b/crates/aionui-team/src/prompts/team_guide.rs index 47adfa22f..6c8cded59 100644 --- a/crates/aionui-team/src/prompts/team_guide.rs +++ b/crates/aionui-team/src/prompts/team_guide.rs @@ -1,168 +1,30 @@ -//! Team Guide Prompt (Layer 1) — injected into solo ACP agents so they know -//! when/how to propose a multi-agent Team. Reproduces AionUi -//! `src/process/team/prompts/teamGuidePrompt.ts` byte-for-byte. +//! Solo Team Guide prompt wrapper. //! -//! Hard constraint (aionui-audit §8 #5, team-prompts.md §5): the template text -//! is treated as raw material — it must not be translated, rewritten, or -//! reordered. Only the interpolated slots (`backend`, `leader_label`) may vary. +//! The canonical template lives in `aionui-team-prompts` so ACP, Aionrs, and +//! Team-side prompt tests share one source of truth. -const EXPLICIT_TEAM_REQUEST_CRITERIA: &str = "\ -- The user explicitly asks to create a Team -- The user explicitly asks for multiple agents, teammates, or parallel workers -- The user says they want to pull in a Team before starting"; +pub const TEAM_GUIDE_PROMPT_TEMPLATE: &str = aionui_team_prompts::guide::TEAM_GUIDE_PROMPT_TEMPLATE; -const EXTREME_COMPLEXITY_CRITERIA: &str = "\ -- The task is so large, risky, or specialized that one agent is unlikely to complete it well alone -- The work needs substantial parallel role separation that cannot be reasonably handled in a normal solo workflow -- This bar is very high: if you can handle the task yourself, stay solo"; - -const STAY_SOLO_CRITERIA: &str = "\ -- Greetings, casual conversation, or general questions -- Single-point tasks: one question, one file, one fix, one translation, one explanation -- Normal coding, writing, research, or analysis tasks that one agent can handle with some effort -- Any task you can reasonably complete yourself, even if it takes multiple turns"; - -const SOLO_DEFAULT_RULE: &str = "Handle the task yourself in the current chat by default. Do NOT proactively recommend Team just because the work spans multiple files, takes multiple rounds, or would benefit from specialization."; - -/// Full Team Guide prompt template with `{leader_cell}` / `{agent_type}` -/// placeholders. Exported for cross-crate snapshot tests and the Wave 5 -/// capability injector; prefer [`build_team_guide_prompt`] for runtime use. -pub const TEAM_GUIDE_PROMPT_TEMPLATE: &str = "## Team Mode - -You can create a multi-agent Team for the user. - -### Default behavior -{solo_default_rule} - -### Only bring up Team in either of these cases -1. The user explicitly wants a Team or multiple agents: -{explicit_team_request_criteria} -2. The task is exceptionally complex and you genuinely believe one agent is unlikely to handle it well alone: -{extreme_complexity_criteria} - -### Otherwise stay solo and do not mention Team -{stay_solo_criteria} - -If case 2 applies, ask at most once whether the user wants to bring in a Team. Keep it brief and optional. If the user says no, ignores it, or prefers solo help, continue solo and do not mention Team again. - -### How to proceed when Team is requested or approved (STRICT — follow every step, do NOT skip) -1. FIRST call `aion_list_models` to check available models for each agent type you plan to use. -2. Explain in one sentence why the Team setup helps this task. -3. Present a team configuration table: role name, responsibility, agent type, and recommended model (from aion_list_models results) for each member. Example format: - | Role | Responsibility | Type | Model | - | Leader | Coordinate and review | {leader_cell} | (default) | - | Developer | Implement features | {agent_type} | (model from list) | - | Tester | Write and run tests | {agent_type} | (model from list) | -4. **Output the table as a normal text message and END YOUR TURN.** Do NOT call `aion_create_team` or any other tool (including ask_user) in this turn. Wait for the user to reply in their next message with explicit confirmation (e.g. \"ok\", \"go ahead\", \"确认\") before proceeding. -5. After user confirms → call `aion_create_team`. The summary MUST include both the goal and the confirmed team configuration. (The system automatically sets the correct agent type — you do NOT need to pass agentType.) -6. After `aion_create_team` returns → the Team has been created and the current conversation has been bound as Leader. **Do NOT call `team_spawn_agent`, `team_send_message`, or any other `team_*` tool in this solo turn.** Output only one brief user-facing handoff in the user's language. It should mean: the Team is ready, send the next message, and I will continue from there. Then END YOUR TURN. Do not mention the Team page, solo turn, `team_*` tools, `TeamRun`, or internal tool state in the user-facing handoff. -7. User declines or wants changes → adjust or proceed solo. Do not mention Team again unless the user asks. - -### Tool constraint -Before team creation: use **only** `aion_create_team` and `aion_list_models`. After `aion_create_team` succeeds: do not call any `team_*` tools in this solo turn. Team tools are only for normal Team runtime after the Team page accepts the user's first Team message and an active `TeamRun` exists."; - -/// Build the Team Guide prompt for a solo agent. -/// -/// * `backend` — agent backend key (`"claude"`, `"gemini"`, `"codex"`, …). Empty -/// string falls back to `"claude"`, matching AionUi `opts.backend || 'claude'`. -/// * `leader_label` — optional display name for a preset assistant (e.g. -/// `"Word Creator"`). When present it renders as `"{label} ({backend})"`, -/// mirroring the `rawLabel ? "${rawLabel} (${agentType})" : agentType` branch -/// in `teamGuidePrompt.ts`. Whitespace-only labels are treated as absent. pub fn build_team_guide_prompt(backend: &str, leader_label: Option<&str>) -> String { - let agent_type = if backend.is_empty() { "claude" } else { backend }; - let raw_label = leader_label.map(str::trim).filter(|s| !s.is_empty()); - let leader_cell = match raw_label { - Some(label) => format!("{label} ({agent_type})"), - None => agent_type.to_owned(), - }; - - TEAM_GUIDE_PROMPT_TEMPLATE - .replace("{solo_default_rule}", SOLO_DEFAULT_RULE) - .replace("{explicit_team_request_criteria}", EXPLICIT_TEAM_REQUEST_CRITERIA) - .replace("{extreme_complexity_criteria}", EXTREME_COMPLEXITY_CRITERIA) - .replace("{stay_solo_criteria}", STAY_SOLO_CRITERIA) - .replace("{leader_cell}", &leader_cell) - .replace("{agent_type}", agent_type) + aionui_team_prompts::guide::build_solo_team_guide_prompt_with_label(backend, leader_label) } #[cfg(test)] mod tests { use super::*; - #[test] - fn team_guide_prompt_plain_backend_matches_snapshot() { - let prompt = build_team_guide_prompt("claude", None); - let expected = "## Team Mode\n\ -\n\ -You can create a multi-agent Team for the user.\n\ -\n\ -### Default behavior\n\ -Handle the task yourself in the current chat by default. Do NOT proactively recommend Team just because the work spans multiple files, takes multiple rounds, or would benefit from specialization.\n\ -\n\ -### Only bring up Team in either of these cases\n\ -1. The user explicitly wants a Team or multiple agents:\n\ -- The user explicitly asks to create a Team\n\ -- The user explicitly asks for multiple agents, teammates, or parallel workers\n\ -- The user says they want to pull in a Team before starting\n\ -2. The task is exceptionally complex and you genuinely believe one agent is unlikely to handle it well alone:\n\ -- The task is so large, risky, or specialized that one agent is unlikely to complete it well alone\n\ -- The work needs substantial parallel role separation that cannot be reasonably handled in a normal solo workflow\n\ -- This bar is very high: if you can handle the task yourself, stay solo\n\ -\n\ -### Otherwise stay solo and do not mention Team\n\ -- Greetings, casual conversation, or general questions\n\ -- Single-point tasks: one question, one file, one fix, one translation, one explanation\n\ -- Normal coding, writing, research, or analysis tasks that one agent can handle with some effort\n\ -- Any task you can reasonably complete yourself, even if it takes multiple turns\n\ -\n\ -If case 2 applies, ask at most once whether the user wants to bring in a Team. Keep it brief and optional. If the user says no, ignores it, or prefers solo help, continue solo and do not mention Team again.\n\ -\n\ -### How to proceed when Team is requested or approved (STRICT — follow every step, do NOT skip)\n\ -1. FIRST call `aion_list_models` to check available models for each agent type you plan to use.\n\ -2. Explain in one sentence why the Team setup helps this task.\n\ -3. Present a team configuration table: role name, responsibility, agent type, and recommended model (from aion_list_models results) for each member. Example format:\n \ -| Role | Responsibility | Type | Model |\n \ -| Leader | Coordinate and review | claude | (default) |\n \ -| Developer | Implement features | claude | (model from list) |\n \ -| Tester | Write and run tests | claude | (model from list) |\n\ -4. **Output the table as a normal text message and END YOUR TURN.** Do NOT call `aion_create_team` or any other tool (including ask_user) in this turn. Wait for the user to reply in their next message with explicit confirmation (e.g. \"ok\", \"go ahead\", \"确认\") before proceeding.\n\ -5. After user confirms → call `aion_create_team`. The summary MUST include both the goal and the confirmed team configuration. (The system automatically sets the correct agent type — you do NOT need to pass agentType.)\n\ -6. After `aion_create_team` returns → the Team has been created and the current conversation has been bound as Leader. **Do NOT call `team_spawn_agent`, `team_send_message`, or any other `team_*` tool in this solo turn.** Output only one brief user-facing handoff in the user's language. It should mean: the Team is ready, send the next message, and I will continue from there. Then END YOUR TURN. Do not mention the Team page, solo turn, `team_*` tools, `TeamRun`, or internal tool state in the user-facing handoff.\n\ -7. User declines or wants changes → adjust or proceed solo. Do not mention Team again unless the user asks.\n\ -\n\ -### Tool constraint\n\ -Before team creation: use **only** `aion_create_team` and `aion_list_models`. After `aion_create_team` succeeds: do not call any `team_*` tools in this solo turn. Team tools are only for normal Team runtime after the Team page accepts the user's first Team message and an active `TeamRun` exists."; - assert_eq!(prompt, expected); - } - #[test] fn team_guide_prompt_hands_off_after_create_team() { let prompt = build_team_guide_prompt("claude", None); - assert!(prompt.contains( - "After `aion_create_team` returns → the Team has been created and the current conversation has been bound as Leader." - )); - assert!(prompt.contains( - "Do NOT call `team_spawn_agent`, `team_send_message`, or any other `team_*` tool in this solo turn." - )); - assert!(prompt.contains( - "Output only one brief user-facing handoff in the user's language. It should mean: the Team is ready, send the next message, and I will continue from there." - )); - assert!(prompt.contains( - "Do not mention the Team page, solo turn, `team_*` tools, `TeamRun`, or internal tool state in the user-facing handoff." + assert!(prompt.contains("aion_create_team")); + assert!(prompt.contains("aion_list_models")); + assert!(prompt.contains("hand off to the created Team conversation")); + assert!(prompt.contains("Do NOT call `team_*` tools from this solo Guide MCP session.")); + assert!(!prompt.contains("Immediately")); + assert!(!prompt.contains( + "use team tools (`team_spawn_agent`, `team_send_message`, `team_members`, `team_task_create`, etc.) to manage your team" )); - assert!( - prompt.contains("After `aion_create_team` succeeds: do not call any `team_*` tools in this solo turn.") - ); - assert!( - !prompt.contains("Your team tools (team_spawn_agent, team_send_message, etc.) are now active."), - "prompt must not claim Team tools are active immediately after creation" - ); - assert!( - !prompt.contains("Immediately proceed to spawn teammates as planned"), - "prompt must not ask the solo agent to spawn teammates in the same solo turn" - ); } #[test] @@ -185,5 +47,6 @@ Before team creation: use **only** `aion_create_team` and `aion_list_models`. Af fn team_guide_prompt_whitespace_label_treated_as_absent() { let prompt = build_team_guide_prompt("codex", Some(" ")); assert!(prompt.contains("| Leader | Coordinate and review | codex | (default) |")); + assert!(!prompt.contains("()")); } } diff --git a/crates/aionui-team/src/prompts/teammate.rs b/crates/aionui-team/src/prompts/teammate.rs deleted file mode 100644 index 7fcdb319c..000000000 --- a/crates/aionui-team/src/prompts/teammate.rs +++ /dev/null @@ -1,426 +0,0 @@ -//! Teammate prompt template + wake payload builder. -//! -//! Template text copied verbatim from AionUi `src/process/team/prompts/teammatePrompt.ts` -//! (aionui-audit §8 #5: prompt text must be reused as-is, no translation, no rewriting). -//! -//! Placeholders enclosed in `{{...}}` are filled by `build_teammate_prompt`: -//! - `{{AGENT_NAME}}`, `{{ROLE_DESC}}`, `{{LEADER_NAME}}`, `{{TEAMMATES}}`, `{{WORKSPACE}}` - -use std::collections::HashMap; - -use crate::types::{MailboxMessage, MailboxMessageType, TaskStatus, TeamAgent, TeamTask, TeammateRole}; - -// --------------------------------------------------------------------------- -// TeammatePromptParams — signature frozen by phase1/interface-contracts.md §5 -// --------------------------------------------------------------------------- - -pub struct TeammatePromptParams<'a> { - pub agent: &'a TeamAgent, - pub team_name: &'a str, - pub leader: &'a TeamAgent, - pub teammates: &'a [TeamAgent], - pub renamed_agents: &'a HashMap, - pub team_workspace: Option<&'a str>, -} - -// --------------------------------------------------------------------------- -// Template constant (verbatim from teammatePrompt.ts) -// --------------------------------------------------------------------------- - -/// Full Teammate system prompt template. -/// -/// Tokens `{{...}}` are substituted by [`build_teammate_prompt`]. The rest of -/// the text MUST NOT be modified (aionui-audit §8 #5). -pub const TEAMMATE_PROMPT_TEMPLATE: &str = r#"# You are a Team Member - -## Your Identity -Name: {{AGENT_NAME}}, Role: {{ROLE_DESC}} - -## 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 -- Briefly introduce yourself and your role on the team, then invite the user to share what they need -- Do NOT open with task board details, idle/waiting status, or coordination mechanics unless they are directly relevant - -## Your Team -Leader: {{LEADER_NAME}} -Teammates: {{TEAMMATES}}{{WORKSPACE}} - -## 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, -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. - -## How to Work -1. Read your unread messages to understand your assignment -2. If you have a clear task assignment in the messages AND no prerequisite is blocking it, start working on it immediately -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 - -## 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. - -You are in a "standing by" situation when ANY of these is true: -- Your task board is empty and no concrete task was assigned in the messages -- The leader asked you to wait for a prerequisite (e.g. "hold until reviewer-1 finishes") -- 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"` -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. - -## Bug Fix Priority -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. - -## Shutdown Requests -If you receive a message with type `shutdown_request`, the leader is asking you to shut down. -- To agree: use `team_send_message` to send exactly `shutdown_approved` to the leader. -- To refuse: use `team_send_message` to send `shutdown_rejected: ` to the leader. - -## Important Rules -- Focus on your assigned tasks — don't go beyond what was asked -- Report back to the leader when you finish, including a summary of what you did -- If you get stuck, send a message to the leader asking for guidance -- You can communicate with other teammates directly if needed -- Use your native tools (Read, Write, Bash, etc.) for implementation work"#; - -// --------------------------------------------------------------------------- -// Role description (mirror of teammatePrompt.ts `roleDescription`) -// --------------------------------------------------------------------------- - -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"), - } -} - -// --------------------------------------------------------------------------- -// Builder -// --------------------------------------------------------------------------- - -/// Build the full Teammate system prompt by filling [`TEAMMATE_PROMPT_TEMPLATE`]. -pub fn build_teammate_prompt(params: &TeammatePromptParams<'_>) -> String { - let teammates_section = if params.teammates.is_empty() { - "(none)".to_string() - } else { - params - .teammates - .iter() - .map(|t| match params.renamed_agents.get(&t.slot_id) { - Some(original) => format!("{} [formerly: {}]", t.name, original), - None => t.name.clone(), - }) - .collect::>() - .join(", ") - }; - - let workspace_section = match params.team_workspace { - Some(ws) => format!( - "\n\n## Workspaces\n\ -- **Team workspace**: `{ws}` — all project work (code, files, tests) happens here.\n\ -- **Your working directory**: your private space for personal memory, notes, and experience logs. Not for project files.\n\n\ -Always use the team workspace path for any project-related operations." - ), - None => String::new(), - }; - - TEAMMATE_PROMPT_TEMPLATE - .replace("{{AGENT_NAME}}", ¶ms.agent.name) - .replace("{{ROLE_DESC}}", &role_description(¶ms.agent.backend)) - .replace("{{LEADER_NAME}}", ¶ms.leader.name) - .replace("{{TEAMMATES}}", &teammates_section) - .replace("{{WORKSPACE}}", &workspace_section) -} - -// --------------------------------------------------------------------------- -// Wake payload — frozen signature from phase1/interface-contracts.md §5 -// --------------------------------------------------------------------------- - -/// Build the payload sent as the first `send_message` content when waking an -/// agent. Combines mailbox messages and current task board into markdown. -/// -/// AionUi's `TeammateManager.wake` uses `formatMessages(...)` alone for the -/// subsequent-wake case. The phase1 contract extends this to also surface the -/// task board so the agent can see outstanding work without calling -/// `team_task_list` first. -pub fn build_wake_payload( - agent: &TeamAgent, - tasks: &[TeamTask], - unread_messages: &[MailboxMessage], - sender_name_lookup: &HashMap, -) -> String { - let mut out = String::with_capacity(1024); - - // -- Unread Messages ------------------------------------------------------ - out.push_str("## Unread Messages\n"); - if unread_messages.is_empty() { - out.push_str("No unread messages.\n"); - } else { - out.push_str(&format_messages(unread_messages, sender_name_lookup)); - out.push('\n'); - } - - // -- Task Board ----------------------------------------------------------- - out.push('\n'); - out.push_str("## Task Board\n"); - if tasks.is_empty() { - out.push_str("No tasks on the board.\n"); - } else { - out.push_str("| ID | Subject | Status | Owner | Blocked By |\n"); - out.push_str("|---|---|---|---|---|\n"); - for t in tasks { - let short_id = if t.id.len() > 8 { &t.id[..8] } else { &t.id }; - let owner = t.owner.as_deref().unwrap_or("-"); - let blocked = if t.blocked_by.is_empty() { - "-".to_string() - } else { - t.blocked_by.join(", ") - }; - out.push_str(&format!( - "| {short_id}… | {} | {} | {} | {} |\n", - t.subject, - task_status_label(t.status), - owner, - blocked, - )); - } - } - - // -- Identity footer ------------------------------------------------------ - out.push('\n'); - out.push_str(&format!( - "You are **{}** (role: {}). Proceed with your work.\n", - agent.name, - match agent.role { - TeammateRole::Lead => "lead", - TeammateRole::Teammate => "teammate", - }, - )); - - out -} - -/// Mirror of AionUi `formatHelpers.ts :: formatMessages`. -/// Sender "user" renders as `[From User]`; known slot_id resolves to agent -/// name via `sender_name_lookup`; unknown slot_id falls back to the raw id. -fn format_messages(messages: &[MailboxMessage], sender_name_lookup: &HashMap) -> String { - messages - .iter() - .map(|m| { - let sender = if m.from_agent_id == "user" { - "User".to_string() - } else { - sender_name_lookup - .get(&m.from_agent_id) - .cloned() - .unwrap_or_else(|| m.from_agent_id.clone()) - }; - let type_tag = match m.msg_type { - MailboxMessageType::Message => "", - MailboxMessageType::IdleNotification => " [idle_notification]", - MailboxMessageType::ShutdownRequest => " [shutdown_request]", - }; - let summary_line = m - .summary - .as_deref() - .map(|s| format!("\nSummary: {s}")) - .unwrap_or_default(); - format!("[From {sender}{type_tag}] {content}{summary_line}", content = m.content,) - }) - .collect::>() - .join("\n") -} - -fn task_status_label(status: TaskStatus) -> &'static str { - match status { - TaskStatus::Pending => "pending", - TaskStatus::InProgress => "in_progress", - TaskStatus::Completed => "completed", - TaskStatus::Deleted => "deleted", - } -} - -// --------------------------------------------------------------------------- -// Tests — snapshot-style (string-contains) per D5c task spec -// --------------------------------------------------------------------------- - -#[cfg(test)] -mod tests { - use super::*; - use crate::types::{MailboxMessageType, TaskStatus, TeammateRole}; - - fn make_agent(slot_id: &str, name: &str, role: TeammateRole, backend: &str) -> TeamAgent { - TeamAgent { - slot_id: slot_id.into(), - name: name.into(), - role, - conversation_id: format!("conv-{slot_id}"), - backend: backend.into(), - model: "default".into(), - custom_agent_id: None, - status: None, - conversation_type: None, - cli_path: None, - } - } - - fn make_task(id: &str, subject: &str, status: TaskStatus, owner: Option<&str>) -> TeamTask { - TeamTask { - id: id.into(), - team_id: "t1".into(), - subject: subject.into(), - description: None, - status, - owner: owner.map(String::from), - blocked_by: vec![], - blocks: vec![], - metadata: None, - created_at: 0, - updated_at: 0, - } - } - - fn make_msg(id: &str, from: &str, to: &str, msg_type: MailboxMessageType, content: &str) -> MailboxMessage { - MailboxMessage { - id: id.into(), - team_id: "t1".into(), - to_agent_id: to.into(), - from_agent_id: from.into(), - msg_type, - content: content.into(), - summary: None, - files: None, - read: false, - created_at: 0, - } - } - - // Test 1: teammate prompt with minimal params (no renamed, no workspace) - #[test] - fn teammate_prompt_minimal_params() { - let lead = make_agent("lead-1", "Captain", TeammateRole::Lead, "claude"); - let agent = make_agent("w1", "Worker1", TeammateRole::Teammate, "gemini"); - let renamed = HashMap::new(); - let params = TeammatePromptParams { - agent: &agent, - team_name: "Alpha", - leader: &lead, - teammates: &[], - renamed_agents: &renamed, - team_workspace: None, - }; - let out = build_teammate_prompt(¶ms); - - assert!(out.contains("# You are a Team Member")); - assert!(out.contains("Name: Worker1, Role: Google Gemini AI assistant")); - assert!(out.contains("Leader: Captain")); - assert!(out.contains("Teammates: (none)")); - assert!(!out.contains("## Workspaces")); - assert!(out.contains("## Standing By (CRITICAL")); - assert!(out.contains("shutdown_approved")); - assert!(out.contains("shutdown_rejected:")); - } - - // Test 2: teammate prompt with renamed teammates and workspace - #[test] - fn teammate_prompt_with_renamed_and_workspace() { - let lead = make_agent("lead-1", "Captain", TeammateRole::Lead, "claude"); - let agent = make_agent("w1", "Worker1", TeammateRole::Teammate, "claude"); - let mate_a = make_agent("w2", "Alice", TeammateRole::Teammate, "claude"); - let mate_b = make_agent("w3", "Bob", TeammateRole::Teammate, "codex"); - let mut renamed = HashMap::new(); - renamed.insert("w3".into(), "Robert".into()); - let params = TeammatePromptParams { - agent: &agent, - team_name: "Alpha", - leader: &lead, - teammates: &[mate_a, mate_b], - renamed_agents: &renamed, - team_workspace: Some("/workspace/team-alpha"), - }; - let out = build_teammate_prompt(¶ms); - - assert!(out.contains("Teammates: Alice, Bob [formerly: Robert]")); - assert!(out.contains("## Workspaces")); - assert!(out.contains("`/workspace/team-alpha`")); - assert!(out.contains("Role: general-purpose AI assistant")); - } - - // Test 3: wake payload with empty mailbox and no tasks - #[test] - fn wake_payload_empty_mailbox_no_tasks() { - let agent = make_agent("lead-1", "Captain", TeammateRole::Lead, "claude"); - let lookup = HashMap::new(); - let out = build_wake_payload(&agent, &[], &[], &lookup); - - assert!(out.contains("## Unread Messages")); - assert!(out.contains("No unread messages.")); - assert!(out.contains("## Task Board")); - assert!(out.contains("No tasks on the board.")); - assert!(out.contains("You are **Captain** (role: lead)")); - } - - // Test 4: wake payload with tasks and messages (mixed types) - #[test] - fn wake_payload_with_tasks_and_messages() { - let agent = make_agent("w1", "Worker1", TeammateRole::Teammate, "claude"); - let mut lookup = HashMap::new(); - lookup.insert("lead-1".into(), "Captain".into()); - lookup.insert("w2".into(), "Alice".into()); - - let msgs = vec![ - make_msg( - "m1", - "lead-1", - "w1", - MailboxMessageType::Message, - "Please implement feature X", - ), - make_msg("m2", "user", "w1", MailboxMessageType::Message, "Direct user request"), - make_msg("m3", "w2", "w1", MailboxMessageType::IdleNotification, "done with task"), - ]; - - let tasks = vec![ - make_task( - "aaaaaaaa-1234-5678-9abc-def012345678", - "Implement X", - TaskStatus::InProgress, - Some("w1"), - ), - make_task( - "bbbbbbbb-1234-5678-9abc-def012345678", - "Review Y", - TaskStatus::Pending, - None, - ), - ]; - - let out = build_wake_payload(&agent, &tasks, &msgs, &lookup); - - // Messages section - assert!(out.contains("[From Captain] Please implement feature X")); - assert!(out.contains("[From User] Direct user request")); - assert!(out.contains("[From Alice [idle_notification]] done with task")); - - // Task board section - assert!(out.contains("aaaaaaaa…")); - assert!(out.contains("Implement X")); - assert!(out.contains("in_progress")); - assert!(out.contains("Review Y")); - assert!(out.contains("pending")); - - // Identity footer - assert!(out.contains("You are **Worker1** (role: teammate)")); - } -} diff --git a/crates/aionui-team/src/service.rs b/crates/aionui-team/src/service.rs index 2ad61abe6..9c968aecd 100644 --- a/crates/aionui-team/src/service.rs +++ b/crates/aionui-team/src/service.rs @@ -4,12 +4,12 @@ pub(crate) mod spawn_support; use std::path::PathBuf; use std::sync::{Arc, Weak}; -use aionui_ai_agent::{AgentError, AgentInstance, IWorkerTaskManager}; +use aionui_ai_agent::{AgentError, AgentInstance, AgentStreamEvent, IWorkerTaskManager}; use aionui_api_types::{ AddAgentRequest, CreateTeamRequest, GuideMcpConfig, TeamAgentResponse, TeamMcpPhase, TeamMcpStatusPayload, TeamResponse, TeamRunAckResponse, TeamRunTargetRole, WebSocketMessage, }; -use aionui_common::{AgentKillReason, generate_id, now_ms}; +use aionui_common::{AgentKillReason, ConversationStatus, generate_id, now_ms}; use aionui_db::models::TeamRow; use aionui_db::{IAgentMetadataRepository, IProviderRepository, ITeamRepository, UpdateTeamParams}; use aionui_realtime::EventBroadcaster; @@ -189,6 +189,7 @@ impl TeamSessionService { return Err(TeamError::InvalidRequest("at least one agent is required".into())); } + let adopted_leader_conversation_id = req.agents.first().and_then(|agent| agent.conversation_id.clone()); let shared_workspace = match req.workspace.as_deref() { Some(workspace) if !workspace.is_empty() => Some(validate_create_workspace_path(workspace)?), _ => None, @@ -249,6 +250,14 @@ impl TeamSessionService { // via POST /api/teams/{id}/session if needed. if let Err(e) = self.ensure_session_inner(&team.id, true).await { warn!(team_id = %team.id, error = %e, "auto ensure_session after create_team failed"); + } else if let Some(conversation_id) = adopted_leader_conversation_id + && let Some(leader) = team + .agents + .iter() + .find(|agent| agent.role == TeammateRole::Lead && agent.conversation_id == conversation_id) + .cloned() + { + self.schedule_deferred_leader_rebuild(user_id.to_owned(), team.id.clone(), leader); } self.build_team_response(&team).await @@ -686,6 +695,77 @@ impl TeamSessionService { Ok(()) } + fn schedule_deferred_leader_rebuild(&self, user_id: String, team_id: String, leader: TeamAgent) { + info!( + team_id = %team_id, + slot_id = %leader.slot_id, + conversation_id = %leader.conversation_id, + "deferred leader Team MCP rebuild scheduled" + ); + let service = self.self_ref.clone(); + tokio::spawn(async move { + let Some(service) = service.upgrade() else { + return; + }; + service.wait_until_agent_not_running(&leader.conversation_id).await; + if let Err(error) = service.rebuild_single_agent_process(&user_id, &team_id, &leader).await { + warn!( + team_id = %team_id, + slot_id = %leader.slot_id, + conversation_id = %leader.conversation_id, + error = %error, + "deferred leader Team MCP rebuild failed" + ); + } + }); + } + + async fn wait_until_agent_not_running(&self, conversation_id: &str) { + let Some(agent) = self.task_manager.get_task(conversation_id) else { + return; + }; + if agent.status() != Some(ConversationStatus::Running) { + return; + } + let mut events = agent.subscribe(); + loop { + match events.recv().await { + Ok(AgentStreamEvent::Finish(_) | AgentStreamEvent::Error(_)) => return, + Ok(_) => {} + Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => { + if agent.status() != Some(ConversationStatus::Running) { + return; + } + } + Err(tokio::sync::broadcast::error::RecvError::Closed) => return, + } + } + } + + async fn rebuild_single_agent_process( + &self, + user_id: &str, + team_id: &str, + agent: &TeamAgent, + ) -> Result<(), TeamError> { + let session = self + .sessions + .get(team_id) + .map(|entry| Arc::clone(&entry.session)) + .ok_or_else(|| TeamError::InvalidRequest(format!("no active session for team {team_id}")))?; + let cfg = session.mcp_stdio_config(&agent.slot_id); + self.provisioner() + .attach_agent_process(user_id, agent, cfg, &self.task_manager) + .await?; + info!( + team_id = %team_id, + slot_id = %agent.slot_id, + conversation_id = %agent.conversation_id, + "deferred leader Team MCP rebuild completed" + ); + Ok(()) + } + /// 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]) { diff --git a/crates/aionui-team/src/session.rs b/crates/aionui-team/src/session.rs index cf7a6a74c..45b64e83d 100644 --- a/crates/aionui-team/src/session.rs +++ b/crates/aionui-team/src/session.rs @@ -305,7 +305,10 @@ impl TeamSession { &available_agent_types, ) } - TeammateRole::Teammate => build_teammate_prompt(&agent, &self.team.name), + 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 { @@ -2989,7 +2992,12 @@ mod tests { .expect("WakeInput"); assert!( - input.first_message.contains("Teammate Agent"), + 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 ); @@ -2997,6 +3005,34 @@ mod tests { 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; diff --git a/crates/aionui-team/tests/mcp_server_integration.rs b/crates/aionui-team/tests/mcp_server_integration.rs index 07c26b413..ece1a5fc2 100644 --- a/crates/aionui-team/tests/mcp_server_integration.rs +++ b/crates/aionui-team/tests/mcp_server_integration.rs @@ -143,6 +143,31 @@ async fn read_response(stream: &mut TcpStream) -> Value { serde_json::from_slice(&frame).unwrap() } +async fn http_rpc(port: u16, slot_id: &str, payload: Value) -> Value { + http_rpc_with_auth(port, slot_id, Some("test-token-123"), payload).await +} + +async fn http_rpc_with_auth(port: u16, slot_id: &str, token: Option<&str>, payload: Value) -> Value { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let body = serde_json::to_string(&payload).unwrap(); + let auth_header = token + .map(|token| format!("Authorization: Bearer {token}\r\n")) + .unwrap_or_default(); + let request = format!( + "POST /mcp HTTP/1.1\r\nHost: 127.0.0.1:{port}\r\nContent-Type: application/json\r\n{auth_header}x-slot-id: {slot_id}\r\nContent-Length: {}\r\n\r\n{body}", + body.len() + ); + let mut stream = TcpStream::connect(format!("127.0.0.1:{port}")).await.unwrap(); + stream.write_all(request.as_bytes()).await.unwrap(); + + let mut buf = Vec::new(); + stream.read_to_end(&mut buf).await.unwrap(); + let response = String::from_utf8_lossy(&buf); + let body = response.split("\r\n\r\n").nth(1).unwrap_or(""); + serde_json::from_str(body).unwrap() +} + async fn call_tool(stream: &mut TcpStream, id: u64, tool: &str, args: Value) -> Value { let req = json!({ "jsonrpc": "2.0", @@ -157,6 +182,22 @@ async fn call_tool(stream: &mut TcpStream, id: u64, tool: &str, args: Value) -> read_response(stream).await } +async fn list_tools(stream: &mut TcpStream, id: u64) -> Vec { + let req = json!({ + "jsonrpc": "2.0", + "id": id, + "method": "tools/list" + }); + send_request(stream, &req).await; + let resp = read_response(stream).await; + resp["result"]["tools"] + .as_array() + .unwrap() + .iter() + .map(|tool| tool["name"].as_str().unwrap().to_owned()) + .collect() +} + fn extract_text(resp: &Value) -> String { resp["result"]["content"][0]["text"].as_str().unwrap_or("").to_string() } @@ -246,27 +287,40 @@ 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 req = json!({ - "jsonrpc": "2.0", - "id": 10, - "method": "tools/list" - }); - 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(), 10); + let names = list_tools(&mut stream, 10).await; + assert_eq!(names.len(), 10); + + assert!(names.contains(&"team_send_message".to_owned())); + assert!(names.contains(&"team_spawn_agent".to_owned())); + assert!(names.contains(&"team_task_create".to_owned())); + assert!(names.contains(&"team_task_update".to_owned())); + assert!(names.contains(&"team_task_list".to_owned())); + assert!(names.contains(&"team_members".to_owned())); + assert!(names.contains(&"team_rename_agent".to_owned())); + assert!(names.contains(&"team_shutdown_agent".to_owned())); + assert!(names.contains(&"team_list_models".to_owned())); + assert!(names.contains(&"team_describe_assistant".to_owned())); + + env.server.stop(); +} + +#[tokio::test] +async fn mcp_tools_list_filters_lead_only_tools() { + let env = setup().await; + let mut stream = connect_and_init(env.server.port(), "test-token-123", "worker-1").await; + + let names = list_tools(&mut stream, 10).await; - let names: Vec<&str> = tools.iter().map(|t| t["name"].as_str().unwrap()).collect(); - assert!(names.contains(&"team_send_message")); - assert!(names.contains(&"team_spawn_agent")); - assert!(names.contains(&"team_task_create")); - assert!(names.contains(&"team_task_update")); - assert!(names.contains(&"team_task_list")); - assert!(names.contains(&"team_members")); - assert!(names.contains(&"team_rename_agent")); - assert!(names.contains(&"team_shutdown_agent")); - assert!(names.contains(&"team_list_models")); - assert!(names.contains(&"team_describe_assistant")); + assert!(!names.contains(&"team_spawn_agent".to_owned())); + assert!(!names.contains(&"team_rename_agent".to_owned())); + assert!(!names.contains(&"team_shutdown_agent".to_owned())); + assert!(names.contains(&"team_send_message".to_owned())); + assert!(names.contains(&"team_task_create".to_owned())); + assert!(names.contains(&"team_task_update".to_owned())); + assert!(names.contains(&"team_task_list".to_owned())); + assert!(names.contains(&"team_members".to_owned())); + assert!(names.contains(&"team_list_models".to_owned())); + assert!(names.contains(&"team_describe_assistant".to_owned())); env.server.stop(); } @@ -688,6 +742,160 @@ async fn tra2_rename_nonexistent_agent() { env.server.stop(); } +#[tokio::test] +async fn mcp_non_lead_cannot_rename_agent() { + let env = setup().await; + let mut stream = connect_and_init(env.server.port(), "test-token-123", "worker-1").await; + + let resp = call_tool( + &mut stream, + 2, + "team_rename_agent", + json!({"slot_id": "worker-1", "new_name": "Renamed"}), + ) + .await; + + assert!(is_error_response(&resp)); + let text = extract_text(&resp); + assert!(text.contains("Only Lead")); + + env.server.stop(); +} + +#[tokio::test] +async fn http_mcp_tools_list_filters_lead_only_tools() { + let env = setup().await; + + let resp = http_rpc( + env.server.http_port(), + "worker-1", + json!({"jsonrpc": "2.0", "id": 10, "method": "tools/list"}), + ) + .await; + let names: Vec = resp["result"]["tools"] + .as_array() + .unwrap() + .iter() + .map(|tool| tool["name"].as_str().unwrap().to_owned()) + .collect(); + + assert!(!names.contains(&"team_spawn_agent".to_owned())); + assert!(!names.contains(&"team_rename_agent".to_owned())); + assert!(!names.contains(&"team_shutdown_agent".to_owned())); + assert!(names.contains(&"team_send_message".to_owned())); + + env.server.stop(); +} + +#[tokio::test] +async fn http_mcp_non_lead_cannot_rename_agent() { + let env = setup().await; + + let resp = http_rpc( + env.server.http_port(), + "worker-1", + json!({ + "jsonrpc": "2.0", + "id": 11, + "method": "tools/call", + "params": { + "name": "team_rename_agent", + "arguments": { + "slot_id": "worker-1", + "new_name": "Renamed" + } + } + }), + ) + .await; + + assert!(resp["result"]["isError"].as_bool().unwrap_or(false)); + let text = resp["result"]["content"][0]["text"].as_str().unwrap_or(""); + assert!(text.contains("Only Lead")); + + env.server.stop(); +} + +#[tokio::test] +async fn http_mcp_rejects_missing_auth_token() { + let env = setup().await; + + let resp = http_rpc_with_auth( + env.server.http_port(), + "worker-1", + None, + json!({"jsonrpc": "2.0", "id": 12, "method": "tools/list"}), + ) + .await; + + assert_eq!(resp["error"]["code"].as_i64(), Some(-32600)); + assert!( + resp["error"]["message"] + .as_str() + .unwrap_or_default() + .contains("Authentication failed") + ); + + env.server.stop(); +} + +#[tokio::test] +async fn http_mcp_rejects_invalid_auth_token() { + let env = setup().await; + + let resp = http_rpc_with_auth( + env.server.http_port(), + "worker-1", + Some("wrong-token"), + json!({"jsonrpc": "2.0", "id": 13, "method": "tools/list"}), + ) + .await; + + assert_eq!(resp["error"]["code"].as_i64(), Some(-32600)); + assert!( + resp["error"]["message"] + .as_str() + .unwrap_or_default() + .contains("Authentication failed") + ); + + env.server.stop(); +} + +#[tokio::test] +async fn http_mcp_rejects_lead_slot_spoof_without_valid_auth() { + let env = setup().await; + + let resp = http_rpc_with_auth( + env.server.http_port(), + "lead-1", + Some("wrong-token"), + json!({ + "jsonrpc": "2.0", + "id": 14, + "method": "tools/call", + "params": { + "name": "team_rename_agent", + "arguments": { + "slot_id": "worker-1", + "new_name": "Spoofed" + } + } + }), + ) + .await; + + assert_eq!(resp["error"]["code"].as_i64(), Some(-32600)); + assert!( + resp["error"]["message"] + .as_str() + .unwrap_or_default() + .contains("Authentication failed") + ); + + env.server.stop(); +} + // --------------------------------------------------------------------------- // Tests: team_shutdown_agent (TSA-1, TSA-4) // --------------------------------------------------------------------------- diff --git a/crates/aionui-team/tests/prompts_events_integration.rs b/crates/aionui-team/tests/prompts_events_integration.rs index b51d9a2df..2b9143dce 100644 --- a/crates/aionui-team/tests/prompts_events_integration.rs +++ b/crates/aionui-team/tests/prompts_events_integration.rs @@ -286,13 +286,20 @@ fn lp3_lead_prompt_contains_task_management_guidance() { #[test] fn tp1_teammate_prompt_contains_execution_guidance() { let agent = make_agent("w1", "Worker1", TeammateRole::Teammate); - let prompt = build_teammate_prompt(&agent, "Alpha"); + let members = vec![make_agent("lead-1", "Lead", TeammateRole::Lead), agent.clone()]; + let prompt = build_teammate_prompt(&agent, "Alpha", &members); - assert!(prompt.contains("execute tasks"), "missing execution guidance"); + assert!(prompt.contains("## Team Governance"), "missing governance"); + assert!( + prompt.contains("You MUST use the `team_*` MCP tools for ALL team coordination."), + "missing canonical coordination rule" + ); + assert!(prompt.contains("## How to Work"), "missing execution guidance"); assert!(prompt.contains("team_send_message"), "missing communication tool"); assert!(prompt.contains("team_task_update"), "missing task update tool"); assert!(prompt.contains("shutdown_request"), "missing shutdown protocol"); assert!(prompt.contains("shutdown_approved"), "missing shutdown_approved"); + assert!(prompt.contains("STOP GENERATING"), "missing stop protocol"); } // -- TP-2: Teammate prompt contains team name -------------------------------- @@ -300,9 +307,10 @@ fn tp1_teammate_prompt_contains_execution_guidance() { #[test] fn tp2_teammate_prompt_contains_team_name() { let agent = make_agent("w1", "Worker1", TeammateRole::Teammate); - let prompt = build_teammate_prompt(&agent, "Project Falcon"); + let members = vec![make_agent("lead-1", "Lead", TeammateRole::Lead), agent.clone()]; + let prompt = build_teammate_prompt(&agent, "Project Falcon", &members); - assert!(prompt.contains("\"Project Falcon\"")); + assert!(prompt.contains("Team: Project Falcon")); } // -- WP-1: Wake payload includes unread messages ----------------------------- diff --git a/crates/aionui-team/tests/session_service_integration.rs b/crates/aionui-team/tests/session_service_integration.rs index 51c9c8481..8414819fe 100644 --- a/crates/aionui-team/tests/session_service_integration.rs +++ b/crates/aionui-team/tests/session_service_integration.rs @@ -11,7 +11,7 @@ use aionui_ai_agent::task_manager::AgentFactory; use aionui_ai_agent::types::BuildTaskOptions; use aionui_ai_agent::{AgentError, IWorkerTaskManager, WorkerTaskManagerImpl}; use aionui_api_types::{AcpBuildExtra, AddAgentRequest, CreateTeamRequest, TeamAgentInput, WebSocketMessage}; -use aionui_common::{AgentKillReason, AgentType, PaginatedResult, ProviderWithModel}; +use aionui_common::{AgentKillReason, AgentType, ConversationStatus, PaginatedResult, ProviderWithModel}; use aionui_db::models::{ AgentMetadataRow, ConversationRow, MessageRow, UpdateAgentHandshakeParams, UpsertAgentMetadataParams, }; @@ -272,10 +272,18 @@ struct FakeConversationPorts { repo: Arc, broadcaster: Arc, workspace_root: std::path::PathBuf, + preset_snapshots: Mutex>, fail_team_temp_create: std::sync::atomic::AtomicBool, fail_leader_workspace_patch: std::sync::atomic::AtomicBool, } +#[derive(Clone)] +struct FakePresetAssistantSnapshot { + rules: String, + skills: Vec, + mcp_server_ids: Vec, +} + impl FakeConversationPorts { fn new(repo: Arc, broadcaster: Arc) -> Self { let workspace_root = @@ -284,10 +292,39 @@ impl FakeConversationPorts { repo, broadcaster, workspace_root, + preset_snapshots: Mutex::new(HashMap::new()), fail_team_temp_create: std::sync::atomic::AtomicBool::new(false), fail_leader_workspace_patch: std::sync::atomic::AtomicBool::new(false), } } + + fn upsert_preset_snapshot(&self, id: &str, snapshot: FakePresetAssistantSnapshot) { + self.preset_snapshots.lock().unwrap().insert(id.to_owned(), snapshot); + } + + fn apply_preset_snapshot(&self, extra: &mut serde_json::Value) { + let Some(preset_id) = extra + .get("preset_assistant_id") + .and_then(serde_json::Value::as_str) + .map(str::to_owned) + else { + return; + }; + let Some(snapshot) = self.preset_snapshots.lock().unwrap().get(&preset_id).cloned() else { + return; + }; + extra["preset_context"] = serde_json::Value::String(snapshot.rules.clone()); + extra["preset_rules"] = serde_json::Value::String(snapshot.rules); + extra["skills"] = + serde_json::Value::Array(snapshot.skills.into_iter().map(serde_json::Value::String).collect()); + extra["mcp_server_ids"] = serde_json::Value::Array( + snapshot + .mcp_server_ids + .into_iter() + .map(serde_json::Value::String) + .collect(), + ); + } } #[async_trait::async_trait] @@ -311,6 +348,7 @@ impl TeamConversationProvisioningPort for FakeConversationPorts { }); let mut extra = request.extra; extra["workspace"] = serde_json::Value::String(workspace.clone()); + self.apply_preset_snapshot(&mut extra); self.repo .create(&ConversationRow { id: id.clone(), @@ -901,17 +939,35 @@ mod mock_agent { pub workspace: String, pub event_tx: broadcast::Sender, pub confirmations: Vec, + pub status: Option>>>, } impl MockAgent { pub fn new(conversation_id: String, workspace: String) -> Self { - Self::with_confirmations(conversation_id, workspace, Vec::new()) + Self::with_confirmations_and_status(conversation_id, workspace, Vec::new(), None) } pub fn with_confirmations( conversation_id: String, workspace: String, confirmations: Vec, + ) -> Self { + Self::with_confirmations_and_status(conversation_id, workspace, confirmations, None) + } + + pub fn with_status( + conversation_id: String, + workspace: String, + status: std::sync::Arc>>, + ) -> Self { + Self::with_confirmations_and_status(conversation_id, workspace, Vec::new(), Some(status)) + } + + fn with_confirmations_and_status( + conversation_id: String, + workspace: String, + confirmations: Vec, + status: Option>>>, ) -> Self { let (event_tx, _) = broadcast::channel(16); Self { @@ -919,6 +975,7 @@ mod mock_agent { workspace, event_tx, confirmations, + status, } } } @@ -935,7 +992,7 @@ mod mock_agent { &self.workspace } fn status(&self) -> Option { - None + self.status.as_ref().and_then(|status| *status.lock().unwrap()) } fn last_activity_at(&self) -> TimestampMs { 0 @@ -1001,6 +1058,27 @@ fn confirmations_factory(count: usize) -> AgentFactory { }) } +fn status_factory_with_event_sender( + status: Arc>>, + event_sender: Arc>>>, +) -> AgentFactory { + use futures_util::FutureExt; + Arc::new(move |opts: BuildTaskOptions| { + let status = status.clone(); + let event_sender = event_sender.clone(); + async move { + let agent = mock_agent::MockAgent::with_status( + opts.context.conversation.conversation_id, + opts.context.workspace.path, + status, + ); + *event_sender.lock().unwrap() = Some(agent.event_tx.clone()); + Ok(aionui_ai_agent::AgentInstance::Mock(Arc::new(agent))) + } + .boxed() + }) +} + fn test_acp_build_options(conversation_id: String, workspace: String) -> BuildTaskOptions { BuildTaskOptions::new(AgentSessionContext { conversation: ConversationContext { @@ -1252,6 +1330,76 @@ async fn ensure_session_recovery_drain_runs_agent_turn_with_team_run_id() { assert!(requests[0].team_run_id.is_some(), "recovery turn must be TeamRun-owned"); } +#[tokio::test] +async fn teammate_first_wake_uses_canonical_prompt_at_service_boundary() { + let (svc, team_repo, turn_port, _conv_repo) = setup_with_recording_turn_port(); + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "Recover Teammate".into(), + agents: two_agent_input(), + workspace: None, + }, + ) + .await + .expect("create team"); + let worker_slot_id = created.agents[1].slot_id.clone(); + svc.stop_session("user1", &created.id) + .await + .expect("stop auto-started session"); + + team_repo + .write_message(&aionui_db::models::MailboxMessageRow { + id: "mailbox-worker-1".into(), + team_id: created.id.clone(), + to_agent_id: worker_slot_id.clone(), + from_agent_id: "user".into(), + msg_type: "message".into(), + content: "do X".into(), + summary: None, + files: None, + read: false, + created_at: aionui_common::now_ms(), + }) + .await + .expect("seed teammate mailbox"); + + svc.ensure_session("user1", &created.id).await.expect("ensure"); + + tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + if turn_port + .requests + .lock() + .unwrap() + .iter() + .any(|request| request.slot_id == worker_slot_id) + { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } + }) + .await + .expect("teammate recovery turn should run"); + + let requests = turn_port.requests.lock().unwrap(); + let worker_request = requests + .iter() + .find(|request| request.slot_id == worker_slot_id) + .expect("worker turn request"); + let first_message = &worker_request.content; + 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")); +} + #[tokio::test] async fn ensure_session_does_not_run_self_message_only_recovery_turn() { let (svc, team_repo, turn_port, _conv_repo) = setup_with_recording_turn_port(); @@ -1569,6 +1717,115 @@ async fn tc_create_team_carries_assistant_identity_into_lead_conversation_extra( assert_eq!(extra["preset_assistant_id"], serde_json::json!("2d23ff1c")); } +fn fake_preset_snapshot(rules: &str, skills: &[&str], mcp_server_ids: &[&str]) -> FakePresetAssistantSnapshot { + FakePresetAssistantSnapshot { + rules: rules.to_owned(), + skills: skills.iter().map(|value| (*value).to_owned()).collect(), + mcp_server_ids: mcp_server_ids.iter().map(|value| (*value).to_owned()).collect(), + } +} + +fn assert_frozen_preset_extra(extra: &serde_json::Value) { + assert_eq!(extra["preset_assistant_id"], serde_json::json!("word-creator")); + assert_eq!(extra["custom_agent_id"], serde_json::json!("word-creator")); + assert_eq!(extra["preset_context"], serde_json::json!("assistant rule body")); + assert_eq!(extra["preset_rules"], serde_json::json!("assistant rule body")); + assert_eq!(extra["skills"], serde_json::json!(["pdf", "cron"])); + assert_eq!(extra["mcp_server_ids"], serde_json::json!(["mcp-docs"])); +} + +#[tokio::test] +async fn team_preset_assistant_snapshot_is_frozen() { + let (svc, _team_repo, conversation_ports, conv_repo) = + setup_with_ports_team_repo_and_conversation_repo(success_factory(), Arc::new(StubAgentMetadataRepo::empty())); + conversation_ports.upsert_preset_snapshot( + "word-creator", + fake_preset_snapshot("assistant rule body", &["pdf", "cron"], &["mcp-docs"]), + ); + + let resp = svc + .create_team( + "user1", + CreateTeamRequest { + name: "Preset Team".into(), + agents: vec![TeamAgentInput { + name: "Lead".into(), + role: "lead".into(), + backend: "claude".into(), + model: "claude-sonnet-4".into(), + custom_agent_id: Some("word-creator".into()), + conversation_id: None, + }], + workspace: None, + }, + ) + .await + .expect("create team"); + + let extra = conv_repo.get_extra(&resp.agents[0].conversation_id).unwrap(); + assert_frozen_preset_extra(&extra); + + conversation_ports.upsert_preset_snapshot( + "word-creator", + fake_preset_snapshot("changed rule body", &["changed"], &["changed-mcp"]), + ); + + let after_live_change = conv_repo.get_extra(&resp.agents[0].conversation_id).unwrap(); + assert_frozen_preset_extra(&after_live_change); +} + +#[tokio::test] +async fn spawned_preset_assistant_snapshot_is_frozen() { + let (svc, _team_repo, conversation_ports, conv_repo) = + setup_with_ports_team_repo_and_conversation_repo(success_factory(), Arc::new(StubAgentMetadataRepo::empty())); + conversation_ports.upsert_preset_snapshot( + "word-creator", + fake_preset_snapshot("assistant rule body", &["pdf", "cron"], &["mcp-docs"]), + ); + + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "Spawn Preset".into(), + agents: two_agent_input(), + workspace: None, + }, + ) + .await + .expect("create team"); + let lead_slot_id = created.lead_agent_id.clone().expect("lead slot"); + svc.ensure_session("user1", &created.id).await.expect("ensure session"); + svc.send_message("user1", &created.id, "start active run", None) + .await + .expect("active run"); + + let spawned = svc + .spawn_agent_in_session( + &created.id, + &lead_slot_id, + SpawnAgentRequest { + name: "Writer".into(), + agent_type: Some("claude".into()), + custom_agent_id: Some("word-creator".into()), + model: Some("claude-sonnet-4".into()), + }, + ) + .await + .expect("spawn preset teammate"); + + let extra = conv_repo.get_extra(&spawned.conversation_id).unwrap(); + assert_frozen_preset_extra(&extra); + + conversation_ports.upsert_preset_snapshot( + "word-creator", + fake_preset_snapshot("changed rule body", &["changed"], &["changed-mcp"]), + ); + + let after_live_change = conv_repo.get_extra(&spawned.conversation_id).unwrap(); + assert_frozen_preset_extra(&after_live_change); +} + #[tokio::test] async fn ta_add_agent_uses_model_fallback_for_acp_backend() { let svc = setup_with_metadata_rows(vec![make_agent_metadata_row( @@ -2959,6 +3216,112 @@ async fn d9_ensure_session_kills_and_rebuilds_every_agent() { } } +#[tokio::test] +async fn d9_create_team_from_running_solo_leader_rebuilds_leader_after_turn_finishes() { + let status = Arc::new(Mutex::new(Some(ConversationStatus::Running))); + let event_sender = Arc::new(Mutex::new(None)); + let (svc, tm, conv_repo) = setup_with_factory_and_metadata_and_conversation_repo( + status_factory_with_event_sender(status.clone(), event_sender.clone()), + Arc::new(StubAgentMetadataRepo::empty()), + ); + let lead_conversation_id = "solo-lead"; + let workspace = "/tmp/aioncore-test-solo-lead"; + conv_repo + .create(&ConversationRow { + id: lead_conversation_id.to_owned(), + user_id: "user1".to_owned(), + name: "Solo Lead".to_owned(), + r#type: "acp".to_owned(), + pinned: false, + pinned_at: None, + source: None, + channel_chat_id: None, + extra: serde_json::json!({ + "backend": "claude", + "current_model_id": "opus", + "workspace": workspace, + }) + .to_string(), + model: None, + status: Some("running".to_owned()), + created_at: aionui_common::now_ms(), + updated_at: aionui_common::now_ms(), + }) + .await + .unwrap(); + tm.get_or_build_task( + lead_conversation_id, + test_acp_build_options(lead_conversation_id.to_owned(), workspace.to_owned()), + ) + .await + .unwrap(); + + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "Guide Upgrade".into(), + agents: vec![TeamAgentInput { + name: "Leader".into(), + role: "lead".into(), + backend: "claude".into(), + model: "opus".into(), + custom_agent_id: None, + conversation_id: Some(lead_conversation_id.to_owned()), + }], + workspace: None, + }, + ) + .await + .unwrap(); + + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + assert!( + tm.snapshot().kill.is_empty(), + "running solo leader must not be killed before the create_team tool turn can finish" + ); + + *status.lock().unwrap() = Some(ConversationStatus::Finished); + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + assert!( + tm.snapshot().kill.is_empty(), + "leader rebuild should wait for the agent terminal event, not poll status changes" + ); + + let sender = event_sender + .lock() + .unwrap() + .clone() + .expect("mock agent should expose its stream event sender"); + sender + .send(aionui_ai_agent::AgentStreamEvent::Finish( + aionui_ai_agent::protocol::events::FinishEventData { session_id: None }, + )) + .unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + let calls = tm.snapshot(); + if calls.kill.len() == 1 && calls.build.len() == 2 { + break calls; + } + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } + }) + .await + .expect("leader should be rebuilt after solo turn finishes"); + + let calls = tm.snapshot(); + assert_eq!( + calls.kill, + vec![(lead_conversation_id.to_owned(), Some(AgentKillReason::TeamMcpRebuild))] + ); + assert_eq!( + calls.build, + vec![lead_conversation_id.to_owned(), lead_conversation_id.to_owned()] + ); + assert_eq!(created.agents.len(), 1); +} + #[tokio::test] async fn d9_ensure_session_persists_team_mcp_stdio_config() { // Each agent's conversation.extra must carry a `team_mcp_stdio_config`