Skip to main content

aether_core/context/
ext.rs

1use crate::events::{AgentEvent, ContextEvent, MessageEvent, ToolEvent, TurnEvent};
2use crate::session::{SessionEvent, UserEvent};
3use llm::types::IsoString;
4use llm::{AssistantReasoning, ChatMessage, Context, ToolCallError, ToolCallResult};
5
6pub fn conversation_messages_from_events(events: &[SessionEvent]) -> Vec<ChatMessage> {
7    Context::from_events(events).messages().iter().filter(|message| !message.is_system()).cloned().collect()
8}
9
10pub trait ContextExt {
11    fn from_events(events: &[SessionEvent]) -> Self
12    where
13        Self: Sized;
14}
15
16impl ContextExt for Context {
17    fn from_events(events: &[SessionEvent]) -> Self {
18        let mut context = Context::new(vec![], vec![]);
19        let mut acc = TurnAccumulator::default();
20        for event in events {
21            match event {
22                SessionEvent::User(e) => apply_user_event(&mut context, e),
23                SessionEvent::Agent(m) => apply_agent_event(&mut context, m, &mut acc),
24                SessionEvent::Control(_) => {}
25            }
26        }
27        context
28    }
29}
30
31#[derive(Default)]
32struct TurnAccumulator {
33    text: String,
34    reasoning: String,
35    tool_results: Vec<Result<ToolCallResult, ToolCallError>>,
36}
37
38fn apply_user_event(ctx: &mut Context, event: &UserEvent) {
39    match event {
40        UserEvent::Message { content } => {
41            ctx.add_message(ChatMessage::User { content: content.clone(), timestamp: IsoString::now() });
42        }
43        UserEvent::ClearContext => {
44            ctx.clear_conversation();
45        }
46    }
47}
48
49fn apply_agent_event(ctx: &mut Context, event: &AgentEvent, acc: &mut TurnAccumulator) {
50    match event {
51        AgentEvent::Message(MessageEvent::Text { chunk, is_complete: true, .. }) => {
52            acc.text.clone_from(chunk);
53        }
54        AgentEvent::Message(MessageEvent::Thought { chunk, is_complete: true, .. }) => {
55            acc.reasoning.clone_from(chunk);
56        }
57        AgentEvent::Tool(ToolEvent::Result { result, .. }) => {
58            acc.tool_results.push(Ok(result.clone()));
59        }
60        AgentEvent::Tool(ToolEvent::Error { error, .. }) => {
61            acc.tool_results.push(Err(error.clone()));
62        }
63        AgentEvent::Turn(TurnEvent::Ended { .. }) => {
64            let text = std::mem::take(&mut acc.text);
65            let reasoning_text = std::mem::take(&mut acc.reasoning);
66            let tools = std::mem::take(&mut acc.tool_results);
67            if !text.is_empty() || !tools.is_empty() {
68                let reasoning = AssistantReasoning::from_parts(reasoning_text, None);
69                ctx.push_assistant_turn(&text, reasoning, tools);
70            }
71        }
72        AgentEvent::Context(ContextEvent::Cleared) => {
73            ctx.clear_conversation();
74            acc.text.clear();
75            acc.reasoning.clear();
76            acc.tool_results.clear();
77        }
78        AgentEvent::Context(ContextEvent::CompactionResult { summary, .. }) => {
79            *ctx = ctx.with_compacted_summary(summary);
80        }
81        _ => {}
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use crate::events::TurnOutcome;
88    use crate::session::{SessionControlEvent, last_agent_from_events};
89
90    use super::*;
91    use llm::ToolCallResult;
92
93    fn system_context() -> Context {
94        Context::new(
95            vec![ChatMessage::System { content: "You are helpful.".to_string(), timestamp: IsoString::now() }],
96            vec![],
97        )
98    }
99
100    fn user_msg(content: &str) -> UserEvent {
101        UserEvent::Message { content: vec![llm::ContentBlock::text(content)] }
102    }
103
104    fn user_session(content: &str) -> SessionEvent {
105        SessionEvent::User(user_msg(content))
106    }
107
108    fn text_complete(chunk: &str) -> AgentEvent {
109        AgentEvent::text("msg_1", chunk, true)
110    }
111
112    fn tool_result(id: &str, name: &str, result: &str) -> AgentEvent {
113        AgentEvent::Tool(ToolEvent::Result {
114            result: ToolCallResult {
115                id: id.to_string(),
116                name: name.to_string(),
117                arguments: "{}".to_string(),
118                result: result.to_string(),
119            },
120            result_meta: None,
121        })
122    }
123
124    fn agent_session(msg: AgentEvent) -> SessionEvent {
125        SessionEvent::Agent(msg)
126    }
127
128    /// Runs a sequence of agent events against a `system_context` and returns the context.
129    fn run_agent_events(events: &[AgentEvent]) -> Context {
130        let mut ctx = system_context();
131        let mut acc = TurnAccumulator::default();
132        for event in events {
133            apply_agent_event(&mut ctx, event, &mut acc);
134        }
135        ctx
136    }
137
138    #[test]
139    fn from_events_ignores_control_events() {
140        let ctx = Context::from_events(&[
141            user_session("Hello"),
142            SessionEvent::Control(SessionControlEvent::AgentSwitched {
143                from: Some("Planner".to_string()),
144                to: Some("Coder".to_string()),
145            }),
146            agent_session(text_complete("Hi there!")),
147            agent_session(AgentEvent::turn_ended(TurnOutcome::Completed)),
148        ]);
149
150        assert_eq!(ctx.message_count(), 2);
151        assert!(matches!(ctx.messages()[0], ChatMessage::User { .. }));
152        assert!(matches!(ctx.messages()[1], ChatMessage::Assistant { .. }));
153    }
154
155    #[test]
156    fn conversation_messages_from_events_filters_system_messages() {
157        let messages = conversation_messages_from_events(&[
158            user_session("Hello"),
159            agent_session(text_complete("Hi there!")),
160            agent_session(AgentEvent::turn_ended(TurnOutcome::Completed)),
161        ]);
162
163        assert_eq!(messages.len(), 2);
164        assert!(messages.iter().all(|message| !message.is_system()));
165    }
166
167    #[test]
168    fn last_agent_from_events_returns_last_switched_agent() {
169        let events = [
170            SessionEvent::Control(SessionControlEvent::AgentSwitched { from: None, to: Some("Planner".to_string()) }),
171            user_session("Hello"),
172            SessionEvent::Control(SessionControlEvent::AgentSwitched {
173                from: Some("Planner".to_string()),
174                to: Some("Coder".to_string()),
175            }),
176        ];
177
178        assert_eq!(last_agent_from_events(Some("Default".to_string()), &events), Some("Coder".to_string()));
179    }
180
181    #[test]
182    fn apply_user_message_adds_user_message() {
183        let mut ctx = system_context();
184        apply_user_event(&mut ctx, &user_msg("Hello"));
185        assert_eq!(ctx.message_count(), 2);
186        match &ctx.messages()[1] {
187            ChatMessage::User { content, .. } => {
188                assert_eq!(content, &vec![llm::ContentBlock::text("Hello")]);
189            }
190            other => panic!("Expected User, got {other:?}"),
191        }
192    }
193
194    #[test]
195    fn apply_user_clear_retains_system_messages() {
196        let mut ctx = system_context();
197        apply_user_event(&mut ctx, &user_msg("Hello"));
198        apply_user_event(&mut ctx, &UserEvent::ClearContext);
199        assert_eq!(ctx.message_count(), 1);
200        assert!(ctx.messages()[0].is_system());
201    }
202
203    #[test]
204    fn apply_agent_produces_assistant_and_tool_results() {
205        let ctx = run_agent_events(&[
206            tool_result("call_1", "read_file", "file contents"),
207            text_complete("Here is the file"),
208            AgentEvent::turn_ended(TurnOutcome::Completed),
209        ]);
210
211        assert_eq!(ctx.message_count(), 3);
212        match &ctx.messages()[1] {
213            ChatMessage::Assistant { content, tool_calls, .. } => {
214                assert_eq!(content, "Here is the file");
215                assert_eq!(tool_calls.len(), 1);
216                assert_eq!(tool_calls[0].name, "read_file");
217            }
218            other => panic!("Expected Assistant, got {other:?}"),
219        }
220        assert!(ctx.messages()[2].is_tool_result());
221    }
222
223    #[test]
224    fn apply_agent_context_cleared() {
225        let mut ctx = system_context();
226        let mut acc = TurnAccumulator::default();
227        apply_user_event(&mut ctx, &user_msg("Hello"));
228        apply_agent_event(&mut ctx, &AgentEvent::Context(ContextEvent::Cleared), &mut acc);
229        assert_eq!(ctx.message_count(), 1);
230        assert!(ctx.messages()[0].is_system());
231    }
232
233    #[test]
234    fn apply_agent_compaction_replaces_with_summary() {
235        let mut ctx = system_context();
236        let mut acc = TurnAccumulator::default();
237        apply_user_event(&mut ctx, &user_msg("Hello"));
238        apply_agent_event(
239            &mut ctx,
240            &AgentEvent::Context(ContextEvent::CompactionResult {
241                summary: "Summary of conversation".to_string(),
242                messages_removed: 1,
243            }),
244            &mut acc,
245        );
246        assert_eq!(ctx.message_count(), 2);
247        assert!(ctx.messages()[0].is_system());
248        assert!(ctx.messages()[1].is_summary());
249    }
250
251    #[test]
252    fn done_without_content_does_not_add_message() {
253        let ctx = run_agent_events(&[AgentEvent::turn_ended(TurnOutcome::Completed)]);
254        assert_eq!(ctx.message_count(), 1);
255    }
256
257    #[test]
258    fn streaming_chunks_are_ignored() {
259        let ctx = run_agent_events(&[AgentEvent::text("msg_1", "partial", false)]);
260        assert_eq!(ctx.message_count(), 1);
261    }
262
263    #[test]
264    fn accumulator_resets_after_done() {
265        let ctx = run_agent_events(&[
266            text_complete("Turn 1"),
267            AgentEvent::turn_ended(TurnOutcome::Completed),
268            AgentEvent::text("msg_2", "Turn 2", true),
269            AgentEvent::turn_ended(TurnOutcome::Completed),
270        ]);
271        assert_eq!(ctx.message_count(), 3);
272    }
273
274    #[test]
275    fn user_event_serde_roundtrip() {
276        let cases: Vec<UserEvent> = vec![user_msg("Hello"), UserEvent::ClearContext];
277        for event in cases {
278            let json = serde_json::to_string(&event).unwrap();
279            let parsed: UserEvent = serde_json::from_str(&json).unwrap();
280            assert_eq!(parsed, event);
281        }
282    }
283
284    #[test]
285    fn from_events_basic_conversation() {
286        let ctx = Context::from_events(&[
287            user_session("Hello"),
288            agent_session(text_complete("Hi there!")),
289            agent_session(AgentEvent::turn_ended(TurnOutcome::Completed)),
290        ]);
291        assert_eq!(ctx.message_count(), 2);
292        assert!(matches!(ctx.messages()[0], ChatMessage::User { .. }));
293        assert!(matches!(ctx.messages()[1], ChatMessage::Assistant { .. }));
294    }
295
296    #[test]
297    fn from_events_with_tool_calls() {
298        let ctx = Context::from_events(&[
299            user_session("Read Cargo.toml"),
300            agent_session(AgentEvent::Tool(ToolEvent::Call {
301                request: llm::ToolCallRequest {
302                    id: "call_1".to_string(),
303                    name: "read_file".to_string(),
304                    arguments: "{}".to_string(),
305                },
306            })),
307            agent_session(tool_result("call_1", "read_file", "file contents")),
308            agent_session(text_complete("Here is the file")),
309            agent_session(AgentEvent::turn_ended(TurnOutcome::Completed)),
310        ]);
311
312        assert_eq!(ctx.message_count(), 3);
313        match &ctx.messages()[1] {
314            ChatMessage::Assistant { tool_calls, .. } => {
315                assert_eq!(tool_calls.len(), 1);
316                assert_eq!(tool_calls[0].name, "read_file");
317            }
318            other => panic!("Expected Assistant, got {other:?}"),
319        }
320        assert!(ctx.messages()[2].is_tool_result());
321    }
322
323    #[test]
324    fn from_events_handles_clear() {
325        let ctx = Context::from_events(&[
326            user_session("Hello"),
327            agent_session(text_complete("Hi!")),
328            agent_session(AgentEvent::turn_ended(TurnOutcome::Completed)),
329            SessionEvent::User(UserEvent::ClearContext),
330            user_session("Start fresh"),
331        ]);
332        assert_eq!(ctx.message_count(), 1);
333        assert!(matches!(ctx.messages()[0], ChatMessage::User { .. }));
334    }
335
336    #[test]
337    fn from_events_handles_compaction() {
338        let ctx = Context::from_events(&[
339            user_session("Hello"),
340            agent_session(text_complete("Hi!")),
341            agent_session(AgentEvent::turn_ended(TurnOutcome::Completed)),
342            agent_session(AgentEvent::Context(ContextEvent::CompactionResult {
343                summary: "Earlier we greeted each other.".to_string(),
344                messages_removed: 2,
345            })),
346            user_session("What did we talk about?"),
347        ]);
348        assert_eq!(ctx.message_count(), 2);
349        assert!(ctx.messages()[0].is_summary());
350    }
351
352    #[test]
353    fn from_events_empty() {
354        let ctx = Context::from_events(&[]);
355        assert_eq!(ctx.message_count(), 0);
356    }
357}