Skip to main content

agent_teams/messaging/
structured.rs

1//! Convenience constructors for [`StructuredMessage`] variants.
2//!
3//! These helpers make it easy to create well-formed structured messages
4//! without manually constructing enum variants with `Option` fields.
5
6use crate::models::StructuredMessage;
7
8/// Create a task assignment message.
9pub fn task_assignment(
10    task_id: impl Into<String>,
11    subject: impl Into<String>,
12    description: impl Into<String>,
13) -> StructuredMessage {
14    StructuredMessage::TaskAssignment {
15        task_id: task_id.into(),
16        subject: subject.into(),
17        description: Some(description.into()),
18        assigned_by: None,
19        timestamp: None,
20    }
21}
22
23/// Create a shutdown request message.
24pub fn shutdown_request(
25    request_id: impl Into<String>,
26    reason: impl Into<String>,
27) -> StructuredMessage {
28    StructuredMessage::ShutdownRequest {
29        request_id: Some(request_id.into()),
30        reason: Some(reason.into()),
31    }
32}
33
34/// Create a shutdown approved response.
35pub fn shutdown_approved(request_id: impl Into<String>) -> StructuredMessage {
36    StructuredMessage::ShutdownApproved {
37        request_id: Some(request_id.into()),
38    }
39}
40
41/// Create an idle notification from an agent.
42pub fn idle_notification(
43    agent: impl Into<String>,
44    last_task_id: Option<String>,
45) -> StructuredMessage {
46    StructuredMessage::IdleNotification {
47        agent: agent.into(),
48        last_task_id,
49        idle_reason: None,
50        timestamp: None,
51    }
52}
53
54/// Create a plan approval request.
55pub fn plan_approval_request(
56    request_id: impl Into<String>,
57    plan: impl Into<String>,
58) -> StructuredMessage {
59    StructuredMessage::PlanApprovalRequest {
60        request_id: Some(request_id.into()),
61        plan: plan.into(),
62    }
63}
64
65/// Create a plan approval response.
66pub fn plan_approval_response(
67    request_id: impl Into<String>,
68    approved: bool,
69    feedback: Option<String>,
70) -> StructuredMessage {
71    StructuredMessage::PlanApprovalResponse {
72        request_id: Some(request_id.into()),
73        approved,
74        feedback,
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use super::*;
81
82    #[test]
83    fn task_assignment_constructor() {
84        let msg = task_assignment("1", "Fix bug", "Auth is broken");
85        match msg {
86            StructuredMessage::TaskAssignment {
87                task_id,
88                subject,
89                description,
90                ..
91            } => {
92                assert_eq!(task_id, "1");
93                assert_eq!(subject, "Fix bug");
94                assert_eq!(description.unwrap(), "Auth is broken");
95            }
96            _ => panic!("wrong variant"),
97        }
98    }
99
100    #[test]
101    fn shutdown_request_constructor() {
102        let msg = shutdown_request("req-1", "All tasks done");
103        match msg {
104            StructuredMessage::ShutdownRequest {
105                request_id,
106                reason,
107            } => {
108                assert_eq!(request_id.unwrap(), "req-1");
109                assert_eq!(reason.unwrap(), "All tasks done");
110            }
111            _ => panic!("wrong variant"),
112        }
113    }
114
115    #[test]
116    fn shutdown_approved_constructor() {
117        let msg = shutdown_approved("req-1");
118        match msg {
119            StructuredMessage::ShutdownApproved { request_id } => {
120                assert_eq!(request_id.unwrap(), "req-1");
121            }
122            _ => panic!("wrong variant"),
123        }
124    }
125
126    #[test]
127    fn idle_notification_constructor() {
128        let msg = idle_notification("worker-1", Some("42".into()));
129        match msg {
130            StructuredMessage::IdleNotification {
131                agent,
132                last_task_id,
133                ..
134            } => {
135                assert_eq!(agent, "worker-1");
136                assert_eq!(last_task_id.unwrap(), "42");
137            }
138            _ => panic!("wrong variant"),
139        }
140    }
141
142    #[test]
143    fn idle_notification_no_task() {
144        let msg = idle_notification("worker-1", None);
145        match msg {
146            StructuredMessage::IdleNotification { last_task_id, .. } => {
147                assert!(last_task_id.is_none());
148            }
149            _ => panic!("wrong variant"),
150        }
151    }
152
153    #[test]
154    fn plan_approval_request_constructor() {
155        let msg = plan_approval_request("req-2", "Step 1: read code\nStep 2: fix bug");
156        match msg {
157            StructuredMessage::PlanApprovalRequest { request_id, plan } => {
158                assert_eq!(request_id.unwrap(), "req-2");
159                assert!(plan.contains("Step 1"));
160            }
161            _ => panic!("wrong variant"),
162        }
163    }
164
165    #[test]
166    fn plan_approval_response_approved() {
167        let msg = plan_approval_response("req-2", true, None);
168        match msg {
169            StructuredMessage::PlanApprovalResponse {
170                request_id,
171                approved,
172                feedback,
173            } => {
174                assert_eq!(request_id.unwrap(), "req-2");
175                assert!(approved);
176                assert!(feedback.is_none());
177            }
178            _ => panic!("wrong variant"),
179        }
180    }
181
182    #[test]
183    fn plan_approval_response_rejected_with_feedback() {
184        let msg = plan_approval_response("req-2", false, Some("Add error handling".into()));
185        match msg {
186            StructuredMessage::PlanApprovalResponse {
187                approved, feedback, ..
188            } => {
189                assert!(!approved);
190                assert_eq!(feedback.unwrap(), "Add error handling");
191            }
192            _ => panic!("wrong variant"),
193        }
194    }
195
196    #[test]
197    fn structured_messages_serialize_correctly() {
198        let messages: Vec<StructuredMessage> = vec![
199            task_assignment("1", "Do thing", "Details"),
200            shutdown_request("r1", "Done"),
201            shutdown_approved("r1"),
202            idle_notification("w1", None),
203            plan_approval_request("r2", "The plan"),
204            plan_approval_response("r2", true, None),
205        ];
206
207        for msg in &messages {
208            let json = serde_json::to_string(msg).unwrap();
209            let parsed: StructuredMessage = serde_json::from_str(&json).unwrap();
210            // Verify round-trip doesn't panic and summary works.
211            let _ = parsed.summary();
212        }
213    }
214}