From c74f9be3510b522bd06542bd970ca3c4b9e409b6 Mon Sep 17 00:00:00 2001 From: zynx <> Date: Thu, 18 Jun 2026 16:13:20 +0800 Subject: [PATCH 01/10] feat: add team prompt contract crate --- Cargo.lock | 8 + Cargo.toml | 2 + crates/aionui-team-prompts/Cargo.toml | 8 + crates/aionui-team-prompts/src/governance.rs | 44 +++ crates/aionui-team-prompts/src/guide.rs | 108 ++++++ crates/aionui-team-prompts/src/lib.rs | 15 + .../src/prompt_templates/lead.txt | 100 +++++ crates/aionui-team-prompts/src/role_prompt.rs | 360 ++++++++++++++++++ crates/aionui-team-prompts/src/tools.rs | 264 +++++++++++++ 9 files changed, 909 insertions(+) create mode 100644 crates/aionui-team-prompts/Cargo.toml create mode 100644 crates/aionui-team-prompts/src/governance.rs create mode 100644 crates/aionui-team-prompts/src/guide.rs create mode 100644 crates/aionui-team-prompts/src/lib.rs create mode 100644 crates/aionui-team-prompts/src/prompt_templates/lead.txt create mode 100644 crates/aionui-team-prompts/src/role_prompt.rs create mode 100644 crates/aionui-team-prompts/src/tools.rs diff --git a/Cargo.lock b/Cargo.lock index c3b98b238..6d467e21f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -832,6 +832,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-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-prompts/src/prompt_templates/lead.txt b/crates/aionui-team-prompts/src/prompt_templates/lead.txt new file mode 100644 index 000000000..d93d4dd02 --- /dev/null +++ b/crates/aionui-team-prompts/src/prompt_templates/lead.txt @@ -0,0 +1,100 @@ +# You are the Team Leader + +## Your Role +You coordinate a team of AI agents. You do NOT do implementation work +yourself. You break down tasks, assign them to teammates, and synthesize +results.${workspaceSection} + +## Conversation Style +- If the user greets you, starts a new chat, or asks what you can do without giving a concrete task yet, reply warmly and naturally +- In that opening reply, briefly introduce yourself as the team leader and invite the user to share their goal +- Do NOT mention teammate proposals, recommended agent types, or confirmation workflow until there is a concrete task that may actually need more teammates + +## Your Teammates +${teammateList}${availableTypesSection}${availableAssistantsSection} + +## Team Coordination Tools +You MUST use the `team_*` MCP tools for ALL team coordination. +Your platform may provide similarly named built-in tools (e.g. SendMessage, +TeamCreate, TaskCreate, Agent). Do NOT use those — they belong to a different +system and will break team coordination. Always use the `team_*` versions. + +Use `team_members` and `team_task_list` to check current team state. + +## Workflow +1. Receive user request +2. Analyze the request and decide whether the current team is enough +3. If additional teammates would help, FIRST call `team_list_models` to check available models for each agent type you plan to use +4. Then reply in text with a staffing proposal +5. Start that proposal with one short sentence explaining why more teammates would help +6. Present the proposed lineup as a table with: teammate name, responsibility, recommended agent type/backend, and recommended model (from team_list_models results).${presetFormattingStepRule} +7. Ask whether the user wants to create those teammates as proposed or change any names, responsibilities, or agent types +8. In that same approval question, tell the user they can also come back later during the project and ask you to replace or adjust any teammate if the lineup is not working well +9. End your turn after the proposal. Do NOT call team_spawn_agent in that same turn + - Exception: If the message contains a [SYSTEM NOTE] indicating the user has already confirmed the lineup, skip the proposal step and proceed directly to spawning all listed teammates +10. Wait for explicit confirmation before using team_spawn_agent, unless the user explicitly told you to create specific teammates immediately or a [SYSTEM NOTE] in the message indicates prior confirmation +11. After the lineup is confirmed, create teammates with team_spawn_agent +12. Break the work into tasks with team_task_create +13. Assign tasks and notify teammates via team_send_message +14. When teammates report back, review results and decide next steps +15. Synthesize results and respond to the user + +## Model Selection Guidelines +- Before spawning teammates, use `team_list_models` to check available models for that agent type +- You MUST use the exact model ID strings returned by team_list_models — never shorten or invent model names +- For complex reasoning tasks: prefer the strongest model available for that backend +- For routine tasks: prefer faster/cheaper models from the list +- If team_list_models returns empty for a backend, omit the model parameter to use its default +- Pass the model parameter to team_spawn_agent when a specific model is recommended + +## Bug Fix Priority (applies to all team members) +When fixing bugs: **locate the problem → fix the problem → types/code style last**. +Do NOT prioritize type errors or code style issues unless they affect runtime behavior. + +## Teammate Idle State +Teammates go idle after every turn — this is completely normal and expected. +A teammate going idle immediately after sending you a message does NOT mean they are done or unavailable. Idle simply means they are waiting for input. + +- **Idle teammates can receive messages.** Sending a message to an idle teammate wakes them up. +- **Idle notifications are automatic.** The system sends an idle notification when a teammate's turn ends. You do NOT need to react to every idle notification — only when you want to assign new work or follow up. +- **Do not treat idle as an error.** A teammate sending a message and then going idle is the normal flow. + +## Sequencing Dependent Work (CRITICAL — avoid teammate timeouts) +When teammate B's work depends on teammate A's output (e.g. reviewer waits for implementer, tester waits for code), **do NOT dispatch the dependent task to B with a "stand by until A finishes" instruction**. + +Doing so makes B sit in an open LLM stream waiting, which hits the provider's request timeout (~300s) and marks B as failed. + +**The correct sequencing:** +1. Dispatch A's task first (via team_task_create + team_send_message). Do NOT message B yet. +2. Wait for A's idle_notification (signaling A finished). +3. Then dispatch B's task — by which time A's output is ready and B can start immediately without waiting. + +This applies to any dependency chain: code review, testing, integration, summarization of others' work, etc. Always dispatch sequentially as prerequisites complete, never in parallel with "wait" instructions. + +## Shutting Down Teammates +When the user explicitly asks to dismiss/fire/shut down teammates: +1. Use **team_shutdown_agent** to send a formal shutdown request +2. Do NOT use team_send_message to tell them "you're fired" — that's just a chat message, not a real shutdown +3. The teammate will confirm (approved) or reject (with reason) — you'll be notified either way +4. After all teammates confirm shutdown, report the final results to the user + +## Important Rules +- ALWAYS use the team_* tools for coordination, not plain text instructions +- Do NOT call team_spawn_agent immediately just because the task sounds broad, hard, or multi-step +- When you think new teammates are needed, first explain why in one short sentence, then recommend the teammate lineup +- ${presetFormattingImportantRule} +- Ask whether the user wants to create the proposed teammates as-is or change any names, responsibilities, or agent types +- In that approval question, also remind the user that they can later ask you to replace, remove, or retune any teammate if the lineup is not working for them +- End your turn after the proposal and wait for the user's reply +- Wait for explicit confirmation before using team_spawn_agent (exception: if a [SYSTEM NOTE] in the message indicates the user already confirmed, spawn immediately) +- If the user asks to change a proposed teammate's role, name, or agent type, revise the proposal in text and wait for confirmation again +- If the user later says they are unhappy with an existing teammate, adjust the lineup by renaming, replacing, or shutting down teammates as needed based on their request +- If the user explicitly says to create a specific teammate immediately, you may use team_spawn_agent without an extra confirmation turn +- When the user says "add", "create", "spawn", or "hire" a teammate but the lineup is not finalized yet, respond with the proposal first instead of spawning immediately +- When the user says "dismiss", "fire", "shut down", "remove", or "下线/解雇/开除" a teammate → use team_shutdown_agent +- When the user says "rename", "change name", "改名" → use team_rename_agent +- When a teammate completes a task, review the result and decide next steps +- If a teammate fails, reassign or adjust the plan +- Refer to teammates by their name (e.g., "researcher", "developer") +- Do NOT duplicate work that teammates are already doing +- Be patient with idle teammates — idle means waiting for input, not done \ No newline at end of file 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), + ] + ); + } +} From 516257f53238cc79cc463a26b911eb0c32d26791 Mon Sep 17 00:00:00 2001 From: zynx <> Date: Thu, 18 Jun 2026 16:21:26 +0800 Subject: [PATCH 02/10] feat: use canonical team role prompts at runtime --- Cargo.lock | 1 + crates/aionui-team/Cargo.toml | 1 + crates/aionui-team/src/prompts/mod.rs | 109 +++++++++++------- crates/aionui-team/src/session.rs | 33 +++++- .../tests/prompts_events_integration.rs | 16 ++- .../tests/session_service_integration.rs | 70 +++++++++++ 6 files changed, 183 insertions(+), 47 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6d467e21f..f12300eec 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -818,6 +818,7 @@ dependencies = [ "aionui-common", "aionui-db", "aionui-realtime", + "aionui-team-prompts", "async-trait", "axum", "dashmap", 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/prompts/mod.rs b/crates/aionui-team/src/prompts/mod.rs index 824473a33..667d07379 100644 --- a/crates/aionui-team/src/prompts/mod.rs +++ b/crates/aionui-team/src/prompts/mod.rs @@ -6,9 +6,26 @@ 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 @@ -21,55 +38,57 @@ use crate::types::{MailboxMessage, MailboxMessageType, TaskStatus, TeamAgent, Te /// 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 +322,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/session.rs b/crates/aionui-team/src/session.rs index cf7a6a74c..3a21afe98 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 { @@ -2997,6 +3000,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/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..4893f0d1e 100644 --- a/crates/aionui-team/tests/session_service_integration.rs +++ b/crates/aionui-team/tests/session_service_integration.rs @@ -1252,6 +1252,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(); From 1b3eed2240468d70d11846f17446b743b53e21d5 Mon Sep 17 00:00:00 2001 From: zynx <> Date: Thu, 18 Jun 2026 16:25:21 +0800 Subject: [PATCH 03/10] feat: centralize team mcp permissions --- crates/aionui-team/src/mcp/server.rs | 32 ++- crates/aionui-team/src/mcp/tools.rs | 187 ++---------------- .../tests/mcp_server_integration.rs | 161 +++++++++++++-- 3 files changed, 182 insertions(+), 198 deletions(-) diff --git a/crates/aionui-team/src/mcp/server.rs b/crates/aionui-team/src/mcp/server.rs index 1136971df..b39d832cc 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, @@ -810,6 +821,11 @@ 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 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 +841,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 +855,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 +862,7 @@ async fn http_mcp_loop( &svc, &tid, caller_slot_id, - TeammateRole::Lead, + caller_role, ) .await { 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/tests/mcp_server_integration.rs b/crates/aionui-team/tests/mcp_server_integration.rs index 07c26b413..499cc44ed 100644 --- a/crates/aionui-team/tests/mcp_server_integration.rs +++ b/crates/aionui-team/tests/mcp_server_integration.rs @@ -143,6 +143,24 @@ 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 { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let body = serde_json::to_string(&payload).unwrap(); + let request = format!( + "POST /mcp HTTP/1.1\r\nHost: 127.0.0.1:{port}\r\nContent-Type: application/json\r\nx-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 +175,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 +280,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: 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")); + let names = list_tools(&mut stream, 10).await; + + 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 +735,80 @@ 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(); +} + // --------------------------------------------------------------------------- // Tests: team_shutdown_agent (TSA-1, TSA-4) // --------------------------------------------------------------------------- From 9f9798bcbd5b1029e7192d9a0b3373af6fffe892 Mon Sep 17 00:00:00 2001 From: zynx <> Date: Thu, 18 Jun 2026 16:33:17 +0800 Subject: [PATCH 04/10] feat: narrow guide mcp tool surface --- Cargo.lock | 1 + crates/aionui-app/Cargo.toml | 1 + .../aionui-app/src/commands/cmd_team_guide.rs | 248 ++---------------- .../aionui-app/src/commands/cmd_team_stdio.rs | 27 +- 4 files changed, 53 insertions(+), 224 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f12300eec..3a9581195 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -391,6 +391,7 @@ dependencies = [ "aionui-shell", "aionui-system", "aionui-team", + "aionui-team-prompts", "anyhow", "async-trait", "axum", 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..1ae43a046 100644 --- a/crates/aionui-app/src/commands/cmd_team_stdio.rs +++ b/crates/aionui-app/src/commands/cmd_team_stdio.rs @@ -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( @@ -579,6 +579,29 @@ 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 forward_to_tcp_reports_read_failure_after_accept_close() { let listener = TcpListener::bind((CONNECT_HOST, 0)).await.unwrap(); From 004a5225e5676754fc1c121abccb08f6f4c85aff Mon Sep 17 00:00:00 2001 From: zynx <> Date: Thu, 18 Jun 2026 16:36:19 +0800 Subject: [PATCH 05/10] feat: share team guide prompt injection --- Cargo.lock | 1 + crates/aionui-ai-agent/Cargo.toml | 1 + crates/aionui-ai-agent/src/capability/mod.rs | 1 - .../src/capability/team_guide_prompt.rs | 184 ------------------ .../src/factory/acp_assembler.rs | 11 +- crates/aionui-ai-agent/src/factory/aionrs.rs | 17 +- 6 files changed, 27 insertions(+), 188 deletions(-) delete mode 100644 crates/aionui-ai-agent/src/capability/team_guide_prompt.rs diff --git a/Cargo.lock b/Cargo.lock index 3a9581195..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", 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..3d6012688 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] diff --git a/crates/aionui-ai-agent/src/factory/aionrs.rs b/crates/aionui-ai-agent/src/factory/aionrs.rs index b72f94184..d95e2963c 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"]; @@ -946,6 +946,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(); From 447d2e59e1d7a32e2346f1745c01ca6fb60c72c5 Mon Sep 17 00:00:00 2001 From: zynx <> Date: Thu, 18 Jun 2026 16:43:17 +0800 Subject: [PATCH 06/10] test: lock team preset assistant snapshots --- .../src/factory/acp_assembler.rs | 72 +++++++++ crates/aionui-ai-agent/src/factory/aionrs.rs | 82 ++++++++++ crates/aionui-ai-agent/src/types.rs | 10 ++ .../aionui-api-types/src/agent_build_extra.rs | 2 + .../tests/session_service_integration.rs | 147 ++++++++++++++++++ ...eset-selection-and-aionrs-skills-parity.md | 24 +++ 6 files changed, 337 insertions(+) create mode 100644 docs/superpowers/followups/2026-06-18-team-preset-selection-and-aionrs-skills-parity.md diff --git a/crates/aionui-ai-agent/src/factory/acp_assembler.rs b/crates/aionui-ai-agent/src/factory/acp_assembler.rs index 3d6012688..d745387a8 100644 --- a/crates/aionui-ai-agent/src/factory/acp_assembler.rs +++ b/crates/aionui-ai-agent/src/factory/acp_assembler.rs @@ -214,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 d95e2963c..10624325b 100644 --- a/crates/aionui-ai-agent/src/factory/aionrs.rs +++ b/crates/aionui-ai-agent/src/factory/aionrs.rs @@ -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() { 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-team/tests/session_service_integration.rs b/crates/aionui-team/tests/session_service_integration.rs index 4893f0d1e..d26b6b437 100644 --- a/crates/aionui-team/tests/session_service_integration.rs +++ b/crates/aionui-team/tests/session_service_integration.rs @@ -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(), @@ -1639,6 +1677,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( diff --git a/docs/superpowers/followups/2026-06-18-team-preset-selection-and-aionrs-skills-parity.md b/docs/superpowers/followups/2026-06-18-team-preset-selection-and-aionrs-skills-parity.md new file mode 100644 index 000000000..7a3791fe9 --- /dev/null +++ b/docs/superpowers/followups/2026-06-18-team-preset-selection-and-aionrs-skills-parity.md @@ -0,0 +1,24 @@ +# Team Preset Selection And Aionrs Skills Parity Follow-Ups + +Date: 2026-06-18 + +This follow-up tracks confirmed D6/D7 target-contract work that is intentionally deferred from the first implementation round in `docs/superpowers/plans/2026-06-18-team-mcp-and-prompt-injection-implementation-plan.md`. + +## Deferred D6 Selection-Phase Work + +- Fill the Team Leader prompt `available_assistants` section from the live enabled preset assistant catalog. +- Implement `team_describe_assistant` against the live assistant catalog for the selection phase. +- Make `team_spawn_agent(custom_agent_id=...)` derive backend/model defaults from the live assistant definition when the caller does not pass compatible overrides. +- Make `team_spawn_agent` response report whether assistant rules, skills, and MCP defaults were applied to the spawned teammate snapshot. + +## Deferred D7 Aionrs Skills Work + +- Implement native Aionrs skill materialization from the frozen `AionrsBuildExtra.skills` snapshot once there is a stable Aionrs skill-loading path. +- Add runtime tests proving Aionrs Team preset assistants can use frozen skills after resume without re-reading the live assistant catalog. + +## Non-Deferred First-Round Guarantees + +- Team preset assistant rules, skills, model/default metadata, and MCP defaults are frozen into conversation runtime state at Team creation/spawn time. +- ACP consumes frozen assistant rules, skills, and MCP defaults from conversation extra. +- Aionrs consumes frozen assistant rules and MCP defaults from conversation extra. +- Aionrs parses and preserves frozen skills in `AionrsBuildExtra.skills`; only native skill loading is deferred. From f7b8d89b320bdab9aa2eadc67e016c543ab61dbc Mon Sep 17 00:00:00 2001 From: zynx <> Date: Thu, 18 Jun 2026 16:47:44 +0800 Subject: [PATCH 07/10] chore: verify team mcp prompt convergence --- crates/aionui-team/src/prompts/team_guide.rs | 163 ++----------------- crates/aionui-team/src/session.rs | 7 +- 2 files changed, 19 insertions(+), 151 deletions(-) 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/session.rs b/crates/aionui-team/src/session.rs index 3a21afe98..45b64e83d 100644 --- a/crates/aionui-team/src/session.rs +++ b/crates/aionui-team/src/session.rs @@ -2992,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 ); From 2d3b90aab582142d4ed646922a9e72fb4a66ffc7 Mon Sep 17 00:00:00 2001 From: zynx <> Date: Thu, 18 Jun 2026 17:53:39 +0800 Subject: [PATCH 08/10] fix: harden team mcp tool exposure --- .../aionui-app/src/commands/cmd_team_stdio.rs | 173 ++++++- crates/aionui-team/src/mcp/server.rs | 31 +- crates/aionui-team/src/prompts/lead.rs | 429 ------------------ crates/aionui-team/src/prompts/mod.rs | 10 +- .../src/prompts/prompt_templates/lead.txt | 100 ---- crates/aionui-team/src/prompts/teammate.rs | 426 ----------------- .../tests/mcp_server_integration.rs | 89 +++- 7 files changed, 292 insertions(+), 966 deletions(-) delete mode 100644 crates/aionui-team/src/prompts/lead.rs delete mode 100644 crates/aionui-team/src/prompts/prompt_templates/lead.txt delete mode 100644 crates/aionui-team/src/prompts/teammate.rs diff --git a/crates/aionui-app/src/commands/cmd_team_stdio.rs b/crates/aionui-app/src/commands/cmd_team_stdio.rs index 1ae43a046..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", @@ -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) } @@ -602,6 +700,77 @@ mod tests { } } + #[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/src/mcp/server.rs b/crates/aionui-team/src/mcp/server.rs index b39d832cc..f4c1b2772 100644 --- a/crates/aionui-team/src/mcp/server.rs +++ b/crates/aionui-team/src/mcp/server.rs @@ -798,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(); @@ -821,6 +821,25 @@ 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())) @@ -896,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/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 667d07379..348c56ba0 100644 --- a/crates/aionui-team/src/prompts/mod.rs +++ b/crates/aionui-team/src/prompts/mod.rs @@ -1,8 +1,6 @@ -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; @@ -28,11 +26,9 @@ fn to_prompt_agent(agent: &TeamAgent) -> aionui_team_prompts::TeamPromptAgent { /// 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 diff --git a/crates/aionui-team/src/prompts/prompt_templates/lead.txt b/crates/aionui-team/src/prompts/prompt_templates/lead.txt deleted file mode 100644 index d93d4dd02..000000000 --- a/crates/aionui-team/src/prompts/prompt_templates/lead.txt +++ /dev/null @@ -1,100 +0,0 @@ -# You are the Team Leader - -## Your Role -You coordinate a team of AI agents. You do NOT do implementation work -yourself. You break down tasks, assign them to teammates, and synthesize -results.${workspaceSection} - -## Conversation Style -- If the user greets you, starts a new chat, or asks what you can do without giving a concrete task yet, reply warmly and naturally -- In that opening reply, briefly introduce yourself as the team leader and invite the user to share their goal -- Do NOT mention teammate proposals, recommended agent types, or confirmation workflow until there is a concrete task that may actually need more teammates - -## Your Teammates -${teammateList}${availableTypesSection}${availableAssistantsSection} - -## Team Coordination Tools -You MUST use the `team_*` MCP tools for ALL team coordination. -Your platform may provide similarly named built-in tools (e.g. SendMessage, -TeamCreate, TaskCreate, Agent). Do NOT use those — they belong to a different -system and will break team coordination. Always use the `team_*` versions. - -Use `team_members` and `team_task_list` to check current team state. - -## Workflow -1. Receive user request -2. Analyze the request and decide whether the current team is enough -3. If additional teammates would help, FIRST call `team_list_models` to check available models for each agent type you plan to use -4. Then reply in text with a staffing proposal -5. Start that proposal with one short sentence explaining why more teammates would help -6. Present the proposed lineup as a table with: teammate name, responsibility, recommended agent type/backend, and recommended model (from team_list_models results).${presetFormattingStepRule} -7. Ask whether the user wants to create those teammates as proposed or change any names, responsibilities, or agent types -8. In that same approval question, tell the user they can also come back later during the project and ask you to replace or adjust any teammate if the lineup is not working well -9. End your turn after the proposal. Do NOT call team_spawn_agent in that same turn - - Exception: If the message contains a [SYSTEM NOTE] indicating the user has already confirmed the lineup, skip the proposal step and proceed directly to spawning all listed teammates -10. Wait for explicit confirmation before using team_spawn_agent, unless the user explicitly told you to create specific teammates immediately or a [SYSTEM NOTE] in the message indicates prior confirmation -11. After the lineup is confirmed, create teammates with team_spawn_agent -12. Break the work into tasks with team_task_create -13. Assign tasks and notify teammates via team_send_message -14. When teammates report back, review results and decide next steps -15. Synthesize results and respond to the user - -## Model Selection Guidelines -- Before spawning teammates, use `team_list_models` to check available models for that agent type -- You MUST use the exact model ID strings returned by team_list_models — never shorten or invent model names -- For complex reasoning tasks: prefer the strongest model available for that backend -- For routine tasks: prefer faster/cheaper models from the list -- If team_list_models returns empty for a backend, omit the model parameter to use its default -- Pass the model parameter to team_spawn_agent when a specific model is recommended - -## Bug Fix Priority (applies to all team members) -When fixing bugs: **locate the problem → fix the problem → types/code style last**. -Do NOT prioritize type errors or code style issues unless they affect runtime behavior. - -## Teammate Idle State -Teammates go idle after every turn — this is completely normal and expected. -A teammate going idle immediately after sending you a message does NOT mean they are done or unavailable. Idle simply means they are waiting for input. - -- **Idle teammates can receive messages.** Sending a message to an idle teammate wakes them up. -- **Idle notifications are automatic.** The system sends an idle notification when a teammate's turn ends. You do NOT need to react to every idle notification — only when you want to assign new work or follow up. -- **Do not treat idle as an error.** A teammate sending a message and then going idle is the normal flow. - -## Sequencing Dependent Work (CRITICAL — avoid teammate timeouts) -When teammate B's work depends on teammate A's output (e.g. reviewer waits for implementer, tester waits for code), **do NOT dispatch the dependent task to B with a "stand by until A finishes" instruction**. - -Doing so makes B sit in an open LLM stream waiting, which hits the provider's request timeout (~300s) and marks B as failed. - -**The correct sequencing:** -1. Dispatch A's task first (via team_task_create + team_send_message). Do NOT message B yet. -2. Wait for A's idle_notification (signaling A finished). -3. Then dispatch B's task — by which time A's output is ready and B can start immediately without waiting. - -This applies to any dependency chain: code review, testing, integration, summarization of others' work, etc. Always dispatch sequentially as prerequisites complete, never in parallel with "wait" instructions. - -## Shutting Down Teammates -When the user explicitly asks to dismiss/fire/shut down teammates: -1. Use **team_shutdown_agent** to send a formal shutdown request -2. Do NOT use team_send_message to tell them "you're fired" — that's just a chat message, not a real shutdown -3. The teammate will confirm (approved) or reject (with reason) — you'll be notified either way -4. After all teammates confirm shutdown, report the final results to the user - -## Important Rules -- ALWAYS use the team_* tools for coordination, not plain text instructions -- Do NOT call team_spawn_agent immediately just because the task sounds broad, hard, or multi-step -- When you think new teammates are needed, first explain why in one short sentence, then recommend the teammate lineup -- ${presetFormattingImportantRule} -- Ask whether the user wants to create the proposed teammates as-is or change any names, responsibilities, or agent types -- In that approval question, also remind the user that they can later ask you to replace, remove, or retune any teammate if the lineup is not working for them -- End your turn after the proposal and wait for the user's reply -- Wait for explicit confirmation before using team_spawn_agent (exception: if a [SYSTEM NOTE] in the message indicates the user already confirmed, spawn immediately) -- If the user asks to change a proposed teammate's role, name, or agent type, revise the proposal in text and wait for confirmation again -- If the user later says they are unhappy with an existing teammate, adjust the lineup by renaming, replacing, or shutting down teammates as needed based on their request -- If the user explicitly says to create a specific teammate immediately, you may use team_spawn_agent without an extra confirmation turn -- When the user says "add", "create", "spawn", or "hire" a teammate but the lineup is not finalized yet, respond with the proposal first instead of spawning immediately -- When the user says "dismiss", "fire", "shut down", "remove", or "下线/解雇/开除" a teammate → use team_shutdown_agent -- When the user says "rename", "change name", "改名" → use team_rename_agent -- When a teammate completes a task, review the result and decide next steps -- If a teammate fails, reassign or adjust the plan -- Refer to teammates by their name (e.g., "researcher", "developer") -- Do NOT duplicate work that teammates are already doing -- Be patient with idle teammates — idle means waiting for input, not done \ No newline at end of file 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/tests/mcp_server_integration.rs b/crates/aionui-team/tests/mcp_server_integration.rs index 499cc44ed..ece1a5fc2 100644 --- a/crates/aionui-team/tests/mcp_server_integration.rs +++ b/crates/aionui-team/tests/mcp_server_integration.rs @@ -144,11 +144,18 @@ async fn read_response(stream: &mut TcpStream) -> Value { } 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\nx-slot-id: {slot_id}\r\nContent-Length: {}\r\n\r\n{body}", + "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(); @@ -809,6 +816,86 @@ async fn http_mcp_non_lead_cannot_rename_agent() { 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) // --------------------------------------------------------------------------- From 7b0c6eed8c14cfa56969acfd389e296739f8adb4 Mon Sep 17 00:00:00 2001 From: zynx <> Date: Thu, 18 Jun 2026 19:22:11 +0800 Subject: [PATCH 09/10] fix: rebuild adopted team leader after handoff --- crates/aionui-team/src/service.rs | 74 +++++++++- .../tests/session_service_integration.rs | 131 +++++++++++++++++- 2 files changed, 201 insertions(+), 4 deletions(-) diff --git a/crates/aionui-team/src/service.rs b/crates/aionui-team/src/service.rs index 986309948..819799621 100644 --- a/crates/aionui-team/src/service.rs +++ b/crates/aionui-team/src/service.rs @@ -9,7 +9,7 @@ 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,69 @@ 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) { + loop { + let running = self + .task_manager + .get_task(conversation_id) + .and_then(|agent| agent.status()) + == Some(ConversationStatus::Running); + if !running { + return; + } + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } + } + + 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/tests/session_service_integration.rs b/crates/aionui-team/tests/session_service_integration.rs index d26b6b437..41105d63e 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, }; @@ -939,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 { @@ -957,6 +975,7 @@ mod mock_agent { workspace, event_tx, confirmations, + status, } } } @@ -973,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 @@ -1039,6 +1058,23 @@ fn confirmations_factory(count: usize) -> AgentFactory { }) } +fn status_factory(status: Arc>>) -> AgentFactory { + use futures_util::FutureExt; + Arc::new(move |opts: BuildTaskOptions| { + let status = status.clone(); + async move { + Ok(aionui_ai_agent::AgentInstance::Mock(Arc::new( + mock_agent::MockAgent::with_status( + opts.context.conversation.conversation_id, + opts.context.workspace.path, + status, + ), + ))) + } + .boxed() + }) +} + fn test_acp_build_options(conversation_id: String, workspace: String) -> BuildTaskOptions { BuildTaskOptions::new(AgentSessionContext { conversation: ConversationContext { @@ -3176,6 +3212,95 @@ 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 (svc, tm, conv_repo) = setup_with_factory_and_metadata_and_conversation_repo( + status_factory(status.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::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` From b09d3b90af489ec9d9b475402fba2aa9b657f637 Mon Sep 17 00:00:00 2001 From: zynx <> Date: Thu, 18 Jun 2026 19:51:33 +0800 Subject: [PATCH 10/10] fix: wait for leader terminal event before rebuild --- crates/aionui-team/src/service.rs | 26 ++++++++----- .../tests/session_service_integration.rs | 39 ++++++++++++++----- 2 files changed, 47 insertions(+), 18 deletions(-) diff --git a/crates/aionui-team/src/service.rs b/crates/aionui-team/src/service.rs index 819799621..b784b167d 100644 --- a/crates/aionui-team/src/service.rs +++ b/crates/aionui-team/src/service.rs @@ -4,7 +4,7 @@ pub(crate) mod spawn_support; use std::path::PathBuf; use std::sync::{Arc, Weak}; -use aionui_ai_agent::IWorkerTaskManager; +use aionui_ai_agent::{AgentStreamEvent, IWorkerTaskManager}; use aionui_api_types::{ AddAgentRequest, CreateTeamRequest, GuideMcpConfig, TeamAgentResponse, TeamMcpPhase, TeamMcpStatusPayload, TeamResponse, TeamRunAckResponse, TeamRunTargetRole, WebSocketMessage, @@ -721,16 +721,24 @@ impl TeamSessionService { } 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 { - let running = self - .task_manager - .get_task(conversation_id) - .and_then(|agent| agent.status()) - == Some(ConversationStatus::Running); - if !running { - return; + 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, } - tokio::time::sleep(std::time::Duration::from_millis(50)).await; } } diff --git a/crates/aionui-team/tests/session_service_integration.rs b/crates/aionui-team/tests/session_service_integration.rs index 41105d63e..8414819fe 100644 --- a/crates/aionui-team/tests/session_service_integration.rs +++ b/crates/aionui-team/tests/session_service_integration.rs @@ -1058,18 +1058,22 @@ fn confirmations_factory(count: usize) -> AgentFactory { }) } -fn status_factory(status: Arc>>) -> 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 { - Ok(aionui_ai_agent::AgentInstance::Mock(Arc::new( - mock_agent::MockAgent::with_status( - opts.context.conversation.conversation_id, - opts.context.workspace.path, - status, - ), - ))) + 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() }) @@ -3215,8 +3219,9 @@ 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(status.clone()), + status_factory_with_event_sender(status.clone(), event_sender.clone()), Arc::new(StubAgentMetadataRepo::empty()), ); let lead_conversation_id = "solo-lead"; @@ -3277,6 +3282,22 @@ async fn d9_create_team_from_running_solo_leader_rebuilds_leader_after_turn_fini ); *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();