1use crate::events::{AgentEvent, ContextEvent, MessageEvent, ToolEvent, TurnEvent, task_created_result};
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::TaskCreated { request, task_id, .. }) => {
61 acc.tool_results.push(Ok(task_created_result(request, task_id)));
62 }
63 AgentEvent::Tool(ToolEvent::Error { error }) => {
64 acc.tool_results.push(Err(error.clone()));
65 }
66 AgentEvent::Turn(TurnEvent::Ended { .. }) => {
67 let text = std::mem::take(&mut acc.text);
68 let reasoning_text = std::mem::take(&mut acc.reasoning);
69 let tools = std::mem::take(&mut acc.tool_results);
70 if !text.is_empty() || !tools.is_empty() {
71 let reasoning = AssistantReasoning::from_parts(reasoning_text, None);
72 ctx.push_assistant_turn(&text, reasoning, tools);
73 }
74 }
75 AgentEvent::Context(ContextEvent::Cleared) => {
76 ctx.clear_conversation();
77 acc.text.clear();
78 acc.reasoning.clear();
79 acc.tool_results.clear();
80 }
81 AgentEvent::Context(ContextEvent::CompactionResult { summary, .. }) => {
82 *ctx = ctx.with_compacted_summary(summary);
83 }
84 AgentEvent::Tool(
85 event @ (ToolEvent::TaskCompleted { .. } | ToolEvent::TaskFailed { .. } | ToolEvent::TaskCancelled { .. }),
86 ) => {
87 if let Some(message) = event.task_context_message() {
88 ctx.add_message(message);
89 }
90 }
91 _ => {}
92 }
93}
94
95#[cfg(test)]
96mod tests {
97 use crate::events::{StreamState, TurnOutcome};
98 use crate::session::{SessionControlEvent, last_agent_from_events};
99
100 use super::*;
101 use llm::ToolCallResult;
102
103 fn system_context() -> Context {
104 Context::new(vec![ChatMessage::system("You are helpful.")], vec![])
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) -> AgentEvent {
116 AgentEvent::text("msg_1", chunk, StreamState::Complete)
117 }
118
119 fn tool_result(id: &str, name: &str, result: &str) -> AgentEvent {
120 AgentEvent::Tool(ToolEvent::Result {
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 })
129 }
130
131 fn agent_session(msg: AgentEvent) -> SessionEvent {
132 SessionEvent::Agent(msg)
133 }
134
135 fn run_agent_events(events: &[AgentEvent]) -> Context {
137 let mut ctx = system_context();
138 let mut acc = TurnAccumulator::default();
139 for event in events {
140 apply_agent_event(&mut ctx, event, &mut acc);
141 }
142 ctx
143 }
144
145 #[test]
146 fn from_events_ignores_control_events() {
147 let ctx = Context::from_events(&[
148 user_session("Hello"),
149 SessionEvent::Control(SessionControlEvent::AgentSwitched {
150 from: Some("Planner".to_string()),
151 to: Some("Coder".to_string()),
152 }),
153 agent_session(text_complete("Hi there!")),
154 agent_session(AgentEvent::turn_ended(TurnOutcome::Completed)),
155 ]);
156
157 assert_eq!(ctx.message_count(), 2);
158 assert!(matches!(ctx.messages()[0], ChatMessage::User { .. }));
159 assert!(matches!(ctx.messages()[1], ChatMessage::Assistant { .. }));
160 }
161
162 #[test]
163 fn conversation_messages_from_events_filters_system_messages() {
164 let messages = conversation_messages_from_events(&[
165 user_session("Hello"),
166 agent_session(text_complete("Hi there!")),
167 agent_session(AgentEvent::turn_ended(TurnOutcome::Completed)),
168 ]);
169
170 assert_eq!(messages.len(), 2);
171 assert!(messages.iter().all(|message| !message.is_system()));
172 }
173
174 #[test]
175 fn last_agent_from_events_returns_last_switched_agent() {
176 let events = [
177 SessionEvent::Control(SessionControlEvent::AgentSwitched { from: None, to: Some("Planner".to_string()) }),
178 user_session("Hello"),
179 SessionEvent::Control(SessionControlEvent::AgentSwitched {
180 from: Some("Planner".to_string()),
181 to: Some("Coder".to_string()),
182 }),
183 ];
184
185 assert_eq!(last_agent_from_events(Some("Default".to_string()), &events), Some("Coder".to_string()));
186 }
187
188 #[test]
189 fn apply_user_message_adds_user_message() {
190 let mut ctx = system_context();
191 apply_user_event(&mut ctx, &user_msg("Hello"));
192 assert_eq!(ctx.message_count(), 2);
193 match &ctx.messages()[1] {
194 ChatMessage::User { content, .. } => {
195 assert_eq!(content, &vec![llm::ContentBlock::text("Hello")]);
196 }
197 other => panic!("Expected User, got {other:?}"),
198 }
199 }
200
201 #[test]
202 fn apply_user_clear_retains_system_messages() {
203 let mut ctx = system_context();
204 apply_user_event(&mut ctx, &user_msg("Hello"));
205 apply_user_event(&mut ctx, &UserEvent::ClearContext);
206 assert_eq!(ctx.message_count(), 1);
207 assert!(ctx.messages()[0].is_system());
208 }
209
210 #[test]
211 fn apply_agent_produces_assistant_and_tool_results() {
212 let ctx = run_agent_events(&[
213 tool_result("call_1", "read_file", "file contents"),
214 text_complete("Here is the file"),
215 AgentEvent::turn_ended(TurnOutcome::Completed),
216 ]);
217
218 assert_eq!(ctx.message_count(), 3);
219 match &ctx.messages()[1] {
220 ChatMessage::Assistant { content, tool_calls, .. } => {
221 assert_eq!(content, "Here is the file");
222 assert_eq!(tool_calls.len(), 1);
223 assert_eq!(tool_calls[0].name, "read_file");
224 }
225 other => panic!("Expected Assistant, got {other:?}"),
226 }
227 assert!(ctx.messages()[2].is_tool_result());
228 }
229
230 #[test]
231 fn apply_agent_context_cleared() {
232 let mut ctx = system_context();
233 let mut acc = TurnAccumulator::default();
234 apply_user_event(&mut ctx, &user_msg("Hello"));
235 apply_agent_event(&mut ctx, &AgentEvent::Context(ContextEvent::Cleared), &mut acc);
236 assert_eq!(ctx.message_count(), 1);
237 assert!(ctx.messages()[0].is_system());
238 }
239
240 #[test]
241 fn apply_agent_compaction_replaces_with_summary() {
242 let mut ctx = system_context();
243 let mut acc = TurnAccumulator::default();
244 apply_user_event(&mut ctx, &user_msg("Hello"));
245 apply_agent_event(
246 &mut ctx,
247 &AgentEvent::Context(ContextEvent::CompactionResult {
248 summary: "Summary of conversation".to_string(),
249 messages_removed: 1,
250 }),
251 &mut acc,
252 );
253 assert_eq!(ctx.message_count(), 2);
254 assert!(ctx.messages()[0].is_system());
255 assert!(ctx.messages()[1].is_summary());
256 }
257
258 #[test]
259 fn done_without_content_does_not_add_message() {
260 let ctx = run_agent_events(&[AgentEvent::turn_ended(TurnOutcome::Completed)]);
261 assert_eq!(ctx.message_count(), 1);
262 }
263
264 #[test]
265 fn streaming_chunks_are_ignored() {
266 let ctx = run_agent_events(&[AgentEvent::text("msg_1", "partial", StreamState::Partial)]);
267 assert_eq!(ctx.message_count(), 1);
268 }
269
270 #[test]
271 fn accumulator_resets_after_done() {
272 let ctx = run_agent_events(&[
273 text_complete("Turn 1"),
274 AgentEvent::turn_ended(TurnOutcome::Completed),
275 AgentEvent::text("msg_2", "Turn 2", StreamState::Complete),
276 AgentEvent::turn_ended(TurnOutcome::Completed),
277 ]);
278 assert_eq!(ctx.message_count(), 3);
279 }
280
281 #[test]
282 fn user_event_serde_roundtrip() {
283 let cases: Vec<UserEvent> = vec![user_msg("Hello"), UserEvent::ClearContext];
284 for event in cases {
285 let json = serde_json::to_string(&event).unwrap();
286 let parsed: UserEvent = serde_json::from_str(&json).unwrap();
287 assert_eq!(parsed, event);
288 }
289 }
290
291 #[test]
292 fn from_events_basic_conversation() {
293 let ctx = Context::from_events(&[
294 user_session("Hello"),
295 agent_session(text_complete("Hi there!")),
296 agent_session(AgentEvent::turn_ended(TurnOutcome::Completed)),
297 ]);
298 assert_eq!(ctx.message_count(), 2);
299 assert!(matches!(ctx.messages()[0], ChatMessage::User { .. }));
300 assert!(matches!(ctx.messages()[1], ChatMessage::Assistant { .. }));
301 }
302
303 #[test]
304 fn from_events_with_tool_calls() {
305 let ctx = Context::from_events(&[
306 user_session("Read Cargo.toml"),
307 agent_session(AgentEvent::Tool(ToolEvent::Call {
308 request: llm::ToolCallRequest {
309 id: "call_1".to_string(),
310 name: "read_file".to_string(),
311 arguments: "{}".to_string(),
312 },
313 })),
314 agent_session(tool_result("call_1", "read_file", "file contents")),
315 agent_session(text_complete("Here is the file")),
316 agent_session(AgentEvent::turn_ended(TurnOutcome::Completed)),
317 ]);
318
319 assert_eq!(ctx.message_count(), 3);
320 match &ctx.messages()[1] {
321 ChatMessage::Assistant { tool_calls, .. } => {
322 assert_eq!(tool_calls.len(), 1);
323 assert_eq!(tool_calls[0].name, "read_file");
324 }
325 other => panic!("Expected Assistant, got {other:?}"),
326 }
327 assert!(ctx.messages()[2].is_tool_result());
328 }
329
330 #[test]
331 fn from_events_preserves_failed_tool_calls() {
332 let ctx = Context::from_events(&[
333 user_session("Read missing.txt"),
334 agent_session(AgentEvent::Tool(ToolEvent::Error {
335 error: ToolCallError {
336 id: "call_1".to_string(),
337 name: "read_file".to_string(),
338 arguments: Some(r#"{"path":"missing.txt"}"#.to_string()),
339 error: "file not found".to_string(),
340 },
341 })),
342 agent_session(AgentEvent::turn_ended(TurnOutcome::Completed)),
343 ]);
344
345 assert_eq!(ctx.message_count(), 3);
346 assert!(matches!(
347 &ctx.messages()[1],
348 ChatMessage::Assistant { tool_calls, .. }
349 if tool_calls.len() == 1 && tool_calls[0].id == "call_1"
350 ));
351 assert!(matches!(
352 &ctx.messages()[2],
353 ChatMessage::ToolCallResult(Err(error))
354 if error.id == "call_1" && error.error == "file not found"
355 ));
356 }
357
358 #[test]
359 fn from_events_handles_clear() {
360 let ctx = Context::from_events(&[
361 user_session("Hello"),
362 agent_session(text_complete("Hi!")),
363 agent_session(AgentEvent::turn_ended(TurnOutcome::Completed)),
364 SessionEvent::User(UserEvent::ClearContext),
365 user_session("Start fresh"),
366 ]);
367 assert_eq!(ctx.message_count(), 1);
368 assert!(matches!(ctx.messages()[0], ChatMessage::User { .. }));
369 }
370
371 #[test]
372 fn from_events_handles_compaction() {
373 let ctx = Context::from_events(&[
374 user_session("Hello"),
375 agent_session(text_complete("Hi!")),
376 agent_session(AgentEvent::turn_ended(TurnOutcome::Completed)),
377 agent_session(AgentEvent::Context(ContextEvent::CompactionResult {
378 summary: "Earlier we greeted each other.".to_string(),
379 messages_removed: 2,
380 })),
381 user_session("What did we talk about?"),
382 ]);
383 assert_eq!(ctx.message_count(), 2);
384 assert!(ctx.messages()[0].is_summary());
385 }
386
387 #[test]
388 fn from_events_empty() {
389 let ctx = Context::from_events(&[]);
390 assert_eq!(ctx.message_count(), 0);
391 }
392}