aether_core/context/
ext.rs

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

aether_core/context/
ext.rs
