perf(agent): optimize approval thread resolution (UUID parsing + lock contention) (#1592)

src/agent/agent_loop.rs

@@ -1055,10 +1055,11 @@ impl Agent {
             } else {
                 drop(sess);
                 self.session_manager
-                    .resolve_thread(
+                    .resolve_thread_with_parsed_uuid(
                         &message.user_id,
                         &message.channel,
                         message.conversation_scope(),
+                        approval_thread_uuid,
                     )
                     .await
             }

src/agent/session_manager.rs

@@ -102,11 +102,30 @@ impl SessionManager {
     /// Resolve an external thread ID to an internal thread.
     ///
     /// Returns the session and thread ID. Creates both if they don't exist.
+    /// Delegates to [`resolve_thread_with_parsed_uuid`](Self::resolve_thread_with_parsed_uuid)
+    /// with `parsed_uuid: None`.
     pub async fn resolve_thread(
         &self,
         user_id: &str,
         channel: &str,
         external_thread_id: Option<&str>,
+    ) -> (Arc<Mutex<Session>>, Uuid) {
+        self.resolve_thread_with_parsed_uuid(user_id, channel, external_thread_id, None)
+            .await
+    }
+
+    /// Like [`resolve_thread`](Self::resolve_thread), but accepts a pre-parsed
+    /// UUID to skip redundant parsing when the caller has already validated
+    /// the external thread ID as a UUID (e.g. the approval routing path).
+    ///
+    /// Uses a single read-lock acquisition for both the key lookup and the UUID
+    /// adoption check to reduce contention under concurrent approval load.
+    pub async fn resolve_thread_with_parsed_uuid(
+        &self,
+        user_id: &str,
+        channel: &str,
+        external_thread_id: Option<&str>,
+        parsed_uuid: Option<Uuid>,
     ) -> (Arc<Mutex<Session>>, Uuid) {
         let session = self.get_or_create_session(user_id).await;
 
@@ -116,51 +135,65 @@ impl SessionManager {
             external_thread_id: external_thread_id.map(String::from),
         };
 
-        // Check if we have a mapping
-        {
+        // Use pre-parsed UUID if available, otherwise parse from string.
+        let ext_uuid = parsed_uuid
+            .or_else(|| external_thread_id.and_then(|ext_tid| Uuid::parse_str(ext_tid).ok()));
+
+        // Validate that parsed_uuid (if provided) is consistent with external_thread_id.
+        #[cfg(debug_assertions)]
+        if let (Some(parsed), Some(ext_tid)) = (&parsed_uuid, external_thread_id) {
+            debug_assert_eq!(
+                Uuid::parse_str(ext_tid).ok().as_ref(),
+                Some(parsed),
+                "parsed_uuid must be the parsed form of external_thread_id"
+            );
+        }
+
+        // Single read lock for both the key lookup and UUID adoption check
+        let adoptable_uuid = {
             let thread_map = self.thread_map.read().await;
+
+            // Fast path: exact key match
             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)
-        {
-            let thread_map = self.thread_map.read().await;
-            let mapped_elsewhere = thread_map.values().any(|&v| v == ext_uuid);
-            drop(thread_map);
+            // UUID adoption check (still under the same read lock).
+            // If external_thread_id is a valid UUID not mapped elsewhere,
+            // it may be a thread created by chat_new_thread_handler or
+            // hydrated from DB that we can adopt.
+            // Only attempt adoption when external_thread_id is Some, preserving
+            // the invariant that None external_thread_id never triggers adoption.
+            if external_thread_id.is_some() {
+                ext_uuid.filter(|&uuid| !thread_map.values().any(|&v| v == uuid))
+            } else {
+                None
+            }
+        }; // Single read lock dropped here
 
-            if !mapped_elsewhere {
-                let sess = session.lock().await;
-                if sess.threads.contains_key(&ext_uuid) {
-                    drop(sess);
+        // If we found an adoptable UUID, verify it exists in session and acquire write lock
+        if let Some(ext_uuid) = adoptable_uuid {
+            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.
+                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 mapped elsewhere while unlocked, fall through to create new thread
             }
         }
 
@@ -909,6 +942,44 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn test_resolve_thread_consolidates_read_path() {
+        // Verify that resolve_thread still correctly handles:
+        // 1. Fast path: key exists in thread_map
+        // 2. UUID adoption: external_thread_id is a UUID in session but not in map
+        // 3. New thread: neither path matches
+        use crate::agent::session::Thread;
+
+        let manager = SessionManager::new();
+
+        // Case 1: Normal resolution creates thread and maps it
+        let (session1, tid1) = manager
+            .resolve_thread("user1", "chan1", Some("ext-1"))
+            .await;
+        // Resolving again with same key should return same thread (fast path)
+        let (_, tid1_again) = manager
+            .resolve_thread("user1", "chan1", Some("ext-1"))
+            .await;
+        assert_eq!(tid1, tid1_again);
+
+        // Case 2: UUID adoption - insert a thread directly into session
+        let adopted_id = Uuid::new_v4();
+        {
+            let mut sess = session1.lock().await;
+            let thread = Thread::with_id(adopted_id, sess.id);
+            sess.threads.insert(adopted_id, thread);
+        }
+        // Resolve with the UUID as external_thread_id -- should adopt it
+        let (_, resolved) = manager
+            .resolve_thread("user1", "chan1", Some(&adopted_id.to_string()))
+            .await;
+        assert_eq!(resolved, adopted_id);
+
+        // Case 3: Different channel gets different thread
+        let (_, tid2) = manager.resolve_thread("user1", "chan2", None).await;
+        assert_ne!(tid1, tid2);
+    }
+
     #[tokio::test]
     async fn test_resolve_thread_finds_existing_session_thread_by_uuid() {
         use crate::agent::session::{Session, Thread};
@@ -947,4 +1018,88 @@ mod tests {
             "should have exactly 1 thread, not a duplicate"
         );
     }
+
+    #[tokio::test]
+    async fn test_resolve_thread_with_pre_parsed_uuid_adopts_thread() {
+        use crate::agent::session::Thread;
+
+        let manager = SessionManager::new();
+        let (session, _) = manager.resolve_thread("user1", "chan1", None).await;
+
+        // Manually insert a thread with a known UUID
+        let known_id = Uuid::new_v4();
+        {
+            let mut sess = session.lock().await;
+            let thread = Thread::with_id(known_id, sess.id);
+            sess.threads.insert(known_id, thread);
+        }
+
+        // Resolve with pre-parsed UUID -- should adopt it without re-parsing
+        let (_, resolved) = manager
+            .resolve_thread_with_parsed_uuid(
+                "user1",
+                "chan1",
+                Some(&known_id.to_string()),
+                Some(known_id),
+            )
+            .await;
+        assert_eq!(resolved, known_id);
+    }
+
+    #[tokio::test]
+    async fn test_resolve_thread_with_parsed_uuid_none_delegates_to_parse() {
+        use crate::agent::session::Thread;
+
+        let manager = SessionManager::new();
+        let (session, _) = manager.resolve_thread("user2", "chan2", None).await;
+
+        // Insert a thread with a known UUID
+        let known_id = Uuid::new_v4();
+        {
+            let mut sess = session.lock().await;
+            let thread = Thread::with_id(known_id, sess.id);
+            sess.threads.insert(known_id, thread);
+        }
+
+        // Resolve with parsed_uuid=None but a valid UUID string -- should
+        // fall back to parsing the string and still adopt the thread
+        let (_, resolved) = manager
+            .resolve_thread_with_parsed_uuid("user2", "chan2", Some(&known_id.to_string()), None)
+            .await;
+        assert_eq!(resolved, known_id);
+    }
+
+    #[tokio::test]
+    async fn test_resolve_thread_with_none_external_thread_id_does_not_adopt() {
+        use crate::agent::session::Thread;
+
+        let manager = SessionManager::new();
+        let (session, default_tid) = manager.resolve_thread("user3", "chan3", None).await;
+
+        // Manually insert a thread with a known UUID (simulating a thread
+        // created by chat_new_thread_handler)
+        let known_id = Uuid::new_v4();
+        {
+            let mut sess = session.lock().await;
+            let thread = Thread::with_id(known_id, sess.id);
+            sess.threads.insert(known_id, thread);
+        }
+
+        // Resolve with external_thread_id=None but parsed_uuid=Some.
+        // This should NOT adopt the UUID — the old code prevented adoption
+        // when external_thread_id was None, and we preserve that invariant.
+        let (_, resolved) = manager
+            .resolve_thread_with_parsed_uuid("user3", "chan3", None, Some(known_id))
+            .await;
+
+        // Should return the existing default thread, not the injected UUID
+        assert_eq!(
+            resolved, default_tid,
+            "should return existing default thread when external_thread_id is None"
+        );
+        assert_ne!(
+            resolved, known_id,
+            "should NOT adopt UUID when external_thread_id is None"
+        );
+    }
 }