Compare commits

..
Author SHA1 Message Date
Claude 61a626a10d style: fix rustfmt formatting in session_manager
https://claude.ai/code/session_01Sk4XcMNNjTAVYLn4VTzeY4
2026-03-23 07:12:02 +00:00
ZakiandClaude Opus 4.6 4a26987d8e perf(agent): eliminate redundant UUID parsing in approval path (#1488)
Pass pre-parsed UUID through the approval resolution flow instead of
parsing the same string twice. The new resolve_thread_with_parsed_uuid
method accepts an optional pre-parsed UUID, and resolve_thread delegates
to it with None for backward compatibility.

Closes #1488

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-22 23:26:31 -07:00
3 changed files with 102 additions and 157 deletions
+2 -1
View File
@@ -1052,10 +1052,11 @@ 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(),
approval_thread_uuid,
) )
.await .await
} }
+69 -3
View File
@@ -107,6 +107,20 @@ impl SessionManager {
user_id: &str, user_id: &str,
channel: &str, channel: &str,
external_thread_id: Option<&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).
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) { ) -> (Arc<Mutex<Session>>, Uuid) {
let session = self.get_or_create_session(user_id).await; let session = self.get_or_create_session(user_id).await;
@@ -133,9 +147,11 @@ impl SessionManager {
// (e.g. created by chat_new_thread_handler or hydrated from DB). // (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 — // We only adopt it if no thread_map entry maps to this UUID —
// otherwise it belongs to a different channel scope. // otherwise it belongs to a different channel scope.
if let Some(ext_tid) = external_thread_id // Use pre-parsed UUID if available, otherwise parse from string.
&& let Ok(ext_uuid) = Uuid::parse_str(ext_tid) let ext_uuid = parsed_uuid
{ .or_else(|| external_thread_id.and_then(|ext_tid| Uuid::parse_str(ext_tid).ok()));
if let Some(ext_uuid) = ext_uuid {
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);
@@ -947,4 +963,54 @@ 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_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);
}
} }
+10 -132
View File
@@ -992,17 +992,9 @@ impl Agent {
{ {
// Put it back and return error // Put it back and return error
let mut sess = session.lock().await; let mut sess = session.lock().await;
match sess.threads.get_mut(&thread_id) { if let Some(thread) = sess.threads.get_mut(&thread_id) {
Some(thread) => {
thread.await_approval(pending); thread.await_approval(pending);
} }
None => {
tracing::warn!(
%thread_id,
"Thread disappeared while restoring pending approval after request ID mismatch"
);
}
}
return Ok(SubmissionResult::error( return Ok(SubmissionResult::error(
"Request ID mismatch. Use the correct request ID.", "Request ID mismatch. Use the correct request ID.",
)); ));
@@ -1023,20 +1015,9 @@ impl Agent {
// Reset thread state to processing // Reset thread state to processing
{ {
let mut sess = session.lock().await; let mut sess = session.lock().await;
match sess.threads.get_mut(&thread_id) { if let Some(thread) = sess.threads.get_mut(&thread_id) {
Some(thread) => {
thread.state = ThreadState::Processing; thread.state = ThreadState::Processing;
} }
None => {
tracing::error!(
%thread_id,
"Thread disappeared while setting state to Processing during approval"
);
return Ok(SubmissionResult::error(
"Internal error: thread no longer exists",
));
}
}
} }
// Execute the approved tool and continue the loop // Execute the approved tool and continue the loop
@@ -1119,9 +1100,9 @@ impl Agent {
// Record sanitized result in thread // Record sanitized result in thread
{ {
let mut sess = session.lock().await; let mut sess = session.lock().await;
match sess.threads.get_mut(&thread_id) { if let Some(thread) = sess.threads.get_mut(&thread_id)
Some(thread) => { && let Some(turn) = thread.last_turn_mut()
if let Some(turn) = thread.last_turn_mut() { {
if is_tool_error { if is_tool_error {
turn.record_tool_error(result_content.clone()); turn.record_tool_error(result_content.clone());
} else { } else {
@@ -1129,14 +1110,6 @@ impl Agent {
} }
} }
} }
None => {
tracing::error!(
%thread_id,
"Thread disappeared while recording tool result during approval"
);
}
}
}
// If tool_auth returned awaiting_token, enter auth mode and // If tool_auth returned awaiting_token, enter auth mode and
// return instructions directly (skip agentic loop continuation). // return instructions directly (skip agentic loop continuation).
@@ -1381,9 +1354,9 @@ impl Agent {
// Record sanitized result in thread // Record sanitized result in thread
{ {
let mut sess = session.lock().await; let mut sess = session.lock().await;
match sess.threads.get_mut(&thread_id) { if let Some(thread) = sess.threads.get_mut(&thread_id)
Some(thread) => { && let Some(turn) = thread.last_turn_mut()
if let Some(turn) = thread.last_turn_mut() { {
if is_deferred_error { if is_deferred_error {
turn.record_tool_error(deferred_content.clone()); turn.record_tool_error(deferred_content.clone());
} else { } else {
@@ -1391,15 +1364,6 @@ impl Agent {
} }
} }
} }
None => {
tracing::error!(
%thread_id,
tool_name = %tc.name,
"Thread disappeared while recording deferred tool result during approval"
);
}
}
}
// Auth detection — defer return until all results are recorded // Auth detection — defer return until all results are recorded
if deferred_auth.is_none() if deferred_auth.is_none()
@@ -1449,20 +1413,9 @@ impl Agent {
{ {
let mut sess = session.lock().await; let mut sess = session.lock().await;
match sess.threads.get_mut(&thread_id) { if let Some(thread) = sess.threads.get_mut(&thread_id) {
Some(thread) => {
thread.await_approval(new_pending); thread.await_approval(new_pending);
} }
None => {
tracing::error!(
%thread_id,
"Thread disappeared while setting up deferred tool approval"
);
return Ok(SubmissionResult::error(
"Internal error: thread no longer exists",
));
}
}
} }
let _ = self let _ = self
@@ -1593,8 +1546,7 @@ impl Agent {
); );
{ {
let mut sess = session.lock().await; let mut sess = session.lock().await;
match sess.threads.get_mut(&thread_id) { if let Some(thread) = sess.threads.get_mut(&thread_id) {
Some(thread) => {
thread.clear_pending_approval(); thread.clear_pending_approval();
thread.complete_turn(&rejection); thread.complete_turn(&rejection);
// User message already persisted at turn start; save rejection response // User message already persisted at turn start; save rejection response
@@ -1606,16 +1558,6 @@ impl Agent {
) )
.await; .await;
} }
None => {
tracing::error!(
%thread_id,
"Thread disappeared during approval rejection"
);
return Ok(SubmissionResult::error(
"Internal error: thread no longer exists",
));
}
}
} }
let _ = self let _ = self
@@ -2156,70 +2098,6 @@ mod tests {
} }
} }
#[tokio::test]
async fn test_approval_on_missing_thread_should_error() {
// Regression for #1487: when a thread disappears from the session
// during approval processing, the code must return a visible error
// rather than silently succeeding.
//
// We can't call process_approval() directly (requires full Agent),
// so we simulate the exact code pattern used in the rejection and
// state-setting paths: lock session, match on get_mut, verify the
// None arm produces an error.
use crate::agent::session::{Session, Thread, ThreadState};
use std::sync::Arc;
use tokio::sync::Mutex;
use uuid::Uuid;
let thread_id = Uuid::new_v4();
let session_id = Uuid::new_v4();
let session = Arc::new(Mutex::new(Session::new("test-user")));
// Scenario 1: Thread never existed
{
let sess = session.lock().await;
let result = match sess.threads.get(&thread_id) {
Some(_) => Ok("processed"),
None => Err("Internal error: thread no longer exists"),
};
assert!(result.is_err());
assert_eq!(
result.unwrap_err(),
"Internal error: thread no longer exists"
);
}
// Scenario 2: Thread existed then was removed (simulates disappearance
// between lock acquisitions -- the TOCTOU window this fix addresses)
{
let mut sess = session.lock().await;
let mut thread = Thread::with_id(thread_id, session_id);
thread.start_turn("pending approval");
thread.state = ThreadState::AwaitingApproval;
sess.threads.insert(thread_id, thread);
}
{
let mut sess = session.lock().await;
// Simulate thread disappearing (e.g., pruned by another task)
sess.threads.remove(&thread_id);
// The rejection path must detect this and return an error
let result = match sess.threads.get_mut(&thread_id) {
Some(thread) => {
thread.clear_pending_approval();
thread.complete_turn("rejected");
Ok("rejection persisted")
}
None => Err("Internal error: thread no longer exists"),
};
assert!(result.is_err());
assert_eq!(
result.unwrap_err(),
"Internal error: thread no longer exists"
);
}
}
#[test] #[test]
fn test_queue_cap_rejects_at_capacity() { fn test_queue_cap_rejects_at_capacity() {
use crate::agent::session::{MAX_PENDING_MESSAGES, Thread, ThreadState}; use crate::agent::session::{MAX_PENDING_MESSAGES, Thread, ThreadState};