Compare commits

...
Author SHA1 Message Date
ZakiandClaude Opus 4.6 09e59c7c28 style(agent): use match as expression for parsed_thread_uuid
Refactor mutable variable + match arm assignment into idiomatic
Rust match-as-expression pattern. Also remove unused uuid::Uuid
import that was left behind after the parse-once optimization.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-22 18:28:09 -07:00
ZakiandClaude Opus 4.6 9cff809950 perf(agent): parse external thread UUID once in resolution path (#1488)
Eliminates redundant UUID parsing in the thread resolution hot path by
having maybe_hydrate_thread return the parsed UUID and threading it
through to resolve_thread via a new resolve_thread_with_parsed_uuid
method.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-21 10:35:25 -07:00
3 changed files with 82 additions and 48 deletions
+16 -11
View File
@@ -10,7 +10,6 @@
use std::sync::Arc; use std::sync::Arc;
use futures::StreamExt; use futures::StreamExt;
use uuid::Uuid;
use crate::agent::context_monitor::ContextMonitor; use crate::agent::context_monitor::ContextMonitor;
use crate::agent::heartbeat::spawn_heartbeat; use crate::agent::heartbeat::spawn_heartbeat;
@@ -1003,17 +1002,23 @@ 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.
if let Some(external_thread_id) = message.conversation_scope() { // Capture the parsed UUID to avoid redundant re-parsing downstream.
let parsed_thread_uuid = if let Some(external_thread_id) = message.conversation_scope() {
tracing::trace!( tracing::trace!(
message_id = %message.id, message_id = %message.id,
thread_id = %external_thread_id, thread_id = %external_thread_id,
"Hydrating thread from DB" "Hydrating thread from DB"
); );
if let Some(rejection) = self.maybe_hydrate_thread(message, external_thread_id).await { match self.maybe_hydrate_thread(message, external_thread_id).await {
return Ok(Some(format!("Error: {}", rejection))); Ok(uuid) => uuid,
Err(rejection) => {
return Ok(Some(format!("Error: {}", rejection)));
}
} }
} } else {
None
};
// Resolve session and thread. Approval submissions are allowed to // Resolve session and thread. Approval submissions are allowed to
// target an already-loaded owned thread by UUID across channels so the // target an already-loaded owned thread by UUID across channels so the
@@ -1023,9 +1028,7 @@ impl Agent {
submission, submission,
Submission::ExecApproval { .. } | Submission::ApprovalResponse { .. } Submission::ExecApproval { .. } | Submission::ApprovalResponse { .. }
) { ) {
message parsed_thread_uuid
.conversation_scope()
.and_then(|thread_id| Uuid::parse_str(thread_id).ok())
} else { } else {
None None
}; };
@@ -1052,19 +1055,21 @@ impl Agent {
} else { } else {
drop(sess); drop(sess);
self.session_manager self.session_manager
.resolve_thread( .resolve_thread_with_parsed_uuid(
&message.user_id, &message.user_id,
&message.channel, &message.channel,
message.conversation_scope(), message.conversation_scope(),
parsed_thread_uuid,
) )
.await .await
} }
} else { } else {
self.session_manager self.session_manager
.resolve_thread( .resolve_thread_with_parsed_uuid(
&message.user_id, &message.user_id,
&message.channel, &message.channel,
message.conversation_scope(), message.conversation_scope(),
parsed_thread_uuid,
) )
.await .await
}; };
+56 -27
View File
@@ -108,81 +108,81 @@ impl SessionManager {
channel: &str, channel: &str,
external_thread_id: Option<&str>, external_thread_id: Option<&str>,
) -> (Arc<Mutex<Session>>, Uuid) { ) -> (Arc<Mutex<Session>>, 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<Uuid>,
) -> (Arc<Mutex<Session>>, 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<Uuid>,
) -> (Arc<Mutex<Session>>, Uuid) {
let session = self.get_or_create_session(user_id).await;
let key = ThreadKey { let key = ThreadKey {
user_id: user_id.to_string(), user_id: user_id.to_string(),
channel: channel.to_string(), channel: channel.to_string(),
external_thread_id: external_thread_id.map(String::from), external_thread_id: external_thread_id.map(String::from),
}; };
// Check if we have a mapping
{ {
let thread_map = self.thread_map.read().await; let thread_map = self.thread_map.read().await;
if let Some(&thread_id) = thread_map.get(&key) { if let Some(&thread_id) = thread_map.get(&key) {
// Verify thread still exists in session
let sess = session.lock().await; let sess = session.lock().await;
if sess.threads.contains_key(&thread_id) { if sess.threads.contains_key(&thread_id) {
return (Arc::clone(&session), thread_id); return (Arc::clone(&session), thread_id);
} }
} }
} }
if let Some(ext_uuid) = parsed_uuid {
// 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 thread_map = self.thread_map.read().await;
let mapped_elsewhere = thread_map.values().any(|&v| v == ext_uuid); let mapped_elsewhere = thread_map.values().any(|&v| v == ext_uuid);
drop(thread_map); drop(thread_map);
if !mapped_elsewhere { if !mapped_elsewhere {
let sess = session.lock().await; let sess = session.lock().await;
if sess.threads.contains_key(&ext_uuid) { if sess.threads.contains_key(&ext_uuid) {
drop(sess); drop(sess);
let mut thread_map = self.thread_map.write().await; 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) { if !thread_map.values().any(|&v| v == ext_uuid) {
thread_map.insert(key, ext_uuid); thread_map.insert(key, ext_uuid);
drop(thread_map); drop(thread_map);
// Ensure undo manager exists
let mut undo_managers = self.undo_managers.write().await; let mut undo_managers = self.undo_managers.write().await;
undo_managers undo_managers
.entry(ext_uuid) .entry(ext_uuid)
.or_insert_with(|| Arc::new(Mutex::new(UndoManager::new()))); .or_insert_with(|| Arc::new(Mutex::new(UndoManager::new())));
return (session, ext_uuid); 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 thread_id = {
let mut sess = session.lock().await; let mut sess = session.lock().await;
let thread = sess.create_thread(); let thread = sess.create_thread();
thread.id thread.id
}; };
// Store mapping
{ {
let mut thread_map = self.thread_map.write().await; let mut thread_map = self.thread_map.write().await;
thread_map.insert(key, thread_id); thread_map.insert(key, thread_id);
} }
// Create undo manager for thread
{ {
let mut undo_managers = self.undo_managers.write().await; let mut undo_managers = self.undo_managers.write().await;
undo_managers.insert(thread_id, Arc::new(Mutex::new(UndoManager::new()))); undo_managers.insert(thread_id, Arc::new(Mutex::new(UndoManager::new())));
} }
(session, thread_id) (session, thread_id)
} }
@@ -947,4 +947,33 @@ mod tests {
"should have exactly 1 thread, not a duplicate" "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);
}
} }
+10 -10
View File
@@ -45,11 +45,11 @@ impl Agent {
&self, &self,
message: &IncomingMessage, message: &IncomingMessage,
external_thread_id: &str, external_thread_id: &str,
) -> Option<String> { ) -> Result<Option<Uuid>, String> {
// Only hydrate UUID-shaped thread IDs (web gateway uses UUIDs) // Only hydrate UUID-shaped thread IDs (web gateway uses UUIDs)
let thread_uuid = match Uuid::parse_str(external_thread_id) { let thread_uuid = match Uuid::parse_str(external_thread_id) {
Ok(id) => id, Ok(id) => id,
Err(_) => return None, Err(_) => return Ok(None),
}; };
// Check if already in memory // Check if already in memory
@@ -60,7 +60,7 @@ impl Agent {
{ {
let sess = session.lock().await; let sess = session.lock().await;
if sess.threads.contains_key(&thread_uuid) { if sess.threads.contains_key(&thread_uuid) {
return None; return Ok(Some(thread_uuid));
} }
} }
@@ -83,9 +83,9 @@ impl Agent {
e e
); );
if requires_preexisting_uuid_thread(&message.channel) { 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 { if !owned {
@@ -99,9 +99,9 @@ impl Agent {
e e
); );
if requires_preexisting_uuid_thread(&message.channel) { 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, exists,
"Rejected message for unavailable thread id" "Rejected message for unavailable thread id"
); );
return Some(FORGED_THREAD_ID_ERROR.to_string()); return Err(FORGED_THREAD_ID_ERROR.to_string());
} }
tracing::warn!( tracing::warn!(
@@ -122,7 +122,7 @@ impl Agent {
exists, exists,
"Skipped hydration for thread id not owned by sender" "Skipped hydration for thread id not owned by sender"
); );
return None; return Ok(Some(thread_uuid));
} }
let db_messages = store let db_messages = store
@@ -169,7 +169,7 @@ impl Agent {
msg_count msg_count
); );
None Ok(Some(thread_uuid))
} }
pub(super) async fn process_user_input( pub(super) async fn process_user_input(