agent_sdk/agent/
team.rs

1use std::sync::Arc;
2
3use tokio::sync::mpsc;
4use tracing::info;
5use uuid::Uuid;
6
7use crate::config::{AgentConfig, LlmConfig};
8use crate::error::SdkResult;
9use crate::traits::llm_client::LlmClient;
10use crate::traits::prompt_builder::{DefaultPromptBuilder, PromptBuilder};
11use crate::mailbox::broker::MessageBroker;
12use crate::storage::AgentPaths;
13use crate::task::store::TaskStore;
14use crate::types::task::Task;
15
16use super::agent_loop::{AgentLoop, AgentLoopResult};
17use super::events::AgentEvent;
18use super::hooks::HookRegistry;
19use super::memory::MemoryStore;
20use super::team_lead::{ExecutionSummary, TeamLead, TeammateSpec};
21
22/// Result of an `AgentTeam::run()` call.
23#[derive(Debug)]
24pub enum TeamResult {
25    /// Task was handled by a single agent (no team needed).
26    Single(AgentLoopResult),
27    /// Task was handled by a team of agents.
28    Team(ExecutionSummary),
29}
30
31impl TeamResult {
32    pub fn total_tokens(&self) -> u64 {
33        match self {
34            Self::Single(r) => r.total_tokens,
35            Self::Team(s) => s.total_tokens_used,
36        }
37    }
38}
39
40/// High-level entry point for the agent SDK.
41///
42/// `AgentTeam` coordinates multiple agent instances working together.
43/// One session acts as the team lead, coordinating work, assigning tasks,
44/// and synthesizing results. Teammates work independently, each in its own
45/// context window, and can communicate directly with each other.
46///
47/// There is no separate planning step — the lead IS the intelligence that
48/// decides how to organize work, just like Claude Code's agent teams.
49///
50/// # Usage patterns
51///
52/// **You describe the team** — add teammates with roles, add tasks, and run:
53/// ```rust,no_run
54/// # use agent_sdk::agent::team::AgentTeam;
55/// # use agent_sdk::config::{LlmConfig, AgentConfig};
56/// # async fn ex() -> anyhow::Result<()> {
57/// let result = AgentTeam::new(LlmConfig::default(), AgentConfig::default())
58///     .add_teammate("security", "Review for security vulnerabilities")
59///     .add_teammate("performance", "Review for performance issues")
60///     .run("Review the auth module")
61///     .await?;
62/// # Ok(()) }
63/// ```
64///
65/// **Single agent** — for simple tasks, skip the team entirely:
66/// ```rust,no_run
67/// # use agent_sdk::agent::team::AgentTeam;
68/// # use agent_sdk::config::{LlmConfig, AgentConfig};
69/// # async fn ex() -> anyhow::Result<()> {
70/// let result = AgentTeam::new(LlmConfig::default(), AgentConfig::default())
71///     .run_single("Explain this codebase")
72///     .await?;
73/// # Ok(()) }
74/// ```
75pub struct AgentTeam {
76    llm_config: LlmConfig,
77    agent_config: AgentConfig,
78    llm_client: Option<Arc<dyn LlmClient>>,
79    prompt_builder: Arc<dyn PromptBuilder>,
80    hooks: HookRegistry,
81    source_root: std::path::PathBuf,
82    work_dir: std::path::PathBuf,
83    event_tx: Option<mpsc::UnboundedSender<AgentEvent>>,
84    /// Explicit teammates to spawn.
85    teammate_specs: Vec<TeammateSpec>,
86    /// Pre-created tasks for the team.
87    tasks: Vec<Task>,
88}
89
90impl AgentTeam {
91    /// Create a new AgentTeam with the given LLM and agent configuration.
92    pub fn new(llm_config: LlmConfig, agent_config: AgentConfig) -> Self {
93        Self {
94            llm_config,
95            agent_config,
96            llm_client: None,
97            prompt_builder: Arc::new(DefaultPromptBuilder),
98            hooks: HookRegistry::new(),
99            source_root: std::path::PathBuf::from("."),
100            work_dir: std::path::PathBuf::from("./output"),
101            event_tx: None,
102            teammate_specs: Vec::new(),
103            tasks: Vec::new(),
104        }
105    }
106
107    /// Set the source root directory (read-only source code).
108    pub fn source_root(mut self, path: impl Into<std::path::PathBuf>) -> Self {
109        self.source_root = path.into();
110        self
111    }
112
113    /// Set the working/output directory.
114    pub fn work_dir(mut self, path: impl Into<std::path::PathBuf>) -> Self {
115        self.work_dir = path.into();
116        self
117    }
118
119    /// Set a custom prompt builder.
120    pub fn prompt_builder(mut self, builder: Arc<dyn PromptBuilder>) -> Self {
121        self.prompt_builder = builder;
122        self
123    }
124
125    /// Set an event channel for monitoring agent activity.
126    pub fn event_channel(mut self, tx: mpsc::UnboundedSender<AgentEvent>) -> Self {
127        self.event_tx = Some(tx);
128        self
129    }
130
131    /// Provide a pre-created LLM client (skips creating one from config).
132    pub fn llm_client(mut self, client: Arc<dyn LlmClient>) -> Self {
133        self.llm_client = Some(client);
134        self
135    }
136
137    /// Add a hook for quality gates (TeammateIdle, TaskCreated, TaskCompleted).
138    pub fn add_hook(mut self, hook: impl super::hooks::Hook + 'static) -> Self {
139        self.hooks.add(hook);
140        self
141    }
142
143    /// Add a named teammate with a specific role.
144    ///
145    /// ```rust,no_run
146    /// # use agent_sdk::agent::team::AgentTeam;
147    /// # use agent_sdk::config::{LlmConfig, AgentConfig};
148    /// AgentTeam::new(LlmConfig::default(), AgentConfig::default())
149    ///     .add_teammate("security-reviewer", "Review for security vulnerabilities")
150    ///     .add_teammate("perf-reviewer", "Review for performance issues");
151    /// ```
152    pub fn add_teammate(
153        mut self,
154        name: impl Into<String>,
155        prompt: impl Into<String>,
156    ) -> Self {
157        self.teammate_specs.push(TeammateSpec {
158            name: name.into(),
159            prompt: prompt.into(),
160            require_plan_approval: false,
161        });
162        self
163    }
164
165    /// Add a teammate that must get plan approval from the lead before implementing.
166    /// The teammate generates a plan, sends it to the lead for review, and only
167    /// proceeds after the lead approves.
168    pub fn add_teammate_with_plan_approval(
169        mut self,
170        name: impl Into<String>,
171        prompt: impl Into<String>,
172    ) -> Self {
173        self.teammate_specs.push(TeammateSpec {
174            name: name.into(),
175            prompt: prompt.into(),
176            require_plan_approval: true,
177        });
178        self
179    }
180
181    /// Add a task for the team to work on.
182    ///
183    /// ```rust,no_run
184    /// # use agent_sdk::agent::team::AgentTeam;
185    /// # use agent_sdk::config::{LlmConfig, AgentConfig};
186    /// # use agent_sdk::types::task::Task;
187    /// let task1 = Task::new("gen", "Create config", "...", "config.rs");
188    /// let task2 = Task::new("gen", "Create server", "...", "server.rs")
189    ///     .with_dependencies(vec![task1.id]);
190    ///
191    /// AgentTeam::new(LlmConfig::default(), AgentConfig::default())
192    ///     .add_task(task1)
193    ///     .add_task(task2);
194    /// ```
195    pub fn add_task(mut self, task: Task) -> Self {
196        self.tasks.push(task);
197        self
198    }
199
200    /// Run the team. The lead spawns teammates, they claim tasks from the
201    /// shared task list, and work until all tasks are done.
202    pub async fn run(mut self, _goal: &str) -> SdkResult<TeamResult> {
203        let client = match self.llm_client.take() {
204            Some(c) => c,
205            None => crate::llm::create_client(&self.llm_config)?,
206        };
207        let paths = AgentPaths::for_work_dir(&self.work_dir)?;
208        let team_name = paths.new_team_name();
209        let team_config_path = paths.team_config_path(&team_name);
210
211        let hooks = Arc::new(std::mem::take(&mut self.hooks));
212        let task_store = Arc::new(TaskStore::new(paths.team_tasks_dir(&team_name)));
213        task_store.init()?;
214
215        // Add tasks to the store
216        for task in &self.tasks {
217            let hook_result = hooks.evaluate(
218                &super::hooks::HookEvent::TaskCreated { task: task.clone() },
219            );
220            if let super::hooks::HookResult::Reject { feedback } = hook_result {
221                self.emit_event(AgentEvent::HookRejected {
222                    event_name: "TaskCreated".to_string(),
223                    feedback,
224                });
225                continue;
226            }
227            task_store.create_task(task)?;
228        }
229
230        std::fs::create_dir_all(paths.team_dir(&team_name)).map_err(crate::error::SdkError::Io)?;
231        let broker = Arc::new(MessageBroker::new(paths.team_mailbox_dir(&team_name))?);
232        let memory = Arc::new(MemoryStore::new(paths.team_memory_dir(&team_name))?);
233
234        let lead = TeamLead {
235            id: Uuid::new_v4(),
236            team_name,
237            team_config_path,
238            task_store,
239            broker,
240            llm_client: client,
241            prompt_builder: self.prompt_builder.clone(),
242            config: self.agent_config.clone(),
243            source_root: self.source_root.clone(),
244            work_dir: self.work_dir.clone(),
245            memory_store: memory,
246            event_tx: self.event_tx.clone(),
247            hooks,
248            teammate_specs: self.teammate_specs.clone(),
249        };
250
251        self.emit_event(AgentEvent::TeamSpawned {
252            teammate_count: self.teammate_specs.len().max(self.agent_config.max_parallel_agents),
253        });
254
255        lead.run().await.map(TeamResult::Team)
256    }
257
258    /// Run as a single agent (no team). For simple, focused tasks.
259    pub async fn run_single(mut self, user_message: &str) -> SdkResult<AgentLoopResult> {
260        let client = match self.llm_client.take() {
261            Some(c) => c,
262            None => crate::llm::create_client(&self.llm_config)?,
263        };
264
265        use crate::tools::command_tools::RunCommandTool;
266        use crate::tools::fs_tools::{ListDirectoryTool, ReadFileTool, WriteFileTool};
267        use crate::tools::registry::ToolRegistry;
268        use crate::tools::search_tools::SearchFilesTool;
269        use crate::tools::web_tools::WebSearchTool;
270
271        let mut tools = ToolRegistry::new();
272        tools.register(Arc::new(ReadFileTool {
273            source_root: self.source_root.clone(),
274            work_dir: self.work_dir.clone(),
275        }));
276        tools.register(Arc::new(WriteFileTool {
277            work_dir: self.work_dir.clone(),
278        }));
279        tools.register(Arc::new(ListDirectoryTool {
280            source_root: self.source_root.clone(),
281            work_dir: self.work_dir.clone(),
282        }));
283        tools.register(Arc::new(SearchFilesTool {
284            source_root: self.source_root.clone(),
285        }));
286        tools.register(Arc::new(WebSearchTool));
287        tools.register(Arc::new(RunCommandTool::with_defaults(
288            self.work_dir.clone(),
289        )));
290
291        let system = crate::prompts::single_agent_system_prompt(
292            &self.source_root,
293            &self.work_dir,
294        );
295
296        let mut agent = AgentLoop::new(
297            Uuid::new_v4(),
298            client,
299            tools,
300            system,
301            self.agent_config.max_loop_iterations,
302        );
303
304        if let Some(ref tx) = self.event_tx {
305            agent.set_event_sink(tx.clone());
306        }
307
308        info!("Running as single agent");
309        agent.run(user_message.to_string()).await
310    }
311
312    fn emit_event(&self, event: AgentEvent) {
313        if let Some(ref tx) = self.event_tx {
314            let _ = tx.send(event);
315        }
316    }
317}
agent_sdk/agent/team.rs

agent_sdk/agent/
team.rs
