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