Skip to main content

punch_types/
troop.rs

1//! # Troop Types
2//!
3//! Shared types for the multi-agent troop coordination system.
4//! Troops are named groups of coordinated fighters (agents) that
5//! work together using various coordination strategies.
6
7use chrono::{DateTime, Utc};
8use serde::{Deserialize, Serialize};
9use uuid::Uuid;
10
11use crate::fighter::FighterId;
12
13/// Unique identifier for a Troop (coordinated agent group).
14#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
15#[serde(transparent)]
16pub struct TroopId(pub Uuid);
17
18impl TroopId {
19    pub fn new() -> Self {
20        Self(Uuid::new_v4())
21    }
22}
23
24impl Default for TroopId {
25    fn default() -> Self {
26        Self::new()
27    }
28}
29
30impl std::fmt::Display for TroopId {
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        write!(f, "{}", self.0)
33    }
34}
35
36/// Coordination strategy determining how a troop distributes work.
37#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
38#[serde(rename_all = "snake_case")]
39pub enum CoordinationStrategy {
40    /// Leader delegates tasks, workers execute and report back.
41    LeaderWorker,
42    /// Tasks distributed evenly across members.
43    RoundRobin,
44    /// All members receive same task, results aggregated.
45    Broadcast,
46    /// Each member processes output of previous member.
47    Pipeline,
48    /// Members vote on decisions, majority wins.
49    Consensus,
50    /// Tasks routed to member with matching capabilities.
51    Specialist,
52}
53
54impl std::fmt::Display for CoordinationStrategy {
55    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
56        match self {
57            Self::LeaderWorker => write!(f, "leader_worker"),
58            Self::RoundRobin => write!(f, "round_robin"),
59            Self::Broadcast => write!(f, "broadcast"),
60            Self::Pipeline => write!(f, "pipeline"),
61            Self::Consensus => write!(f, "consensus"),
62            Self::Specialist => write!(f, "specialist"),
63        }
64    }
65}
66
67/// Current operational status of a Troop.
68#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
69#[serde(rename_all = "snake_case")]
70pub enum TroopStatus {
71    /// Troop is being assembled, not yet operational.
72    Forming,
73    /// Troop is active and ready to receive tasks.
74    Active,
75    /// Troop is temporarily paused.
76    Paused,
77    /// Troop has been dissolved.
78    Disbanded,
79}
80
81impl std::fmt::Display for TroopStatus {
82    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
83        match self {
84            Self::Forming => write!(f, "forming"),
85            Self::Active => write!(f, "active"),
86            Self::Paused => write!(f, "paused"),
87            Self::Disbanded => write!(f, "disbanded"),
88        }
89    }
90}
91
92/// A named group of coordinated agents.
93#[derive(Debug, Clone, Serialize, Deserialize)]
94pub struct Troop {
95    /// Unique identifier for this troop.
96    pub id: TroopId,
97    /// Human-readable name for this troop.
98    pub name: String,
99    /// The leader fighter who coordinates the troop.
100    pub leader: FighterId,
101    /// All member fighters (including the leader).
102    pub members: Vec<FighterId>,
103    /// How tasks are distributed among members.
104    pub strategy: CoordinationStrategy,
105    /// Current operational status.
106    pub status: TroopStatus,
107    /// When the troop was formed.
108    pub created_at: DateTime<Utc>,
109}
110
111/// Status of a single subtask within a swarm task.
112#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
113#[serde(rename_all = "snake_case")]
114pub enum SubtaskStatus {
115    /// Waiting to be assigned.
116    Pending,
117    /// Assigned to a fighter and running.
118    Running,
119    /// Completed successfully.
120    Completed,
121    /// Failed with an error.
122    Failed(String),
123}
124
125/// A subtask within a larger swarm task.
126#[derive(Debug, Clone, Serialize, Deserialize)]
127pub struct SwarmSubtask {
128    /// Unique identifier for this subtask.
129    pub id: Uuid,
130    /// Description of what this subtask should accomplish.
131    pub description: String,
132    /// The fighter assigned to this subtask, if any.
133    pub assigned_to: Option<FighterId>,
134    /// Current status.
135    pub status: SubtaskStatus,
136    /// Result content, if completed.
137    pub result: Option<String>,
138    /// Dependencies: IDs of subtasks that must complete first.
139    pub depends_on: Vec<Uuid>,
140}
141
142/// A complex task decomposed into subtasks for swarm execution.
143#[derive(Debug, Clone, Serialize, Deserialize)]
144pub struct SwarmTask {
145    /// Unique identifier.
146    pub id: Uuid,
147    /// The original task description.
148    pub description: String,
149    /// Decomposed subtasks.
150    pub subtasks: Vec<SwarmSubtask>,
151    /// Overall progress (0.0 to 1.0).
152    pub progress: f64,
153    /// When the swarm task was created.
154    pub created_at: DateTime<Utc>,
155    /// Aggregated result, if all subtasks completed.
156    pub aggregated_result: Option<String>,
157}
158
159/// Priority levels for inter-agent messages.
160#[derive(
161    Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize,
162)]
163#[serde(rename_all = "snake_case")]
164pub enum MessagePriority {
165    /// Low priority, can be deferred.
166    Low,
167    /// Normal priority (default).
168    #[default]
169    Normal,
170    /// High priority, process promptly.
171    High,
172    /// Critical, process immediately.
173    Critical,
174}
175
176/// Types of channels for inter-agent messaging.
177#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
178#[serde(rename_all = "snake_case")]
179pub enum MessageChannel {
180    /// Point-to-point between two fighters.
181    Direct,
182    /// One-to-all in a troop.
183    Broadcast,
184    /// One-to-some (subset of troop).
185    Multicast(Vec<FighterId>),
186    /// Send and wait for response (with timeout in milliseconds).
187    Request { timeout_ms: u64 },
188    /// Continuous data flow between agents.
189    Stream,
190}
191
192/// Types of messages exchanged between agents.
193#[derive(Debug, Clone, Serialize, Deserialize)]
194#[serde(rename_all = "snake_case", tag = "type")]
195pub enum AgentMessageType {
196    /// Assign work to a fighter.
197    TaskAssignment { task: String },
198    /// Report task completion.
199    TaskResult { result: String, success: bool },
200    /// Heartbeat / progress update.
201    StatusUpdate { progress: f64, detail: String },
202    /// Share context/knowledge between agents.
203    DataShare {
204        key: String,
205        value: serde_json::Value,
206    },
207    /// Request a vote from peers.
208    VoteRequest {
209        proposal: String,
210        options: Vec<String>,
211    },
212    /// Respond to a vote request.
213    VoteResponse { proposal: String, vote: String },
214    /// Task escalation to leader.
215    Escalation {
216        reason: String,
217        original_task: String,
218    },
219}
220
221/// An inter-agent message.
222#[derive(Debug, Clone, Serialize, Deserialize)]
223pub struct AgentMessage {
224    /// Unique message ID.
225    pub id: Uuid,
226    /// The sending fighter.
227    pub from: FighterId,
228    /// The receiving fighter.
229    pub to: FighterId,
230    /// Channel type.
231    pub channel: MessageChannel,
232    /// Message content.
233    pub content: AgentMessageType,
234    /// Priority level.
235    pub priority: MessagePriority,
236    /// When the message was sent.
237    pub timestamp: DateTime<Utc>,
238    /// Whether the message has been delivered.
239    pub delivered: bool,
240}
241
242/// Restart strategy for the Supervisor pattern.
243#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
244#[serde(rename_all = "snake_case")]
245pub enum RestartStrategy {
246    /// Only restart the failed worker.
247    OneForOne,
248    /// Restart all workers if one fails.
249    AllForOne,
250}
251
252/// A bid from an agent in the Auction pattern.
253#[derive(Debug, Clone, Serialize, Deserialize)]
254pub struct AuctionBid {
255    /// The bidding fighter.
256    pub fighter_id: FighterId,
257    /// Estimated time to complete (in seconds).
258    pub estimated_time_secs: u64,
259    /// Confidence level (0.0 to 1.0).
260    pub confidence: f64,
261    /// When the bid was submitted.
262    pub submitted_at: DateTime<Utc>,
263}
264
265/// Selection criteria for the Scatter-Gather pattern.
266#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
267#[serde(rename_all = "snake_case")]
268pub enum SelectionCriteria {
269    /// Select the fastest response.
270    Fastest,
271    /// Select based on highest reported quality.
272    HighestQuality,
273    /// Select based on consensus among responses.
274    Consensus,
275}
276
277#[cfg(test)]
278mod tests {
279    use super::*;
280
281    #[test]
282    fn test_troop_id_display() {
283        let uuid = Uuid::nil();
284        let id = TroopId(uuid);
285        assert_eq!(id.to_string(), uuid.to_string());
286    }
287
288    #[test]
289    fn test_troop_id_new_is_unique() {
290        let id1 = TroopId::new();
291        let id2 = TroopId::new();
292        assert_ne!(id1, id2);
293    }
294
295    #[test]
296    fn test_troop_id_default() {
297        let id = TroopId::default();
298        assert_ne!(id.0, Uuid::nil());
299    }
300
301    #[test]
302    fn test_troop_id_serde_transparent() {
303        let uuid = Uuid::new_v4();
304        let id = TroopId(uuid);
305        let json = serde_json::to_string(&id).expect("serialize");
306        assert_eq!(json, format!("\"{}\"", uuid));
307        let deser: TroopId = serde_json::from_str(&json).expect("deserialize");
308        assert_eq!(deser, id);
309    }
310
311    #[test]
312    fn test_troop_id_copy_clone() {
313        let id = TroopId::new();
314        let copied = id;
315        let cloned = id.clone();
316        assert_eq!(id, copied);
317        assert_eq!(id, cloned);
318    }
319
320    #[test]
321    fn test_troop_id_hash() {
322        let id = TroopId::new();
323        let mut set = std::collections::HashSet::new();
324        set.insert(id);
325        set.insert(id);
326        assert_eq!(set.len(), 1);
327    }
328
329    #[test]
330    fn test_coordination_strategy_display() {
331        assert_eq!(
332            CoordinationStrategy::LeaderWorker.to_string(),
333            "leader_worker"
334        );
335        assert_eq!(CoordinationStrategy::RoundRobin.to_string(), "round_robin");
336        assert_eq!(CoordinationStrategy::Broadcast.to_string(), "broadcast");
337        assert_eq!(CoordinationStrategy::Pipeline.to_string(), "pipeline");
338        assert_eq!(CoordinationStrategy::Consensus.to_string(), "consensus");
339        assert_eq!(CoordinationStrategy::Specialist.to_string(), "specialist");
340    }
341
342    #[test]
343    fn test_coordination_strategy_serde_roundtrip() {
344        let strategies = vec![
345            CoordinationStrategy::LeaderWorker,
346            CoordinationStrategy::RoundRobin,
347            CoordinationStrategy::Broadcast,
348            CoordinationStrategy::Pipeline,
349            CoordinationStrategy::Consensus,
350            CoordinationStrategy::Specialist,
351        ];
352        for strategy in &strategies {
353            let json = serde_json::to_string(strategy).expect("serialize");
354            let deser: CoordinationStrategy = serde_json::from_str(&json).expect("deserialize");
355            assert_eq!(&deser, strategy);
356        }
357    }
358
359    #[test]
360    fn test_troop_status_display() {
361        assert_eq!(TroopStatus::Forming.to_string(), "forming");
362        assert_eq!(TroopStatus::Active.to_string(), "active");
363        assert_eq!(TroopStatus::Paused.to_string(), "paused");
364        assert_eq!(TroopStatus::Disbanded.to_string(), "disbanded");
365    }
366
367    #[test]
368    fn test_troop_status_serde_roundtrip() {
369        let statuses = vec![
370            TroopStatus::Forming,
371            TroopStatus::Active,
372            TroopStatus::Paused,
373            TroopStatus::Disbanded,
374        ];
375        for status in &statuses {
376            let json = serde_json::to_string(status).expect("serialize");
377            let deser: TroopStatus = serde_json::from_str(&json).expect("deserialize");
378            assert_eq!(&deser, status);
379        }
380    }
381
382    #[test]
383    fn test_troop_serde_roundtrip() {
384        let troop = Troop {
385            id: TroopId::new(),
386            name: "Alpha Squad".to_string(),
387            leader: FighterId::new(),
388            members: vec![FighterId::new(), FighterId::new()],
389            strategy: CoordinationStrategy::LeaderWorker,
390            status: TroopStatus::Active,
391            created_at: Utc::now(),
392        };
393        let json = serde_json::to_string(&troop).expect("serialize");
394        let deser: Troop = serde_json::from_str(&json).expect("deserialize");
395        assert_eq!(deser.id, troop.id);
396        assert_eq!(deser.name, "Alpha Squad");
397        assert_eq!(deser.members.len(), 2);
398    }
399
400    #[test]
401    fn test_message_priority_default() {
402        assert_eq!(MessagePriority::default(), MessagePriority::Normal);
403    }
404
405    #[test]
406    fn test_message_priority_ordering() {
407        assert!(MessagePriority::Low < MessagePriority::Normal);
408        assert!(MessagePriority::Normal < MessagePriority::High);
409        assert!(MessagePriority::High < MessagePriority::Critical);
410    }
411
412    #[test]
413    fn test_message_channel_serde() {
414        let channels = vec![
415            MessageChannel::Direct,
416            MessageChannel::Broadcast,
417            MessageChannel::Multicast(vec![FighterId::new()]),
418            MessageChannel::Request { timeout_ms: 5000 },
419            MessageChannel::Stream,
420        ];
421        for channel in &channels {
422            let json = serde_json::to_string(channel).expect("serialize");
423            let deser: MessageChannel = serde_json::from_str(&json).expect("deserialize");
424            assert_eq!(&deser, channel);
425        }
426    }
427
428    #[test]
429    fn test_agent_message_type_task_assignment() {
430        let msg = AgentMessageType::TaskAssignment {
431            task: "analyze code".to_string(),
432        };
433        let json = serde_json::to_string(&msg).expect("serialize");
434        assert!(json.contains("task_assignment"));
435        let deser: AgentMessageType = serde_json::from_str(&json).expect("deserialize");
436        match deser {
437            AgentMessageType::TaskAssignment { task } => {
438                assert_eq!(task, "analyze code");
439            }
440            _ => panic!("wrong variant"),
441        }
442    }
443
444    #[test]
445    fn test_agent_message_type_vote_request() {
446        let msg = AgentMessageType::VoteRequest {
447            proposal: "merge PR?".to_string(),
448            options: vec!["yes".to_string(), "no".to_string()],
449        };
450        let json = serde_json::to_string(&msg).expect("serialize");
451        let deser: AgentMessageType = serde_json::from_str(&json).expect("deserialize");
452        match deser {
453            AgentMessageType::VoteRequest { proposal, options } => {
454                assert_eq!(proposal, "merge PR?");
455                assert_eq!(options.len(), 2);
456            }
457            _ => panic!("wrong variant"),
458        }
459    }
460
461    #[test]
462    fn test_agent_message_serde() {
463        let msg = AgentMessage {
464            id: Uuid::new_v4(),
465            from: FighterId::new(),
466            to: FighterId::new(),
467            channel: MessageChannel::Direct,
468            content: AgentMessageType::StatusUpdate {
469                progress: 0.5,
470                detail: "halfway done".to_string(),
471            },
472            priority: MessagePriority::Normal,
473            timestamp: Utc::now(),
474            delivered: false,
475        };
476        let json = serde_json::to_string(&msg).expect("serialize");
477        let deser: AgentMessage = serde_json::from_str(&json).expect("deserialize");
478        assert_eq!(deser.id, msg.id);
479        assert!(!deser.delivered);
480    }
481
482    #[test]
483    fn test_subtask_status_serde() {
484        let statuses = vec![
485            SubtaskStatus::Pending,
486            SubtaskStatus::Running,
487            SubtaskStatus::Completed,
488            SubtaskStatus::Failed("error".to_string()),
489        ];
490        for status in &statuses {
491            let json = serde_json::to_string(status).expect("serialize");
492            let deser: SubtaskStatus = serde_json::from_str(&json).expect("deserialize");
493            assert_eq!(&deser, status);
494        }
495    }
496
497    #[test]
498    fn test_swarm_task_progress() {
499        let task = SwarmTask {
500            id: Uuid::new_v4(),
501            description: "big task".to_string(),
502            subtasks: vec![],
503            progress: 0.75,
504            created_at: Utc::now(),
505            aggregated_result: None,
506        };
507        assert!((task.progress - 0.75).abs() < f64::EPSILON);
508    }
509
510    #[test]
511    fn test_auction_bid_serde() {
512        let bid = AuctionBid {
513            fighter_id: FighterId::new(),
514            estimated_time_secs: 30,
515            confidence: 0.9,
516            submitted_at: Utc::now(),
517        };
518        let json = serde_json::to_string(&bid).expect("serialize");
519        let deser: AuctionBid = serde_json::from_str(&json).expect("deserialize");
520        assert_eq!(deser.estimated_time_secs, 30);
521        assert!((deser.confidence - 0.9).abs() < f64::EPSILON);
522    }
523
524    #[test]
525    fn test_selection_criteria_serde() {
526        let criteria = vec![
527            SelectionCriteria::Fastest,
528            SelectionCriteria::HighestQuality,
529            SelectionCriteria::Consensus,
530        ];
531        for c in &criteria {
532            let json = serde_json::to_string(c).expect("serialize");
533            let deser: SelectionCriteria = serde_json::from_str(&json).expect("deserialize");
534            assert_eq!(&deser, c);
535        }
536    }
537
538    #[test]
539    fn test_restart_strategy_serde() {
540        let strategies = vec![RestartStrategy::OneForOne, RestartStrategy::AllForOne];
541        for s in &strategies {
542            let json = serde_json::to_string(s).expect("serialize");
543            let deser: RestartStrategy = serde_json::from_str(&json).expect("deserialize");
544            assert_eq!(&deser, s);
545        }
546    }
547
548    #[test]
549    fn test_escalation_message() {
550        let msg = AgentMessageType::Escalation {
551            reason: "too complex".to_string(),
552            original_task: "analyze codebase".to_string(),
553        };
554        let json = serde_json::to_string(&msg).expect("serialize");
555        let deser: AgentMessageType = serde_json::from_str(&json).expect("deserialize");
556        match deser {
557            AgentMessageType::Escalation {
558                reason,
559                original_task,
560            } => {
561                assert_eq!(reason, "too complex");
562                assert_eq!(original_task, "analyze codebase");
563            }
564            _ => panic!("wrong variant"),
565        }
566    }
567
568    #[test]
569    fn test_data_share_message() {
570        let msg = AgentMessageType::DataShare {
571            key: "context".to_string(),
572            value: serde_json::json!({"files": ["main.rs"]}),
573        };
574        let json = serde_json::to_string(&msg).expect("serialize");
575        let deser: AgentMessageType = serde_json::from_str(&json).expect("deserialize");
576        match deser {
577            AgentMessageType::DataShare { key, value } => {
578                assert_eq!(key, "context");
579                assert!(value.get("files").is_some());
580            }
581            _ => panic!("wrong variant"),
582        }
583    }
584
585    #[test]
586    fn test_swarm_subtask_with_dependencies() {
587        let dep_id = Uuid::new_v4();
588        let subtask = SwarmSubtask {
589            id: Uuid::new_v4(),
590            description: "step 2".to_string(),
591            assigned_to: Some(FighterId::new()),
592            status: SubtaskStatus::Pending,
593            result: None,
594            depends_on: vec![dep_id],
595        };
596        let json = serde_json::to_string(&subtask).expect("serialize");
597        let deser: SwarmSubtask = serde_json::from_str(&json).expect("deserialize");
598        assert_eq!(deser.depends_on.len(), 1);
599        assert_eq!(deser.depends_on[0], dep_id);
600    }
601}