a3s_code_core/
agent_teams.rs

1//! Agent Teams — Peer-to-peer multi-agent coordination
2//!
3//! Enables multiple `AgentSession` instances to collaborate on complex tasks
4//! through a shared task board and message passing. Each agent has a role
5//! (Lead, Worker, Reviewer) and can post/claim tasks on the board.
6//!
7//! ## Architecture
8//!
9//! ```text
10//! AgentTeam
11//!   +-- TeamTaskBoard (shared task queue)
12//!   +-- TeamMember[] (role + session reference)
13//!   +-- mpsc channels (peer-to-peer messaging)
14//! ```
15//!
16//! ## Usage
17//!
18//! ```rust,no_run
19//! use a3s_code_core::agent_teams::{AgentTeam, TeamConfig, TeamRole};
20//!
21//! # async fn run() -> anyhow::Result<()> {
22//! let config = TeamConfig::default();
23//! let mut team = AgentTeam::new("refactor-auth", config);
24//!
25//! // Add members (each wraps an AgentSession)
26//! team.add_member("lead", TeamRole::Lead);
27//! team.add_member("worker-1", TeamRole::Worker);
28//! team.add_member("reviewer", TeamRole::Reviewer);
29//!
30//! // Post a task to the board
31//! team.task_board().post("Refactor auth module", "lead", None);
32//!
33//! // Worker claims and works on it
34//! let task = team.task_board().claim("worker-1");
35//! # Ok(())
36//! # }
37//! ```
38
39use std::collections::HashMap;
40use std::sync::{Arc, RwLock};
41use tokio::sync::mpsc;
42
43/// Team configuration.
44#[derive(Debug, Clone)]
45pub struct TeamConfig {
46    /// Maximum concurrent tasks on the board.
47    /// Default: 50
48    pub max_tasks: usize,
49    /// Message channel buffer size.
50    /// Default: 128
51    pub channel_buffer: usize,
52}
53
54impl Default for TeamConfig {
55    fn default() -> Self {
56        Self {
57            max_tasks: 50,
58            channel_buffer: 128,
59        }
60    }
61}
62
63/// Role of a team member.
64#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
65pub enum TeamRole {
66    /// Decomposes goals into tasks, assigns work.
67    Lead,
68    /// Executes assigned tasks.
69    Worker,
70    /// Reviews completed work, provides feedback.
71    Reviewer,
72}
73
74impl std::fmt::Display for TeamRole {
75    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
76        match self {
77            TeamRole::Lead => write!(f, "lead"),
78            TeamRole::Worker => write!(f, "worker"),
79            TeamRole::Reviewer => write!(f, "reviewer"),
80        }
81    }
82}
83
84/// A message passed between team members.
85#[derive(Debug, Clone)]
86pub struct TeamMessage {
87    /// Sender member ID.
88    pub from: String,
89    /// Recipient member ID.
90    pub to: String,
91    /// Message content.
92    pub content: String,
93    /// Optional task ID this message relates to.
94    pub task_id: Option<String>,
95    /// Timestamp (Unix epoch seconds).
96    pub timestamp: i64,
97}
98
99/// Task status on the board.
100#[derive(Debug, Clone, Copy, PartialEq, Eq)]
101pub enum TaskStatus {
102    /// Waiting to be claimed.
103    Open,
104    /// Claimed by a worker.
105    InProgress,
106    /// Work done, awaiting review.
107    InReview,
108    /// Approved by reviewer.
109    Done,
110    /// Rejected, needs rework.
111    Rejected,
112}
113
114impl std::fmt::Display for TaskStatus {
115    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
116        match self {
117            TaskStatus::Open => write!(f, "open"),
118            TaskStatus::InProgress => write!(f, "in_progress"),
119            TaskStatus::InReview => write!(f, "in_review"),
120            TaskStatus::Done => write!(f, "done"),
121            TaskStatus::Rejected => write!(f, "rejected"),
122        }
123    }
124}
125
126/// A task on the team board.
127#[derive(Debug, Clone)]
128pub struct TeamTask {
129    /// Unique task ID.
130    pub id: String,
131    /// Task description.
132    pub description: String,
133    /// Who posted it.
134    pub posted_by: String,
135    /// Who is working on it (if claimed).
136    pub assigned_to: Option<String>,
137    /// Current status.
138    pub status: TaskStatus,
139    /// Optional result/output when completed.
140    pub result: Option<String>,
141    /// Created timestamp.
142    pub created_at: i64,
143    /// Last updated timestamp.
144    pub updated_at: i64,
145}
146
147/// Shared task board for team coordination.
148#[derive(Debug)]
149pub struct TeamTaskBoard {
150    tasks: RwLock<Vec<TeamTask>>,
151    max_tasks: usize,
152    next_id: RwLock<u64>,
153}
154
155impl TeamTaskBoard {
156    /// Create a new task board.
157    pub fn new(max_tasks: usize) -> Self {
158        Self {
159            tasks: RwLock::new(Vec::new()),
160            max_tasks,
161            next_id: RwLock::new(1),
162        }
163    }
164
165    /// Post a new task to the board. Returns the task ID.
166    pub fn post(
167        &self,
168        description: &str,
169        posted_by: &str,
170        assign_to: Option<&str>,
171    ) -> Option<String> {
172        let mut tasks = self.tasks.write().unwrap();
173        if tasks.len() >= self.max_tasks {
174            return None;
175        }
176
177        let mut id_counter = self.next_id.write().unwrap();
178        let id = format!("task-{}", *id_counter);
179        *id_counter += 1;
180
181        let now = chrono::Utc::now().timestamp();
182        let status = if assign_to.is_some() {
183            TaskStatus::InProgress
184        } else {
185            TaskStatus::Open
186        };
187
188        tasks.push(TeamTask {
189            id: id.clone(),
190            description: description.to_string(),
191            posted_by: posted_by.to_string(),
192            assigned_to: assign_to.map(|s| s.to_string()),
193            status,
194            result: None,
195            created_at: now,
196            updated_at: now,
197        });
198
199        Some(id)
200    }
201
202    /// Claim the next open task for a member. Returns the task if available.
203    pub fn claim(&self, member_id: &str) -> Option<TeamTask> {
204        let mut tasks = self.tasks.write().unwrap();
205        let task = tasks.iter_mut().find(|t| t.status == TaskStatus::Open)?;
206        task.assigned_to = Some(member_id.to_string());
207        task.status = TaskStatus::InProgress;
208        task.updated_at = chrono::Utc::now().timestamp();
209        Some(task.clone())
210    }
211
212    /// Mark a task as complete with a result.
213    pub fn complete(&self, task_id: &str, result: &str) -> bool {
214        let mut tasks = self.tasks.write().unwrap();
215        if let Some(task) = tasks.iter_mut().find(|t| t.id == task_id) {
216            task.status = TaskStatus::InReview;
217            task.result = Some(result.to_string());
218            task.updated_at = chrono::Utc::now().timestamp();
219            true
220        } else {
221            false
222        }
223    }
224
225    /// Approve a task (reviewer action).
226    pub fn approve(&self, task_id: &str) -> bool {
227        let mut tasks = self.tasks.write().unwrap();
228        if let Some(task) = tasks
229            .iter_mut()
230            .find(|t| t.id == task_id && t.status == TaskStatus::InReview)
231        {
232            task.status = TaskStatus::Done;
233            task.updated_at = chrono::Utc::now().timestamp();
234            true
235        } else {
236            false
237        }
238    }
239
240    /// Reject a task back to open (reviewer action).
241    pub fn reject(&self, task_id: &str) -> bool {
242        let mut tasks = self.tasks.write().unwrap();
243        if let Some(task) = tasks
244            .iter_mut()
245            .find(|t| t.id == task_id && t.status == TaskStatus::InReview)
246        {
247            task.status = TaskStatus::Rejected;
248            task.assigned_to = None;
249            task.updated_at = chrono::Utc::now().timestamp();
250            true
251        } else {
252            false
253        }
254    }
255
256    /// Get all tasks with a given status.
257    pub fn by_status(&self, status: TaskStatus) -> Vec<TeamTask> {
258        self.tasks
259            .read()
260            .unwrap()
261            .iter()
262            .filter(|t| t.status == status)
263            .cloned()
264            .collect()
265    }
266
267    /// Get all tasks assigned to a member.
268    pub fn by_assignee(&self, member_id: &str) -> Vec<TeamTask> {
269        self.tasks
270            .read()
271            .unwrap()
272            .iter()
273            .filter(|t| t.assigned_to.as_deref() == Some(member_id))
274            .cloned()
275            .collect()
276    }
277
278    /// Get a task by ID.
279    pub fn get(&self, task_id: &str) -> Option<TeamTask> {
280        self.tasks
281            .read()
282            .unwrap()
283            .iter()
284            .find(|t| t.id == task_id)
285            .cloned()
286    }
287
288    /// Number of tasks on the board.
289    pub fn len(&self) -> usize {
290        self.tasks.read().unwrap().len()
291    }
292
293    /// Whether the board is empty.
294    pub fn is_empty(&self) -> bool {
295        self.tasks.read().unwrap().is_empty()
296    }
297
298    /// Summary stats: (open, in_progress, in_review, done, rejected).
299    pub fn stats(&self) -> (usize, usize, usize, usize, usize) {
300        let tasks = self.tasks.read().unwrap();
301        let open = tasks
302            .iter()
303            .filter(|t| t.status == TaskStatus::Open)
304            .count();
305        let progress = tasks
306            .iter()
307            .filter(|t| t.status == TaskStatus::InProgress)
308            .count();
309        let review = tasks
310            .iter()
311            .filter(|t| t.status == TaskStatus::InReview)
312            .count();
313        let done = tasks
314            .iter()
315            .filter(|t| t.status == TaskStatus::Done)
316            .count();
317        let rejected = tasks
318            .iter()
319            .filter(|t| t.status == TaskStatus::Rejected)
320            .count();
321        (open, progress, review, done, rejected)
322    }
323}
324
325/// A team member.
326#[derive(Debug, Clone)]
327pub struct TeamMember {
328    /// Unique member ID.
329    pub id: String,
330    /// Member role.
331    pub role: TeamRole,
332}
333
334/// Multi-agent team coordinator.
335pub struct AgentTeam {
336    /// Team name.
337    name: String,
338    /// Configuration.
339    config: TeamConfig,
340    /// Registered members.
341    members: HashMap<String, TeamMember>,
342    /// Shared task board.
343    task_board: Arc<TeamTaskBoard>,
344    /// Message senders per member.
345    senders: HashMap<String, mpsc::Sender<TeamMessage>>,
346    /// Message receivers per member (taken on first access).
347    receivers: HashMap<String, mpsc::Receiver<TeamMessage>>,
348}
349
350impl AgentTeam {
351    /// Create a new team.
352    pub fn new(name: &str, config: TeamConfig) -> Self {
353        Self {
354            name: name.to_string(),
355            config,
356            members: HashMap::new(),
357            task_board: Arc::new(TeamTaskBoard::new(50)),
358            senders: HashMap::new(),
359            receivers: HashMap::new(),
360        }
361    }
362
363    /// Team name.
364    pub fn name(&self) -> &str {
365        &self.name
366    }
367
368    /// Add a member to the team.
369    pub fn add_member(&mut self, id: &str, role: TeamRole) {
370        let (tx, rx) = mpsc::channel(self.config.channel_buffer);
371        self.members.insert(
372            id.to_string(),
373            TeamMember {
374                id: id.to_string(),
375                role,
376            },
377        );
378        self.senders.insert(id.to_string(), tx);
379        self.receivers.insert(id.to_string(), rx);
380    }
381
382    /// Remove a member from the team.
383    pub fn remove_member(&mut self, id: &str) -> bool {
384        self.senders.remove(id);
385        self.receivers.remove(id);
386        self.members.remove(id).is_some()
387    }
388
389    /// Get a reference to the shared task board.
390    pub fn task_board(&self) -> &TeamTaskBoard {
391        &self.task_board
392    }
393
394    /// Get a cloneable Arc to the task board.
395    pub fn task_board_arc(&self) -> Arc<TeamTaskBoard> {
396        Arc::clone(&self.task_board)
397    }
398
399    /// Send a message to a team member.
400    pub async fn send_message(
401        &self,
402        from: &str,
403        to: &str,
404        content: &str,
405        task_id: Option<&str>,
406    ) -> bool {
407        let sender = match self.senders.get(to) {
408            Some(s) => s,
409            None => return false,
410        };
411
412        let msg = TeamMessage {
413            from: from.to_string(),
414            to: to.to_string(),
415            content: content.to_string(),
416            task_id: task_id.map(|s| s.to_string()),
417            timestamp: chrono::Utc::now().timestamp(),
418        };
419
420        sender.send(msg).await.is_ok()
421    }
422
423    /// Take the message receiver for a member (can only be called once per member).
424    pub fn take_receiver(&mut self, member_id: &str) -> Option<mpsc::Receiver<TeamMessage>> {
425        self.receivers.remove(member_id)
426    }
427
428    /// Broadcast a message to all members except the sender.
429    pub async fn broadcast(&self, from: &str, content: &str, task_id: Option<&str>) {
430        for (id, sender) in &self.senders {
431            if id == from {
432                continue;
433            }
434            let msg = TeamMessage {
435                from: from.to_string(),
436                to: id.clone(),
437                content: content.to_string(),
438                task_id: task_id.map(|s| s.to_string()),
439                timestamp: chrono::Utc::now().timestamp(),
440            };
441            let _ = sender.send(msg).await;
442        }
443    }
444
445    /// Get all members.
446    pub fn members(&self) -> Vec<&TeamMember> {
447        self.members.values().collect()
448    }
449
450    /// Get members by role.
451    pub fn members_by_role(&self, role: TeamRole) -> Vec<&TeamMember> {
452        self.members.values().filter(|m| m.role == role).collect()
453    }
454
455    /// Number of members.
456    pub fn member_count(&self) -> usize {
457        self.members.len()
458    }
459}
460
461impl std::fmt::Debug for AgentTeam {
462    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
463        f.debug_struct("AgentTeam")
464            .field("name", &self.name)
465            .field("members", &self.members.len())
466            .field("tasks", &self.task_board.len())
467            .finish()
468    }
469}
470
471#[cfg(test)]
472mod tests {
473    use super::*;
474
475    #[test]
476    fn test_team_creation() {
477        let team = AgentTeam::new("test-team", TeamConfig::default());
478        assert_eq!(team.name(), "test-team");
479        assert_eq!(team.member_count(), 0);
480    }
481
482    #[test]
483    fn test_add_remove_members() {
484        let mut team = AgentTeam::new("test", TeamConfig::default());
485        team.add_member("lead", TeamRole::Lead);
486        team.add_member("w1", TeamRole::Worker);
487        team.add_member("w2", TeamRole::Worker);
488        team.add_member("rev", TeamRole::Reviewer);
489        assert_eq!(team.member_count(), 4);
490        assert_eq!(team.members_by_role(TeamRole::Worker).len(), 2);
491
492        assert!(team.remove_member("w2"));
493        assert_eq!(team.member_count(), 3);
494        assert!(!team.remove_member("nonexistent"));
495    }
496
497    #[test]
498    fn test_task_board_post_and_claim() {
499        let board = TeamTaskBoard::new(10);
500        let id = board.post("Fix auth bug", "lead", None).unwrap();
501        assert_eq!(board.len(), 1);
502
503        let task = board.claim("worker-1").unwrap();
504        assert_eq!(task.id, id);
505        assert_eq!(task.assigned_to.as_deref(), Some("worker-1"));
506        assert_eq!(task.status, TaskStatus::InProgress);
507
508        // No more open tasks
509        assert!(board.claim("worker-2").is_none());
510    }
511
512    #[test]
513    fn test_task_board_workflow() {
514        let board = TeamTaskBoard::new(10);
515        let id = board.post("Write tests", "lead", None).unwrap();
516
517        // Claim
518        board.claim("worker-1");
519
520        // Complete
521        assert!(board.complete(&id, "Added 5 tests"));
522        let task = board.get(&id).unwrap();
523        assert_eq!(task.status, TaskStatus::InReview);
524
525        // Approve
526        assert!(board.approve(&id));
527        let task = board.get(&id).unwrap();
528        assert_eq!(task.status, TaskStatus::Done);
529    }
530
531    #[test]
532    fn test_task_board_reject() {
533        let board = TeamTaskBoard::new(10);
534        let id = board.post("Refactor module", "lead", None).unwrap();
535        board.claim("worker-1");
536        board.complete(&id, "Done");
537
538        assert!(board.reject(&id));
539        let task = board.get(&id).unwrap();
540        assert_eq!(task.status, TaskStatus::Rejected);
541        assert!(task.assigned_to.is_none());
542    }
543
544    #[test]
545    fn test_task_board_max_capacity() {
546        let board = TeamTaskBoard::new(2);
547        assert!(board.post("Task 1", "lead", None).is_some());
548        assert!(board.post("Task 2", "lead", None).is_some());
549        assert!(board.post("Task 3", "lead", None).is_none()); // Full
550    }
551
552    #[test]
553    fn test_task_board_stats() {
554        let board = TeamTaskBoard::new(10);
555        board.post("T1", "lead", None);
556        board.post("T2", "lead", None);
557        let id3 = board.post("T3", "lead", Some("w1")).unwrap();
558        board.complete(&id3, "done");
559
560        let (open, progress, review, done, rejected) = board.stats();
561        assert_eq!(open, 2);
562        assert_eq!(progress, 0);
563        assert_eq!(review, 1);
564        assert_eq!(done, 0);
565        assert_eq!(rejected, 0);
566    }
567
568    #[test]
569    fn test_task_board_by_assignee() {
570        let board = TeamTaskBoard::new(10);
571        board.post("T1", "lead", Some("w1"));
572        board.post("T2", "lead", Some("w2"));
573        board.post("T3", "lead", Some("w1"));
574
575        let w1_tasks = board.by_assignee("w1");
576        assert_eq!(w1_tasks.len(), 2);
577    }
578
579    #[tokio::test]
580    async fn test_send_message() {
581        let mut team = AgentTeam::new("msg-test", TeamConfig::default());
582        team.add_member("lead", TeamRole::Lead);
583        team.add_member("worker", TeamRole::Worker);
584
585        let mut rx = team.take_receiver("worker").unwrap();
586
587        assert!(
588            team.send_message("lead", "worker", "Please fix the bug", Some("task-1"))
589                .await
590        );
591
592        let msg = rx.recv().await.unwrap();
593        assert_eq!(msg.from, "lead");
594        assert_eq!(msg.to, "worker");
595        assert_eq!(msg.content, "Please fix the bug");
596        assert_eq!(msg.task_id.as_deref(), Some("task-1"));
597    }
598
599    #[tokio::test]
600    async fn test_broadcast() {
601        let mut team = AgentTeam::new("broadcast-test", TeamConfig::default());
602        team.add_member("lead", TeamRole::Lead);
603        team.add_member("w1", TeamRole::Worker);
604        team.add_member("w2", TeamRole::Worker);
605
606        let mut rx1 = team.take_receiver("w1").unwrap();
607        let mut rx2 = team.take_receiver("w2").unwrap();
608
609        team.broadcast("lead", "New task available", None).await;
610
611        let m1 = rx1.recv().await.unwrap();
612        let m2 = rx2.recv().await.unwrap();
613        assert_eq!(m1.content, "New task available");
614        assert_eq!(m2.content, "New task available");
615    }
616
617    #[test]
618    fn test_role_display() {
619        assert_eq!(TeamRole::Lead.to_string(), "lead");
620        assert_eq!(TeamRole::Worker.to_string(), "worker");
621        assert_eq!(TeamRole::Reviewer.to_string(), "reviewer");
622    }
623
624    #[test]
625    fn test_task_status_display() {
626        assert_eq!(TaskStatus::Open.to_string(), "open");
627        assert_eq!(TaskStatus::InProgress.to_string(), "in_progress");
628        assert_eq!(TaskStatus::InReview.to_string(), "in_review");
629        assert_eq!(TaskStatus::Done.to_string(), "done");
630        assert_eq!(TaskStatus::Rejected.to_string(), "rejected");
631    }
632}
a3s_code_core/agent_teams.rs

a3s_code_core/
agent_teams.rs
