Skip to main content

aether_core/events/
agent_message.rs

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/// Message from the agent to the user.
10#[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    /// Context compaction has been triggered.
65    ContextCompactionStarted {
66        message_count: usize,
67    },
68
69    /// Context was compacted to reduce token usage.
70    ContextCompactionResult {
71        summary: String,
72        messages_removed: usize,
73    },
74
75    /// Context usage update for UI display.
76    #[serde(rename = "context_usage")]
77    ContextUsageUpdate {
78        /// Current usage ratio (0.0 - 1.0), if context window is known.
79        usage_ratio: Option<f64>,
80        /// Maximum context limit, if known.
81        context_limit: Option<u32>,
82        /// Input tokens on the most recent API call (the current context size).
83        input_tokens: u32,
84        /// Output tokens on the most recent API call.
85        output_tokens: u32,
86        /// Prompt tokens served from cache on the most recent API call.
87        cache_read_tokens: Option<u32>,
88        /// Prompt tokens written to cache on the most recent API call.
89        cache_creation_tokens: Option<u32>,
90        /// Reasoning tokens spent on the most recent API call.
91        reasoning_tokens: Option<u32>,
92        /// Cumulative input tokens since the agent started.
93        total_input_tokens: u64,
94        /// Cumulative output tokens since the agent started.
95        total_output_tokens: u64,
96        /// Cumulative cache-read tokens since the agent started.
97        total_cache_read_tokens: u64,
98        /// Cumulative cache-creation tokens since the agent started.
99        total_cache_creation_tokens: u64,
100        /// Cumulative reasoning tokens since the agent started.
101        total_reasoning_tokens: u64,
102    },
103
104    /// Agent is auto-continuing because LLM stopped with a resumable stop reason.
105    AutoContinue {
106        /// Current attempt number (1-indexed).
107        attempt: u32,
108        /// Maximum allowed attempts.
109        max_attempts: u32,
110    },
111
112    /// Agent is retrying after a transient LLM provider failure
113    Retrying {
114        /// Current retry attempt number (1-indexed).
115        attempt: u32,
116        /// Maximum allowed attempts.
117        max_attempts: u32,
118        /// Backoff delay in milliseconds before the retry fires.
119        delay_ms: u64,
120        /// The error that triggered the retry, formatted for display.
121        error: String,
122    },
123
124    /// The model was successfully switched.
125    ModelSwitched {
126        previous: String,
127        new: String,
128    },
129
130    /// The agent context was cleared and reset to its blank state.
131    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}