//! Session and thread model for turn-based agent interactions. //! //! A Session contains one or more Threads. Each Thread represents a //! conversation/interaction sequence with the agent. Threads contain //! Turns, which are request/response pairs. //! //! This model supports: //! - Undo: Roll back to a previous turn //! - Interrupt: Cancel the current turn mid-execution //! - Compaction: Summarize old turns to save context //! - Resume: Continue from a saved checkpoint use std::collections::{HashMap, HashSet}; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use uuid::Uuid; use crate::channels::web::util::truncate_preview; use crate::llm::{ChatMessage, ToolCall}; /// A session containing one or more threads. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Session { /// Unique session ID. pub id: Uuid, /// User ID that owns this session. pub user_id: String, /// Active thread ID. pub active_thread: Option, /// All threads in this session. pub threads: HashMap, /// When the session was created. pub created_at: DateTime, /// When the session was last active. pub last_active_at: DateTime, /// Session metadata. pub metadata: serde_json::Value, /// Tools that have been auto-approved for this session ("always approve"). #[serde(default)] pub auto_approved_tools: HashSet, } impl Session { /// Create a new session. pub fn new(user_id: impl Into) -> Self { let now = Utc::now(); Self { id: Uuid::new_v4(), user_id: user_id.into(), active_thread: None, threads: HashMap::new(), created_at: now, last_active_at: now, metadata: serde_json::Value::Null, auto_approved_tools: HashSet::new(), } } /// Check if a tool has been auto-approved for this session. pub fn is_tool_auto_approved(&self, tool_name: &str) -> bool { self.auto_approved_tools.contains(tool_name) } /// Add a tool to the auto-approved set. pub fn auto_approve_tool(&mut self, tool_name: impl Into) { self.auto_approved_tools.insert(tool_name.into()); } /// Create a new thread in this session. pub fn create_thread(&mut self) -> &mut Thread { let thread = Thread::new(self.id); let thread_id = thread.id; self.active_thread = Some(thread_id); self.last_active_at = Utc::now(); self.threads.entry(thread_id).or_insert(thread) } /// Get the active thread. pub fn active_thread(&self) -> Option<&Thread> { self.active_thread.and_then(|id| self.threads.get(&id)) } /// Get the active thread mutably. pub fn active_thread_mut(&mut self) -> Option<&mut Thread> { self.active_thread.and_then(|id| self.threads.get_mut(&id)) } /// Get or create the active thread. pub fn get_or_create_thread(&mut self) -> &mut Thread { match self.active_thread { None => self.create_thread(), Some(id) => { if self.threads.contains_key(&id) { // Entry existence confirmed by contains_key above. // get_mut borrows self.threads mutably, so we can't // combine the check and access into if-let without // conflicting with the self.create_thread() fallback. self.threads.get_mut(&id).unwrap() // safety: contains_key guard above } else { // Stale active_thread ID: create a new thread, which // updates self.active_thread to the new thread's ID. self.create_thread() } } } } /// Switch to a different thread. pub fn switch_thread(&mut self, thread_id: Uuid) -> bool { if self.threads.contains_key(&thread_id) { self.active_thread = Some(thread_id); self.last_active_at = Utc::now(); true } else { false } } } /// State of a thread. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] pub enum ThreadState { /// Thread is idle, waiting for input. Idle, /// Thread is processing a turn. Processing, /// Thread is waiting for user approval. AwaitingApproval, /// Thread has completed (no more turns expected). Completed, /// Thread was interrupted. Interrupted, } /// Pending auth token request. /// /// When `tool_auth` returns `awaiting_token`, the thread enters auth mode. /// The next user message is intercepted before entering the normal pipeline /// (no logging, no turn creation, no history) and routed directly to the /// credential store. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct PendingAuth { /// Extension name to authenticate. pub extension_name: String, } /// Pending tool approval request stored on a thread. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct PendingApproval { /// Unique request ID. pub request_id: Uuid, /// Tool name requiring approval. pub tool_name: String, /// Tool parameters (original values, used for execution). pub parameters: serde_json::Value, /// Redacted tool parameters (sensitive values replaced with `[REDACTED]`). /// Used for display in approval UI, logs, and SSE broadcasts. #[serde(default)] pub display_parameters: serde_json::Value, /// Description of what the tool will do. pub description: String, /// Tool call ID from LLM (for proper context continuation). pub tool_call_id: String, /// Context messages at the time of the request (to resume from). pub context_messages: Vec, /// Remaining tool calls from the same assistant message that were not /// executed yet when approval was requested. #[serde(default)] pub deferred_tool_calls: Vec, /// User timezone at the time the approval was requested, so it persists /// through the approval flow even if the approval message lacks timezone. #[serde(default)] pub user_timezone: Option, } /// A conversation thread within a session. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Thread { /// Unique thread ID. pub id: Uuid, /// Parent session ID. pub session_id: Uuid, /// Current state. pub state: ThreadState, /// Turns in this thread. pub turns: Vec, /// When the thread was created. pub created_at: DateTime, /// When the thread was last updated. pub updated_at: DateTime, /// Thread metadata (e.g., title, tags). pub metadata: serde_json::Value, /// Pending approval request (when state is AwaitingApproval). #[serde(default)] pub pending_approval: Option, /// Pending auth token request (thread is in auth mode). #[serde(default)] pub pending_auth: Option, } impl Thread { /// Create a new thread. pub fn new(session_id: Uuid) -> Self { let now = Utc::now(); Self { id: Uuid::new_v4(), session_id, state: ThreadState::Idle, turns: Vec::new(), created_at: now, updated_at: now, metadata: serde_json::Value::Null, pending_approval: None, pending_auth: None, } } /// Create a thread with a specific ID (for DB hydration). pub fn with_id(id: Uuid, session_id: Uuid) -> Self { let now = Utc::now(); Self { id, session_id, state: ThreadState::Idle, turns: Vec::new(), created_at: now, updated_at: now, metadata: serde_json::Value::Null, pending_approval: None, pending_auth: None, } } /// Get the current turn number (1-indexed for display). pub fn turn_number(&self) -> usize { self.turns.len() + 1 } /// Get the last turn. pub fn last_turn(&self) -> Option<&Turn> { self.turns.last() } /// Get the last turn mutably. pub fn last_turn_mut(&mut self) -> Option<&mut Turn> { self.turns.last_mut() } /// Start a new turn with user input. pub fn start_turn(&mut self, user_input: impl Into) -> &mut Turn { let turn_number = self.turns.len(); let turn = Turn::new(turn_number, user_input); self.turns.push(turn); self.state = ThreadState::Processing; self.updated_at = Utc::now(); // turn_number was len() before push, so it's a valid index after push &mut self.turns[turn_number] } /// Complete the current turn with a response. pub fn complete_turn(&mut self, response: impl Into) { if let Some(turn) = self.turns.last_mut() { turn.complete(response); } self.state = ThreadState::Idle; self.updated_at = Utc::now(); } /// Fail the current turn with an error. pub fn fail_turn(&mut self, error: impl Into) { if let Some(turn) = self.turns.last_mut() { turn.fail(error); } self.state = ThreadState::Idle; self.updated_at = Utc::now(); } /// Mark the thread as awaiting approval with pending request details. pub fn await_approval(&mut self, pending: PendingApproval) { self.state = ThreadState::AwaitingApproval; self.pending_approval = Some(pending); self.updated_at = Utc::now(); } /// Take the pending approval (clearing it from the thread). pub fn take_pending_approval(&mut self) -> Option { self.pending_approval.take() } /// Clear pending approval and return to idle state. pub fn clear_pending_approval(&mut self) { self.pending_approval = None; self.state = ThreadState::Idle; self.updated_at = Utc::now(); } /// Enter auth mode: next user message will be routed directly to /// the credential store, bypassing the normal pipeline entirely. pub fn enter_auth_mode(&mut self, extension_name: String) { self.pending_auth = Some(PendingAuth { extension_name }); self.updated_at = Utc::now(); } /// Take the pending auth (clearing auth mode). pub fn take_pending_auth(&mut self) -> Option { self.pending_auth.take() } /// Interrupt the current turn. pub fn interrupt(&mut self) { if let Some(turn) = self.turns.last_mut() { turn.interrupt(); } self.state = ThreadState::Interrupted; self.updated_at = Utc::now(); } /// Resume after interruption. pub fn resume(&mut self) { if self.state == ThreadState::Interrupted { self.state = ThreadState::Idle; self.updated_at = Utc::now(); } } /// Get all messages for context building, including tool call history. /// /// Emits the full LLM-compatible message sequence per turn: /// `user → [assistant_with_tool_calls → tool_result*] → assistant` /// /// This ensures the LLM sees prior tool executions and won't re-attempt /// completed actions in subsequent turns. pub fn messages(&self) -> Vec { let mut messages = Vec::new(); for turn in &self.turns { if turn.image_content_parts.is_empty() { messages.push(ChatMessage::user(&turn.user_input)); } else { messages.push(ChatMessage::user_with_parts( &turn.user_input, turn.image_content_parts.clone(), )); } if !turn.tool_calls.is_empty() { // Build ToolCall objects with synthetic stable IDs let tool_calls: Vec = turn .tool_calls .iter() .enumerate() .map(|(i, tc)| ToolCall { id: format!("turn{}_{}", turn.turn_number, i), name: tc.name.clone(), arguments: tc.parameters.clone(), }) .collect(); // Assistant message declaring the tool calls (no text content) messages.push(ChatMessage::assistant_with_tool_calls(None, tool_calls)); // Individual tool result messages, truncated to limit context size. for (i, tc) in turn.tool_calls.iter().enumerate() { let call_id = format!("turn{}_{}", turn.turn_number, i); let content = if let Some(ref err) = tc.error { // .error already contains the full error text; // pass through without wrapping to avoid double-prefix. truncate_preview(err, 1000) } else if let Some(ref res) = tc.result { let raw = match res { serde_json::Value::String(s) => s.clone(), other => other.to_string(), }; truncate_preview(&raw, 1000) } else { "OK".to_string() }; messages.push(ChatMessage::tool_result(call_id, &tc.name, content)); } } if let Some(ref response) = turn.response { messages.push(ChatMessage::assistant(response)); } } messages } /// Truncate turns to a specific count (keeping most recent). pub fn truncate_turns(&mut self, keep: usize) { if self.turns.len() > keep { let drain_count = self.turns.len() - keep; self.turns.drain(0..drain_count); // Re-number remaining turns for (i, turn) in self.turns.iter_mut().enumerate() { turn.turn_number = i; } } } /// Restore thread state from a checkpoint's messages. /// /// Clears existing turns and rebuilds from the message sequence. /// Handles the full message pattern including tool messages: /// `user → [assistant_with_tool_calls → tool_result*] → assistant` /// /// Also supports the legacy pattern (user/assistant pairs only) for /// backward compatibility with old checkpoint data. pub fn restore_from_messages(&mut self, messages: Vec) { self.turns.clear(); self.state = ThreadState::Idle; let mut iter = messages.into_iter().peekable(); let mut turn_number = 0; while let Some(msg) = iter.next() { if msg.role == crate::llm::Role::User { let mut turn = Turn::new(turn_number, &msg.content); // Consume tool call sequences (assistant_with_tool_calls + tool_results). // A single turn may contain multiple rounds of tool calls, so we // track the cumulative base index into turn.tool_calls. while let Some(next) = iter.peek() { if next.role == crate::llm::Role::Assistant && next.tool_calls.is_some() { let call_base_idx = turn.tool_calls.len(); if let Some(assistant_msg) = iter.next() && let Some(ref tcs) = assistant_msg.tool_calls { for tc in tcs { turn.record_tool_call(&tc.name, tc.arguments.clone()); } } // Consume the corresponding tool_result messages, // indexing relative to this batch's base offset. let mut pos = 0; while let Some(tr) = iter.peek() { if tr.role != crate::llm::Role::Tool { break; } if let Some(tool_msg) = iter.next() { let idx = call_base_idx + pos; if idx < turn.tool_calls.len() { // Store as result — the error/success distinction // is for the live turn only; restored context just // needs the content the LLM originally saw. turn.tool_calls[idx].result = Some(serde_json::Value::String(tool_msg.content.clone())); } } pos += 1; } } else { break; } } // Check if next is the final assistant response for this turn let is_final_assistant = iter.peek().is_some_and(|n| { n.role == crate::llm::Role::Assistant && n.tool_calls.is_none() }); if is_final_assistant && let Some(response) = iter.next() { turn.complete(&response.content); } self.turns.push(turn); turn_number += 1; } else { // Skip non-user messages that aren't anchored to a turn continue; } } self.updated_at = Utc::now(); } } /// State of a turn. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] pub enum TurnState { /// Turn is being processed. Processing, /// Turn completed successfully. Completed, /// Turn failed with an error. Failed, /// Turn was interrupted. Interrupted, } /// A single turn (request/response pair) in a thread. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Turn { /// Turn number (0-indexed). pub turn_number: usize, /// User input that started this turn. pub user_input: String, /// Agent response (if completed). pub response: Option, /// Tool calls made during this turn. pub tool_calls: Vec, /// Turn state. pub state: TurnState, /// When the turn started. pub started_at: DateTime, /// When the turn completed. pub completed_at: Option>, /// Error message (if failed). pub error: Option, /// Transient image content parts for multimodal LLM input. /// Not serialized — images are only needed for the current LLM call. /// The text description in `user_input` persists for compaction/context. #[serde(skip)] pub image_content_parts: Vec, } impl Turn { /// Create a new turn. pub fn new(turn_number: usize, user_input: impl Into) -> Self { Self { turn_number, user_input: user_input.into(), response: None, tool_calls: Vec::new(), state: TurnState::Processing, started_at: Utc::now(), completed_at: None, error: None, image_content_parts: Vec::new(), } } /// Complete this turn. pub fn complete(&mut self, response: impl Into) { self.response = Some(response.into()); self.state = TurnState::Completed; self.completed_at = Some(Utc::now()); // Free image data — only needed for the initial LLM call, not subsequent turns self.image_content_parts.clear(); } /// Fail this turn. pub fn fail(&mut self, error: impl Into) { self.error = Some(error.into()); self.state = TurnState::Failed; self.completed_at = Some(Utc::now()); self.image_content_parts.clear(); } /// Interrupt this turn. pub fn interrupt(&mut self) { self.state = TurnState::Interrupted; self.completed_at = Some(Utc::now()); self.image_content_parts.clear(); } /// Record a tool call. pub fn record_tool_call(&mut self, name: impl Into, params: serde_json::Value) { self.tool_calls.push(TurnToolCall { name: name.into(), parameters: params, result: None, error: None, }); } /// Record tool call result. pub fn record_tool_result(&mut self, result: serde_json::Value) { if let Some(call) = self.tool_calls.last_mut() { call.result = Some(result); } } /// Record tool call error. pub fn record_tool_error(&mut self, error: impl Into) { if let Some(call) = self.tool_calls.last_mut() { call.error = Some(error.into()); } } } /// Record of a tool call made during a turn. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct TurnToolCall { /// Tool name. pub name: String, /// Parameters passed to the tool. pub parameters: serde_json::Value, /// Result from the tool (if successful). pub result: Option, /// Error from the tool (if failed). pub error: Option, } #[cfg(test)] mod tests { use super::*; #[test] fn test_session_creation() { let mut session = Session::new("user-123"); assert!(session.active_thread.is_none()); session.create_thread(); assert!(session.active_thread.is_some()); } #[test] fn test_thread_turns() { let mut thread = Thread::new(Uuid::new_v4()); thread.start_turn("Hello"); assert_eq!(thread.state, ThreadState::Processing); assert_eq!(thread.turns.len(), 1); thread.complete_turn("Hi there!"); assert_eq!(thread.state, ThreadState::Idle); assert_eq!(thread.turns[0].response, Some("Hi there!".to_string())); } #[test] fn test_thread_messages() { let mut thread = Thread::new(Uuid::new_v4()); thread.start_turn("First message"); thread.complete_turn("First response"); thread.start_turn("Second message"); thread.complete_turn("Second response"); let messages = thread.messages(); assert_eq!(messages.len(), 4); } #[test] fn test_turn_tool_calls() { let mut turn = Turn::new(0, "Test input"); turn.record_tool_call("echo", serde_json::json!({"message": "test"})); turn.record_tool_result(serde_json::json!("test")); assert_eq!(turn.tool_calls.len(), 1); assert!(turn.tool_calls[0].result.is_some()); } #[test] fn test_restore_from_messages() { let mut thread = Thread::new(Uuid::new_v4()); // First add some turns thread.start_turn("Original message"); thread.complete_turn("Original response"); // Now restore from different messages let messages = vec![ ChatMessage::user("Hello"), ChatMessage::assistant("Hi there!"), ChatMessage::user("How are you?"), ChatMessage::assistant("I'm good!"), ]; thread.restore_from_messages(messages); assert_eq!(thread.turns.len(), 2); assert_eq!(thread.turns[0].user_input, "Hello"); assert_eq!(thread.turns[0].response, Some("Hi there!".to_string())); assert_eq!(thread.turns[1].user_input, "How are you?"); assert_eq!(thread.turns[1].response, Some("I'm good!".to_string())); assert_eq!(thread.state, ThreadState::Idle); } #[test] fn test_restore_from_messages_incomplete_turn() { let mut thread = Thread::new(Uuid::new_v4()); // Messages with incomplete last turn (no assistant response) let messages = vec![ ChatMessage::user("Hello"), ChatMessage::assistant("Hi there!"), ChatMessage::user("How are you?"), ]; thread.restore_from_messages(messages); assert_eq!(thread.turns.len(), 2); assert_eq!(thread.turns[1].user_input, "How are you?"); assert!(thread.turns[1].response.is_none()); } #[test] fn test_enter_auth_mode() { let mut thread = Thread::new(Uuid::new_v4()); assert!(thread.pending_auth.is_none()); thread.enter_auth_mode("telegram".to_string()); assert!(thread.pending_auth.is_some()); assert_eq!( thread.pending_auth.as_ref().unwrap().extension_name, "telegram" ); } #[test] fn test_take_pending_auth() { let mut thread = Thread::new(Uuid::new_v4()); thread.enter_auth_mode("notion".to_string()); let pending = thread.take_pending_auth(); assert!(pending.is_some()); assert_eq!(pending.unwrap().extension_name, "notion"); // Should be cleared after take assert!(thread.pending_auth.is_none()); assert!(thread.take_pending_auth().is_none()); } #[test] fn test_pending_auth_serialization() { let mut thread = Thread::new(Uuid::new_v4()); thread.enter_auth_mode("openai".to_string()); let json = serde_json::to_string(&thread).expect("should serialize"); assert!(json.contains("pending_auth")); assert!(json.contains("openai")); let restored: Thread = serde_json::from_str(&json).expect("should deserialize"); assert!(restored.pending_auth.is_some()); assert_eq!(restored.pending_auth.unwrap().extension_name, "openai"); } #[test] fn test_pending_auth_default_none() { // Deserialization of old data without pending_auth should default to None let mut thread = Thread::new(Uuid::new_v4()); thread.pending_auth = None; let json = serde_json::to_string(&thread).expect("serialize"); // Remove the pending_auth field to simulate old data let json = json.replace(",\"pending_auth\":null", ""); let restored: Thread = serde_json::from_str(&json).expect("should deserialize"); assert!(restored.pending_auth.is_none()); } #[test] fn test_thread_with_id() { let specific_id = Uuid::new_v4(); let session_id = Uuid::new_v4(); let thread = Thread::with_id(specific_id, session_id); assert_eq!(thread.id, specific_id); assert_eq!(thread.session_id, session_id); assert_eq!(thread.state, ThreadState::Idle); assert!(thread.turns.is_empty()); } #[test] fn test_thread_with_id_restore_messages() { let thread_id = Uuid::new_v4(); let session_id = Uuid::new_v4(); let mut thread = Thread::with_id(thread_id, session_id); let messages = vec![ ChatMessage::user("Hello from DB"), ChatMessage::assistant("Restored response"), ]; thread.restore_from_messages(messages); assert_eq!(thread.id, thread_id); assert_eq!(thread.turns.len(), 1); assert_eq!(thread.turns[0].user_input, "Hello from DB"); assert_eq!( thread.turns[0].response, Some("Restored response".to_string()) ); } #[test] fn test_restore_from_messages_empty() { let mut thread = Thread::new(Uuid::new_v4()); // Add a turn first, then restore with empty vec thread.start_turn("hello"); thread.complete_turn("hi"); assert_eq!(thread.turns.len(), 1); thread.restore_from_messages(Vec::new()); // Should clear all turns and stay idle assert!(thread.turns.is_empty()); assert_eq!(thread.state, ThreadState::Idle); } #[test] fn test_restore_from_messages_only_assistant_messages() { let mut thread = Thread::new(Uuid::new_v4()); // Only assistant messages (no user messages to anchor turns) let messages = vec![ ChatMessage::assistant("I'm here"), ChatMessage::assistant("Still here"), ]; thread.restore_from_messages(messages); // Assistant-only messages have no user turn to attach to, so // they should be skipped entirely. assert!(thread.turns.is_empty()); } #[test] fn test_restore_from_messages_multiple_user_messages_in_a_row() { let mut thread = Thread::new(Uuid::new_v4()); // Two user messages with no assistant response between them let messages = vec![ ChatMessage::user("first"), ChatMessage::user("second"), ChatMessage::assistant("reply to second"), ]; thread.restore_from_messages(messages); // First user message becomes a turn with no response, // second user message pairs with the assistant response. assert_eq!(thread.turns.len(), 2); assert_eq!(thread.turns[0].user_input, "first"); assert!(thread.turns[0].response.is_none()); assert_eq!(thread.turns[1].user_input, "second"); assert_eq!( thread.turns[1].response, Some("reply to second".to_string()) ); } #[test] fn test_thread_switch() { let mut session = Session::new("user-1"); let t1_id = session.create_thread().id; let t2_id = session.create_thread().id; // After creating two threads, active should be the last one assert_eq!(session.active_thread, Some(t2_id)); // Switch back to the first assert!(session.switch_thread(t1_id)); assert_eq!(session.active_thread, Some(t1_id)); // Switching to a nonexistent thread should fail let fake_id = Uuid::new_v4(); assert!(!session.switch_thread(fake_id)); // Active thread should remain unchanged assert_eq!(session.active_thread, Some(t1_id)); } #[test] fn test_get_or_create_thread_idempotent() { let mut session = Session::new("user-1"); let tid1 = session.get_or_create_thread().id; let tid2 = session.get_or_create_thread().id; // Should return the same thread (not create a new one each time) assert_eq!(tid1, tid2); assert_eq!(session.threads.len(), 1); } #[test] fn test_truncate_turns() { let mut thread = Thread::new(Uuid::new_v4()); for i in 0..5 { thread.start_turn(format!("msg-{}", i)); thread.complete_turn(format!("resp-{}", i)); } assert_eq!(thread.turns.len(), 5); thread.truncate_turns(3); assert_eq!(thread.turns.len(), 3); // Should keep the most recent turns assert_eq!(thread.turns[0].user_input, "msg-2"); assert_eq!(thread.turns[1].user_input, "msg-3"); assert_eq!(thread.turns[2].user_input, "msg-4"); // Turn numbers should be re-indexed assert_eq!(thread.turns[0].turn_number, 0); assert_eq!(thread.turns[1].turn_number, 1); assert_eq!(thread.turns[2].turn_number, 2); } #[test] fn test_truncate_turns_noop_when_fewer() { let mut thread = Thread::new(Uuid::new_v4()); thread.start_turn("only one"); thread.complete_turn("response"); thread.truncate_turns(10); assert_eq!(thread.turns.len(), 1); assert_eq!(thread.turns[0].user_input, "only one"); } #[test] fn test_thread_interrupt_and_resume() { let mut thread = Thread::new(Uuid::new_v4()); thread.start_turn("do something"); assert_eq!(thread.state, ThreadState::Processing); thread.interrupt(); assert_eq!(thread.state, ThreadState::Interrupted); let last_turn = thread.last_turn().unwrap(); assert_eq!(last_turn.state, TurnState::Interrupted); assert!(last_turn.completed_at.is_some()); thread.resume(); assert_eq!(thread.state, ThreadState::Idle); } #[test] fn test_resume_only_from_interrupted() { let mut thread = Thread::new(Uuid::new_v4()); // Idle thread: resume should be a no-op assert_eq!(thread.state, ThreadState::Idle); thread.resume(); assert_eq!(thread.state, ThreadState::Idle); // Processing thread: resume should not change state thread.start_turn("work"); assert_eq!(thread.state, ThreadState::Processing); thread.resume(); assert_eq!(thread.state, ThreadState::Processing); } #[test] fn test_turn_fail() { let mut thread = Thread::new(Uuid::new_v4()); thread.start_turn("risky operation"); thread.fail_turn("connection timed out"); assert_eq!(thread.state, ThreadState::Idle); let turn = thread.last_turn().unwrap(); assert_eq!(turn.state, TurnState::Failed); assert_eq!(turn.error, Some("connection timed out".to_string())); assert!(turn.response.is_none()); assert!(turn.completed_at.is_some()); } #[test] fn test_messages_with_incomplete_last_turn() { let mut thread = Thread::new(Uuid::new_v4()); thread.start_turn("first"); thread.complete_turn("first reply"); thread.start_turn("second (in progress)"); let messages = thread.messages(); // Should have 3 messages: user, assistant, user (no assistant for in-progress) assert_eq!(messages.len(), 3); assert_eq!(messages[0].content, "first"); assert_eq!(messages[1].content, "first reply"); assert_eq!(messages[2].content, "second (in progress)"); } #[test] fn test_thread_serialization_round_trip() { let mut thread = Thread::new(Uuid::new_v4()); thread.start_turn("hello"); thread.complete_turn("world"); let json = serde_json::to_string(&thread).unwrap(); let restored: Thread = serde_json::from_str(&json).unwrap(); assert_eq!(restored.id, thread.id); assert_eq!(restored.session_id, thread.session_id); assert_eq!(restored.turns.len(), 1); assert_eq!(restored.turns[0].user_input, "hello"); assert_eq!(restored.turns[0].response, Some("world".to_string())); } #[test] fn test_session_serialization_round_trip() { let mut session = Session::new("user-ser"); session.create_thread(); session.auto_approve_tool("echo"); let json = serde_json::to_string(&session).unwrap(); let restored: Session = serde_json::from_str(&json).unwrap(); assert_eq!(restored.user_id, "user-ser"); assert_eq!(restored.threads.len(), 1); assert!(restored.is_tool_auto_approved("echo")); assert!(!restored.is_tool_auto_approved("shell")); } #[test] fn test_auto_approved_tools() { let mut session = Session::new("user-1"); assert!(!session.is_tool_auto_approved("shell")); session.auto_approve_tool("shell"); assert!(session.is_tool_auto_approved("shell")); // Idempotent session.auto_approve_tool("shell"); assert_eq!(session.auto_approved_tools.len(), 1); } #[test] fn test_turn_tool_call_error() { let mut turn = Turn::new(0, "test"); turn.record_tool_call("http", serde_json::json!({"url": "example.com"})); turn.record_tool_error("timeout"); assert_eq!(turn.tool_calls.len(), 1); assert_eq!(turn.tool_calls[0].error, Some("timeout".to_string())); assert!(turn.tool_calls[0].result.is_none()); } #[test] fn test_turn_number_increments() { let mut thread = Thread::new(Uuid::new_v4()); // Before any turns, turn_number() is 1 (1-indexed for display) assert_eq!(thread.turn_number(), 1); thread.start_turn("first"); thread.complete_turn("done"); assert_eq!(thread.turn_number(), 2); thread.start_turn("second"); assert_eq!(thread.turn_number(), 3); } #[test] fn test_complete_turn_on_empty_thread() { let mut thread = Thread::new(Uuid::new_v4()); // Completing a turn when there are no turns should be a safe no-op thread.complete_turn("phantom response"); assert_eq!(thread.state, ThreadState::Idle); assert!(thread.turns.is_empty()); } #[test] fn test_fail_turn_on_empty_thread() { let mut thread = Thread::new(Uuid::new_v4()); // Failing a turn when there are no turns should be a safe no-op thread.fail_turn("phantom error"); assert_eq!(thread.state, ThreadState::Idle); assert!(thread.turns.is_empty()); } #[test] fn test_pending_approval_flow() { let mut thread = Thread::new(Uuid::new_v4()); let approval = PendingApproval { request_id: Uuid::new_v4(), tool_name: "shell".to_string(), parameters: serde_json::json!({"command": "rm -rf /"}), display_parameters: serde_json::json!({"command": "rm -rf /"}), description: "dangerous command".to_string(), tool_call_id: "call_123".to_string(), context_messages: vec![ChatMessage::user("do it")], deferred_tool_calls: vec![], user_timezone: None, }; thread.await_approval(approval); assert_eq!(thread.state, ThreadState::AwaitingApproval); assert!(thread.pending_approval.is_some()); let taken = thread.take_pending_approval(); assert!(taken.is_some()); assert_eq!(taken.unwrap().tool_name, "shell"); assert!(thread.pending_approval.is_none()); } #[test] fn test_clear_pending_approval() { let mut thread = Thread::new(Uuid::new_v4()); let approval = PendingApproval { request_id: Uuid::new_v4(), tool_name: "http".to_string(), parameters: serde_json::json!({}), display_parameters: serde_json::json!({}), description: "test".to_string(), tool_call_id: "call_456".to_string(), context_messages: vec![], deferred_tool_calls: vec![], user_timezone: None, }; thread.await_approval(approval); thread.clear_pending_approval(); assert_eq!(thread.state, ThreadState::Idle); assert!(thread.pending_approval.is_none()); } #[test] fn test_active_thread_accessors() { let mut session = Session::new("user-1"); assert!(session.active_thread().is_none()); assert!(session.active_thread_mut().is_none()); let tid = session.create_thread().id; assert!(session.active_thread().is_some()); assert_eq!(session.active_thread().unwrap().id, tid); // Mutably modify through accessor session.active_thread_mut().unwrap().start_turn("test"); assert_eq!( session.active_thread().unwrap().state, ThreadState::Processing ); } // Regression tests for #568: tool call history must survive hydration. #[test] fn test_messages_includes_tool_calls() { let mut thread = Thread::new(Uuid::new_v4()); thread.start_turn("Search for X"); { let turn = thread.turns.last_mut().unwrap(); turn.record_tool_call("memory_search", serde_json::json!({"query": "X"})); turn.record_tool_result(serde_json::json!("Found X in doc.md")); } thread.complete_turn("I found X in doc.md."); let messages = thread.messages(); // user + assistant_with_tool_calls + tool_result + assistant = 4 assert_eq!(messages.len(), 4); assert_eq!(messages[0].role, crate::llm::Role::User); assert_eq!(messages[0].content, "Search for X"); assert_eq!(messages[1].role, crate::llm::Role::Assistant); assert!(messages[1].tool_calls.is_some()); let tcs = messages[1].tool_calls.as_ref().unwrap(); assert_eq!(tcs.len(), 1); assert_eq!(tcs[0].name, "memory_search"); assert_eq!(messages[2].role, crate::llm::Role::Tool); assert!(messages[2].content.contains("Found X")); assert_eq!(messages[3].role, crate::llm::Role::Assistant); assert_eq!(messages[3].content, "I found X in doc.md."); } #[test] fn test_messages_multiple_tool_calls_per_turn() { let mut thread = Thread::new(Uuid::new_v4()); thread.start_turn("Do two things"); { let turn = thread.turns.last_mut().unwrap(); turn.record_tool_call("echo", serde_json::json!({"msg": "a"})); turn.record_tool_result(serde_json::json!("a")); turn.record_tool_call("time", serde_json::json!({})); turn.record_tool_error("timeout"); } thread.complete_turn("Done."); let messages = thread.messages(); // user + assistant_with_calls(2) + tool_result + tool_result + assistant = 5 assert_eq!(messages.len(), 5); let tcs = messages[1].tool_calls.as_ref().unwrap(); assert_eq!(tcs.len(), 2); // First tool: success assert_eq!(messages[2].content, "a"); // Second tool: error (passed through directly, no wrapping) assert!(messages[3].content.contains("timeout")); } #[test] fn test_restore_from_messages_with_tool_calls() { let mut thread = Thread::new(Uuid::new_v4()); // Build a message sequence with tool calls let tc = ToolCall { id: "call_0".to_string(), name: "search".to_string(), arguments: serde_json::json!({"q": "test"}), }; let messages = vec![ ChatMessage::user("Find test"), ChatMessage::assistant_with_tool_calls(None, vec![tc]), ChatMessage::tool_result("call_0", "search", "result: found"), ChatMessage::assistant("Found it."), ]; thread.restore_from_messages(messages); assert_eq!(thread.turns.len(), 1); let turn = &thread.turns[0]; assert_eq!(turn.user_input, "Find test"); assert_eq!(turn.tool_calls.len(), 1); assert_eq!(turn.tool_calls[0].name, "search"); assert_eq!( turn.tool_calls[0].result, Some(serde_json::Value::String("result: found".to_string())) ); assert_eq!(turn.response, Some("Found it.".to_string())); } #[test] fn test_restore_from_messages_with_tool_error() { let mut thread = Thread::new(Uuid::new_v4()); let tc = ToolCall { id: "call_0".to_string(), name: "http".to_string(), arguments: serde_json::json!({}), }; let messages = vec![ ChatMessage::user("Fetch URL"), ChatMessage::assistant_with_tool_calls(None, vec![tc]), ChatMessage::tool_result("call_0", "http", "Error: timeout"), ChatMessage::assistant("The request timed out."), ]; thread.restore_from_messages(messages); // restore_from_messages stores all tool content as result (not error), // because it can't reliably distinguish errors from results that happen // to start with "Error: ". The content is preserved for LLM context. let turn = &thread.turns[0]; assert_eq!( turn.tool_calls[0].result, Some(serde_json::Value::String("Error: timeout".to_string())) ); } #[test] fn test_messages_round_trip_with_tools() { // Build a thread with tool calls, get messages(), restore, get messages() again // The two message sequences should be equivalent. let mut thread = Thread::new(Uuid::new_v4()); thread.start_turn("Do search"); { let turn = thread.turns.last_mut().unwrap(); turn.record_tool_call("search", serde_json::json!({"q": "test"})); turn.record_tool_result(serde_json::json!("found")); } thread.complete_turn("Here are results."); let messages_original = thread.messages(); // Restore into a new thread let mut thread2 = Thread::new(Uuid::new_v4()); thread2.restore_from_messages(messages_original.clone()); let messages_restored = thread2.messages(); // Same number of messages assert_eq!(messages_original.len(), messages_restored.len()); // Same roles for (orig, rest) in messages_original.iter().zip(messages_restored.iter()) { assert_eq!(orig.role, rest.role); } // Same final response assert_eq!( messages_original.last().unwrap().content, messages_restored.last().unwrap().content ); } #[test] fn test_restore_multi_stage_tool_calls() { let mut thread = Thread::new(Uuid::new_v4()); let tc1 = ToolCall { id: "call_a".to_string(), name: "search".to_string(), arguments: serde_json::json!({"q": "data"}), }; let tc2 = ToolCall { id: "call_b".to_string(), name: "write".to_string(), arguments: serde_json::json!({"path": "out.txt"}), }; let messages = vec![ ChatMessage::user("Find and save"), ChatMessage::assistant_with_tool_calls(None, vec![tc1]), ChatMessage::tool_result("call_a", "search", "found data"), ChatMessage::assistant_with_tool_calls(None, vec![tc2]), ChatMessage::tool_result("call_b", "write", "written"), ChatMessage::assistant("Done, saved to out.txt"), ]; thread.restore_from_messages(messages); assert_eq!(thread.turns.len(), 1); let turn = &thread.turns[0]; assert_eq!(turn.tool_calls.len(), 2); assert_eq!(turn.tool_calls[0].name, "search"); assert_eq!(turn.tool_calls[1].name, "write"); assert_eq!( turn.tool_calls[0].result, Some(serde_json::Value::String("found data".to_string())) ); assert_eq!( turn.tool_calls[1].result, Some(serde_json::Value::String("written".to_string())) ); assert_eq!(turn.response, Some("Done, saved to out.txt".to_string())); } #[test] fn test_messages_truncates_large_tool_results() { let mut thread = Thread::new(Uuid::new_v4()); thread.start_turn("Read big file"); { let turn = thread.turns.last_mut().unwrap(); turn.record_tool_call("read_file", serde_json::json!({"path": "big.txt"})); let big_result = "x".repeat(2000); turn.record_tool_result(serde_json::json!(big_result)); } thread.complete_turn("Here's the file content."); let messages = thread.messages(); let tool_result_content = &messages[2].content; assert!( tool_result_content.len() <= 1010, "Tool result should be truncated, got {} chars", tool_result_content.len() ); assert!(tool_result_content.ends_with("...")); } }