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