1use acp_utils::notifications::{
2 SubAgentEvent, SubAgentToolCallUpdate, SubAgentToolError, SubAgentToolRequest, SubAgentToolResult,
3};
4use llm::{ToolCallError, ToolCallRequest, ToolCallResult};
5use mcp_utils::display_meta::ToolResultMeta;
6use serde::{Deserialize, Serialize};
7
8#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
10#[serde(tag = "type", rename_all = "snake_case")]
11pub enum AgentMessage {
12 Text {
13 message_id: String,
14 chunk: String,
15 is_complete: bool,
16 model_name: String,
17 },
18
19 Thought {
20 message_id: String,
21 chunk: String,
22 is_complete: bool,
23 model_name: String,
24 },
25
26 ToolCall {
27 request: ToolCallRequest,
28 model_name: String,
29 },
30
31 ToolCallUpdate {
32 tool_call_id: String,
33 chunk: String,
34 model_name: String,
35 },
36
37 ToolProgress {
38 request: ToolCallRequest,
39 progress: f64,
40 total: Option<f64>,
41 message: Option<String>,
42 },
43
44 ToolResult {
45 result: ToolCallResult,
46 result_meta: Option<ToolResultMeta>,
47 model_name: String,
48 },
49
50 ToolError {
51 error: ToolCallError,
52 model_name: String,
53 },
54
55 Error {
56 message: String,
57 },
58
59 Cancelled {
60 message: String,
61 },
62
63 ContextCompactionStarted {
65 message_count: usize,
66 },
67
68 ContextCompactionResult {
70 summary: String,
71 messages_removed: usize,
72 },
73
74 #[serde(rename = "context_usage")]
76 ContextUsageUpdate {
77 usage_ratio: Option<f64>,
79 context_limit: Option<u32>,
81 input_tokens: u32,
83 output_tokens: u32,
85 cache_read_tokens: Option<u32>,
87 cache_creation_tokens: Option<u32>,
89 reasoning_tokens: Option<u32>,
91 total_input_tokens: u64,
93 total_output_tokens: u64,
95 total_cache_read_tokens: u64,
97 total_cache_creation_tokens: u64,
99 total_reasoning_tokens: u64,
101 },
102
103 AutoContinue {
105 attempt: u32,
107 max_attempts: u32,
109 },
110
111 Retrying {
113 attempt: u32,
115 max_attempts: u32,
117 delay_ms: u64,
119 error: String,
121 },
122
123 ModelSwitched {
125 previous: String,
126 new: String,
127 },
128
129 ContextCleared,
131
132 Done,
133}
134
135impl From<&AgentMessage> for SubAgentEvent {
136 fn from(msg: &AgentMessage) -> Self {
137 match msg {
138 AgentMessage::ToolCall { request, .. } => SubAgentEvent::ToolCall {
139 request: SubAgentToolRequest {
140 id: request.id.clone(),
141 name: request.name.clone(),
142 arguments: request.arguments.clone(),
143 },
144 },
145 AgentMessage::ToolCallUpdate { tool_call_id, chunk, .. } => SubAgentEvent::ToolCallUpdate {
146 update: SubAgentToolCallUpdate { id: tool_call_id.clone(), chunk: chunk.clone() },
147 },
148 AgentMessage::ToolResult { result, result_meta, .. } => SubAgentEvent::ToolResult {
149 result: SubAgentToolResult {
150 id: result.id.clone(),
151 name: result.name.clone(),
152 result_meta: result_meta.clone(),
153 },
154 },
155 AgentMessage::ToolError { error, .. } => {
156 SubAgentEvent::ToolError { error: SubAgentToolError { id: error.id.clone(), name: error.name.clone() } }
157 }
158 AgentMessage::Done => SubAgentEvent::Done,
159 _ => SubAgentEvent::Other,
160 }
161 }
162}
163
164impl AgentMessage {
165 pub fn text(message_id: &str, chunk: &str, is_complete: bool, model_name: &str) -> Self {
166 AgentMessage::Text {
167 message_id: message_id.to_string(),
168 chunk: chunk.to_string(),
169 is_complete,
170 model_name: model_name.to_string(),
171 }
172 }
173
174 pub fn thought(message_id: &str, chunk: &str, is_complete: bool, model_name: &str) -> Self {
175 AgentMessage::Thought {
176 message_id: message_id.to_string(),
177 chunk: chunk.to_string(),
178 is_complete,
179 model_name: model_name.to_string(),
180 }
181 }
182}
183
184#[cfg(test)]
185mod tests {
186 use super::AgentMessage;
187 use acp_utils::notifications::SubAgentEvent;
188 use llm::ToolCallResult;
189 use mcp_utils::display_meta::ToolDisplayMeta;
190
191 #[test]
192 fn test_model_switched_serde_roundtrip() {
193 let msg = AgentMessage::ModelSwitched {
194 previous: "anthropic:claude-3.5-sonnet".to_string(),
195 new: "ollama:llama3.2".to_string(),
196 };
197 let json = serde_json::to_string(&msg).unwrap();
198 let parsed: AgentMessage = serde_json::from_str(&json).unwrap();
199 assert_eq!(parsed, msg);
200 }
201
202 #[test]
203 fn test_thought_serde_roundtrip() {
204 let msg = AgentMessage::Thought {
205 message_id: "msg_1".to_string(),
206 chunk: "thinking".to_string(),
207 is_complete: false,
208 model_name: "test-model".to_string(),
209 };
210 let json = serde_json::to_string(&msg).unwrap();
211 let parsed: AgentMessage = serde_json::from_str(&json).unwrap();
212 assert_eq!(parsed, msg);
213 }
214
215 #[test]
216 fn test_thought_complete_serde_roundtrip() {
217 let msg = AgentMessage::Thought {
218 message_id: "msg_1".to_string(),
219 chunk: "full reasoning".to_string(),
220 is_complete: true,
221 model_name: "test-model".to_string(),
222 };
223 let json = serde_json::to_string(&msg).unwrap();
224 let parsed: AgentMessage = serde_json::from_str(&json).unwrap();
225 assert_eq!(parsed, msg);
226 }
227
228 #[test]
229 fn test_tool_result_serializes_result_meta() {
230 let msg = AgentMessage::ToolResult {
231 result: ToolCallResult {
232 id: "call_1".to_string(),
233 name: "coding__read_file".to_string(),
234 arguments: r#"{"filePath":"Cargo.toml"}"#.to_string(),
235 result: "ok".to_string(),
236 },
237 result_meta: Some(ToolDisplayMeta::new("Read file", "Cargo.toml, 156 lines").into()),
238 model_name: "test-model".to_string(),
239 };
240
241 let json = serde_json::to_value(&msg).unwrap();
242 assert_eq!(json["type"], "tool_result");
243 assert_eq!(json["result_meta"]["display"]["title"], "Read file");
244 assert_eq!(json["result_meta"]["display"]["value"], "Cargo.toml, 156 lines");
245
246 let parsed: AgentMessage = serde_json::from_value(json).unwrap();
247 assert_eq!(parsed, msg);
248 }
249
250 #[test]
251 fn test_sub_agent_tool_result_includes_display_fields() {
252 let msg = AgentMessage::ToolResult {
253 result: ToolCallResult {
254 id: "call_1".to_string(),
255 name: "coding__read_file".to_string(),
256 arguments: r#"{"filePath":"Cargo.toml"}"#.to_string(),
257 result: "ok".to_string(),
258 },
259 result_meta: Some(ToolDisplayMeta::new("Read file", "Cargo.toml, 156 lines").into()),
260 model_name: "test-model".to_string(),
261 };
262
263 let event: SubAgentEvent = (&msg).into();
264 match event {
265 SubAgentEvent::ToolResult { result } => {
266 assert_eq!(result.id, "call_1");
267 assert_eq!(result.name, "coding__read_file");
268 let result_meta = result.result_meta.expect("result_meta should be present");
269 assert_eq!(result_meta.display.title, "Read file");
270 assert_eq!(result_meta.display.value, "Cargo.toml, 156 lines");
271 }
272 other => panic!("Expected ToolResult, got {other:?}"),
273 }
274 }
275
276 #[test]
277 fn test_sub_agent_tool_call_update_includes_updated_fields() {
278 let msg = AgentMessage::ToolCallUpdate {
279 tool_call_id: "call_1".to_string(),
280 chunk: r#"{"filePath":"Cargo.toml"}"#.to_string(),
281 model_name: "test-model".to_string(),
282 };
283
284 let event: SubAgentEvent = (&msg).into();
285 match event {
286 SubAgentEvent::ToolCallUpdate { update } => {
287 assert_eq!(update.id, "call_1");
288 assert_eq!(update.chunk, r#"{"filePath":"Cargo.toml"}"#);
289 }
290 other => panic!("Expected ToolCallUpdate, got {other:?}"),
291 }
292 }
293
294 #[test]
295 fn test_done_serializes_as_object() {
296 let json = serde_json::to_value(&AgentMessage::Done).unwrap();
297 assert_eq!(json["type"], "done");
298 assert_eq!(json.as_object().unwrap().len(), 1);
299
300 let parsed: AgentMessage = serde_json::from_value(json).unwrap();
301 assert_eq!(parsed, AgentMessage::Done);
302 }
303
304 #[test]
305 fn test_tool_result_roundtrip_with_type_tag() {
306 let msg = AgentMessage::ToolResult {
307 result: ToolCallResult {
308 id: "call_1".to_string(),
309 name: "coding__read_file".to_string(),
310 arguments: "{}".to_string(),
311 result: "ok".to_string(),
312 },
313 result_meta: None,
314 model_name: "test".to_string(),
315 };
316 let json = serde_json::to_string(&msg).unwrap();
317 assert!(json.contains(r#""type":"tool_result""#), "missing type tag: {json}");
318
319 let parsed: AgentMessage = serde_json::from_str(&json).unwrap();
320 assert_eq!(parsed, msg);
321 }
322
323 #[test]
324 fn test_context_usage_serializes_with_type_tag() {
325 let msg = AgentMessage::ContextUsageUpdate {
326 usage_ratio: Some(0.5),
327 context_limit: Some(200_000),
328 input_tokens: 1000,
329 output_tokens: 200,
330 cache_read_tokens: None,
331 cache_creation_tokens: None,
332 reasoning_tokens: None,
333 total_input_tokens: 3000,
334 total_output_tokens: 600,
335 total_cache_read_tokens: 0,
336 total_cache_creation_tokens: 0,
337 total_reasoning_tokens: 0,
338 };
339 let json = serde_json::to_value(&msg).unwrap();
340 assert_eq!(json["type"], "context_usage");
341
342 let parsed: AgentMessage = serde_json::from_value(json).unwrap();
343 assert_eq!(parsed, msg);
344 }
345}