aether_core/context/
ext.rs

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

aether_core/context/
ext.rs
