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/// Context/token usage reported after an LLM call.
10#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize, JsonSchema)]
11pub struct ContextUsage {
12    /// Current usage ratio (0.0 - 1.0), if context window is known.
13    pub usage_ratio: Option<f64>,
14    /// Maximum context limit, if known.
15    pub context_limit: Option<u32>,
16    /// Input tokens on the most recent API call (the current context size).
17    pub input_tokens: u32,
18    /// Output tokens on the most recent API call.
19    pub output_tokens: u32,
20    /// Prompt tokens served from cache on the most recent API call.
21    pub cache_read_tokens: Option<u32>,
22    /// Prompt tokens written to cache on the most recent API call.
23    pub cache_creation_tokens: Option<u32>,
24    /// Reasoning tokens spent on the most recent API call.
25    pub reasoning_tokens: Option<u32>,
26    /// Cumulative input tokens since the agent started.
27    pub total_input_tokens: u64,
28    /// Cumulative output tokens since the agent started.
29    pub total_output_tokens: u64,
30    /// Cumulative cache-read tokens since the agent started.
31    pub total_cache_read_tokens: u64,
32    /// Cumulative cache-creation tokens since the agent started.
33    pub total_cache_creation_tokens: u64,
34    /// Cumulative reasoning tokens since the agent started.
35    pub total_reasoning_tokens: u64,
36}
37
38impl ContextUsage {
39    /// Sum of cumulative input + output tokens.
40    pub fn total_tokens(&self) -> u64 {
41        self.total_input_tokens + self.total_output_tokens
42    }
43}
44
45/// Message from the agent to the user.
46#[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    /// Context compaction has been triggered.
101    ContextCompactionStarted {
102        message_count: usize,
103    },
104
105    /// Context was compacted to reduce token usage.
106    ContextCompactionResult {
107        summary: String,
108        messages_removed: usize,
109    },
110
111    /// Context usage update for UI display.
112    #[serde(rename = "context_usage")]
113    ContextUsageUpdate {
114        #[serde(flatten)]
115        usage: ContextUsage,
116    },
117
118    /// Agent is auto-continuing because LLM stopped with a resumable stop reason.
119    AutoContinue {
120        /// Current attempt number (1-indexed).
121        attempt: u32,
122        /// Maximum allowed attempts.
123        max_attempts: u32,
124    },
125
126    /// Agent is retrying after a transient LLM provider failure
127    Retrying {
128        /// Current retry attempt number (1-indexed).
129        attempt: u32,
130        /// Maximum allowed attempts.
131        max_attempts: u32,
132        /// Backoff delay in milliseconds before the retry fires.
133        delay_ms: u64,
134        /// The error that triggered the retry, formatted for display.
135        error: String,
136    },
137
138    /// The model was successfully switched.
139    ModelSwitched {
140        previous: String,
141        new: String,
142    },
143
144    /// The agent context was cleared and reset to its blank state.
145    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}