diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index a0e8278f..a9a3481d 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -1003,15 +1003,22 @@ impl Agent { } } - // Hydrate thread from DB if it's a historical thread not in memory + // Hydrate thread from DB if it's a historical thread not in memory. + // Capture the parsed UUID to avoid redundant re-parsing downstream. + let mut parsed_thread_uuid: Option = None; if let Some(external_thread_id) = message.conversation_scope() { tracing::trace!( message_id = %message.id, thread_id = %external_thread_id, "Hydrating thread from DB" ); - if let Some(rejection) = self.maybe_hydrate_thread(message, external_thread_id).await { - return Ok(Some(format!("Error: {}", rejection))); + match self.maybe_hydrate_thread(message, external_thread_id).await { + Ok(uuid) => { + parsed_thread_uuid = uuid; + } + Err(rejection) => { + return Ok(Some(format!("Error: {}", rejection))); + } } } @@ -1023,9 +1030,7 @@ impl Agent { submission, Submission::ExecApproval { .. } | Submission::ApprovalResponse { .. } ) { - message - .conversation_scope() - .and_then(|thread_id| Uuid::parse_str(thread_id).ok()) + parsed_thread_uuid } else { None }; @@ -1052,19 +1057,21 @@ impl Agent { } else { drop(sess); self.session_manager - .resolve_thread( + .resolve_thread_with_parsed_uuid( &message.user_id, &message.channel, message.conversation_scope(), + parsed_thread_uuid, ) .await } } else { self.session_manager - .resolve_thread( + .resolve_thread_with_parsed_uuid( &message.user_id, &message.channel, message.conversation_scope(), + parsed_thread_uuid, ) .await }; diff --git a/src/agent/session_manager.rs b/src/agent/session_manager.rs index 3bf20697..94b85bf9 100644 --- a/src/agent/session_manager.rs +++ b/src/agent/session_manager.rs @@ -108,81 +108,81 @@ impl SessionManager { channel: &str, external_thread_id: Option<&str>, ) -> (Arc>, Uuid) { - let session = self.get_or_create_session(user_id).await; + let parsed_uuid = external_thread_id.and_then(|tid| Uuid::parse_str(tid).ok()); + self.resolve_thread_inner(user_id, channel, external_thread_id, parsed_uuid) + .await + } + /// Resolve an external thread ID to an internal thread, accepting a + /// pre-parsed UUID to avoid redundant parsing in hot paths. + pub async fn resolve_thread_with_parsed_uuid( + &self, + user_id: &str, + channel: &str, + external_thread_id: Option<&str>, + parsed_uuid: Option, + ) -> (Arc>, Uuid) { + self.resolve_thread_inner(user_id, channel, external_thread_id, parsed_uuid) + .await + } + + /// Inner implementation for thread resolution. + async fn resolve_thread_inner( + &self, + user_id: &str, + channel: &str, + external_thread_id: Option<&str>, + parsed_uuid: Option, + ) -> (Arc>, Uuid) { + let session = self.get_or_create_session(user_id).await; let key = ThreadKey { user_id: user_id.to_string(), channel: channel.to_string(), external_thread_id: external_thread_id.map(String::from), }; - - // Check if we have a mapping { let thread_map = self.thread_map.read().await; if let Some(&thread_id) = thread_map.get(&key) { - // Verify thread still exists in session let sess = session.lock().await; if sess.threads.contains_key(&thread_id) { return (Arc::clone(&session), thread_id); } } } - - // Check if external_thread_id is itself a known thread UUID that - // exists in the session but was never registered in the thread_map - // (e.g. created by chat_new_thread_handler or hydrated from DB). - // We only adopt it if no thread_map entry maps to this UUID — - // otherwise it belongs to a different channel scope. - if let Some(ext_tid) = external_thread_id - && let Ok(ext_uuid) = Uuid::parse_str(ext_tid) - { + if let Some(ext_uuid) = parsed_uuid { let thread_map = self.thread_map.read().await; let mapped_elsewhere = thread_map.values().any(|&v| v == ext_uuid); drop(thread_map); - if !mapped_elsewhere { let sess = session.lock().await; if sess.threads.contains_key(&ext_uuid) { drop(sess); - let mut thread_map = self.thread_map.write().await; - // Re-check after acquiring write lock to prevent race condition - // where another task mapped this UUID between our read and write. if !thread_map.values().any(|&v| v == ext_uuid) { thread_map.insert(key, ext_uuid); drop(thread_map); - // Ensure undo manager exists let mut undo_managers = self.undo_managers.write().await; undo_managers .entry(ext_uuid) .or_insert_with(|| Arc::new(Mutex::new(UndoManager::new()))); return (session, ext_uuid); } - // If it was mapped elsewhere while we were unlocked, fall through - // to create a new thread, preserving channel isolation. } } } - - // Create new thread (always create a new one for a new key) let thread_id = { let mut sess = session.lock().await; let thread = sess.create_thread(); thread.id }; - - // Store mapping { let mut thread_map = self.thread_map.write().await; thread_map.insert(key, thread_id); } - - // Create undo manager for thread { let mut undo_managers = self.undo_managers.write().await; undo_managers.insert(thread_id, Arc::new(Mutex::new(UndoManager::new()))); } - (session, thread_id) } @@ -947,4 +947,33 @@ mod tests { "should have exactly 1 thread, not a duplicate" ); } + + #[tokio::test] + async fn test_resolve_thread_with_parsed_uuid_matches_resolve_thread() { + use crate::agent::session::{Session, Thread}; + let manager = SessionManager::new(); + let tid = Uuid::new_v4(); + let session = Arc::new(Mutex::new(Session::new("user-parsed"))); + { + let mut sess = session.lock().await; + let thread = Thread::with_id(tid, sess.id); + sess.threads.insert(tid, thread); + } + { + let mut sessions = manager.sessions.write().await; + sessions.insert("user-parsed".to_string(), Arc::clone(&session)); + } + let tid_str = tid.to_string(); + let (_, resolved_normal) = manager + .resolve_thread("user-parsed", "gateway", Some(&tid_str)) + .await; + let (_, resolved_parsed) = manager + .resolve_thread_with_parsed_uuid("user-parsed", "gateway", Some(&tid_str), Some(tid)) + .await; + assert_eq!( + resolved_normal, resolved_parsed, + "resolve_thread and resolve_thread_with_parsed_uuid must return the same thread" + ); + assert_eq!(resolved_normal, tid); + } } diff --git a/src/agent/thread_ops.rs b/src/agent/thread_ops.rs index 0fb968f1..0e8df7b7 100644 --- a/src/agent/thread_ops.rs +++ b/src/agent/thread_ops.rs @@ -45,11 +45,11 @@ impl Agent { &self, message: &IncomingMessage, external_thread_id: &str, - ) -> Option { + ) -> Result, String> { // Only hydrate UUID-shaped thread IDs (web gateway uses UUIDs) let thread_uuid = match Uuid::parse_str(external_thread_id) { Ok(id) => id, - Err(_) => return None, + Err(_) => return Ok(None), }; // Check if already in memory @@ -60,7 +60,7 @@ impl Agent { { let sess = session.lock().await; if sess.threads.contains_key(&thread_uuid) { - return None; + return Ok(Some(thread_uuid)); } } @@ -83,9 +83,9 @@ impl Agent { e ); if requires_preexisting_uuid_thread(&message.channel) { - return Some(FORGED_THREAD_ID_ERROR.to_string()); + return Err(FORGED_THREAD_ID_ERROR.to_string()); } - return None; + return Ok(Some(thread_uuid)); } }; if !owned { @@ -99,9 +99,9 @@ impl Agent { e ); if requires_preexisting_uuid_thread(&message.channel) { - return Some(FORGED_THREAD_ID_ERROR.to_string()); + return Err(FORGED_THREAD_ID_ERROR.to_string()); } - return None; + return Ok(Some(thread_uuid)); } }; @@ -113,7 +113,7 @@ impl Agent { exists, "Rejected message for unavailable thread id" ); - return Some(FORGED_THREAD_ID_ERROR.to_string()); + return Err(FORGED_THREAD_ID_ERROR.to_string()); } tracing::warn!( @@ -122,7 +122,7 @@ impl Agent { exists, "Skipped hydration for thread id not owned by sender" ); - return None; + return Ok(Some(thread_uuid)); } let db_messages = store @@ -169,7 +169,7 @@ impl Agent { msg_count ); - None + Ok(Some(thread_uuid)) } pub(super) async fn process_user_input(