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