diff --git a/src/agent/thread_ops.rs b/src/agent/thread_ops.rs index 9956a7e151..e0de224904 100644 --- a/src/agent/thread_ops.rs +++ b/src/agent/thread_ops.rs @@ -1012,6 +1012,10 @@ impl Agent { } /// Process an approval or rejection of a pending tool execution. + // Nested `if` blocks are intentional: collapsing them would produce + // `if let … && …` (let-chains), which require `#![feature(let_chains)]` + // and are not available on our MSRV. + #[allow(clippy::collapsible_if)] pub(super) async fn process_approval( &self, message: &IncomingMessage, @@ -1021,68 +1025,122 @@ impl Agent { approved: bool, always: bool, ) -> Result { - // Get pending approval for this thread - let pending = { + // Take pending approval, verify request ID, auto-approve (if always), + // and transition to Processing — all under a single lock acquisition. + // This prevents a TOCTOU race where the thread could be pruned between + // separate lock scopes, leaving a tool permanently auto-approved without + // execution ever starting (#1486). + let (pending, rejection) = { let mut sess = session.lock().await; - let thread = sess - .threads - .get_mut(&thread_id) - .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; - if thread.state != ThreadState::AwaitingApproval { - // Stale or duplicate approval (tool already executed) — silently ignore. - tracing::debug!( - %thread_id, - state = ?thread.state, - "Ignoring stale approval: thread not in AwaitingApproval state" - ); - return Ok(SubmissionResult::ok_with_message("")); - } + // First borrow: validate state and take pending approval + let taken = { + let thread = sess.threads.get_mut(&thread_id).ok_or_else(|| { + Error::from(crate::error::JobError::NotFound { id: thread_id }) + })?; + + if thread.state != ThreadState::AwaitingApproval { + // Stale or duplicate approval (tool already executed) — silently ignore. + tracing::debug!( + %thread_id, + state = ?thread.state, + "Ignoring stale approval: thread not in AwaitingApproval state" + ); + return Ok(SubmissionResult::ok_with_message("")); + } - thread.take_pending_approval() - }; + let taken = match thread.take_pending_approval() { + Some(p) => p, + None => { + tracing::debug!( + %thread_id, + "Ignoring stale approval: no pending approval found" + ); + return Ok(SubmissionResult::ok_with_message("")); + } + }; + + // Verify request ID while still holding the lock — atomic with take + if let Some(req_id) = request_id { + if req_id != taken.request_id { + // Restore atomically under same lock + thread.await_approval(taken); + return Ok(SubmissionResult::error( + "Request ID mismatch. Use the correct request ID.", + )); + } + } + + taken + // Inner borrow of `thread` ends here + }; - let pending = match pending { - Some(p) => p, - None => { - tracing::debug!( - %thread_id, - "Ignoring stale approval: no pending approval found" + // State transition while we still hold the session lock. + // If we dropped the lock first, a concurrent prune could + // remove the thread between lock scopes (#1486). + if approved { + if always { + sess.auto_approve_tool(&taken.tool_name); + tracing::info!( + "Auto-approved tool '{}' for session {}", + taken.tool_name, + sess.id + ); + } + // Re-borrow thread after session-level mutation. The thread + // cannot have disappeared — we hold the session lock throughout. + if let Some(thread) = sess.threads.get_mut(&thread_id) { + thread.state = ThreadState::Processing; + } else { + // Should be unreachable: lock held continuously since the + // thread was validated above. Roll back auto-approve. + if always { + sess.auto_approved_tools.remove(&taken.tool_name); + } + tracing::error!( + %thread_id, + "Thread disappeared while holding session lock during approval" + ); + return Ok(SubmissionResult::error( + "Internal error: thread no longer exists", + )); + } + (taken, None) + } else { + // Rejection: clear pending approval and complete the turn + // under the same lock scope that took the approval, preventing + // a TOCTOU race where a concurrent prune could remove the + // thread while state=AwaitingApproval but pending_approval=None. + let rejection = format!( + "Tool '{}' was rejected. The agent will not execute this tool.\n\n\ + You can continue the conversation or try a different approach.", + taken.tool_name ); - return Ok(SubmissionResult::ok_with_message("")); + if let Some(thread) = sess.threads.get_mut(&thread_id) { + thread.clear_pending_approval(); + thread.complete_turn(&rejection); + } else { + // Should be unreachable: lock held continuously since the + // thread was validated above. + tracing::error!( + %thread_id, + "Thread disappeared while holding session lock during rejection" + ); + return Ok(SubmissionResult::error( + "Internal error: thread no longer exists", + )); + } + (taken, Some(rejection)) } + // Lock dropped here — all approval/rejection bookkeeping is complete }; - // Verify request ID if provided - if let Some(req_id) = request_id - && req_id != pending.request_id - { - // Put it back and return error - let mut sess = session.lock().await; - if let Some(thread) = sess.threads.get_mut(&thread_id) { - thread.await_approval(pending); - } - return Ok(SubmissionResult::error( - "Request ID mismatch. Use the correct request ID.", - )); - } - if approved { - // If always, add to auto-approved set and persist to settings. + // Defense-in-depth: don't persist AlwaysAllow for tools that + // declare ApprovalRequirement::Always (the UI hides the + // "Always" button for locked tools, but a crafted client + // could send it). if always { - let mut sess = session.lock().await; - sess.auto_approve_tool(&pending.tool_name); - tracing::info!( - "Auto-approved tool '{}' for session {}", - pending.tool_name, - sess.id - ); - drop(sess); - - // Defense-in-depth: don't persist AlwaysAllow for tools that - // declare ApprovalRequirement::Always (the UI hides the - // "Always" button for locked tools, but a crafted client - // could send it). let tool_ref = self.tools().get(&pending.tool_name).await; let is_locked = tool_ref .as_ref() @@ -1122,14 +1180,6 @@ impl Agent { ), } } - } // else (not locked) - } - - // Reset thread state to processing - { - let mut sess = session.lock().await; - if let Some(thread) = sess.threads.get_mut(&thread_id) { - thread.state = ThreadState::Processing; } } @@ -1180,20 +1230,20 @@ impl Agent { ) .await; - if let Ok(ref output) = tool_result - && !output.is_empty() - { - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::ToolResult { - name: pending.tool_name.clone(), - preview: output.clone(), - }, - &message.metadata, - ) - .await; + if let Ok(ref output) = tool_result { + if !output.is_empty() { + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::ToolResult { + name: pending.tool_name.clone(), + preview: output.clone(), + }, + &message.metadata, + ) + .await; + } } // Build context including the tool result @@ -1213,15 +1263,26 @@ impl Agent { // Record sanitized result in thread { let mut sess = session.lock().await; - if let Some(thread) = sess.threads.get_mut(&thread_id) - && let Some(turn) = thread.last_turn_mut() - { - if is_tool_error { - turn.record_tool_error_for(&pending.tool_call_id, result_content.clone()); - } else { - turn.record_tool_result_for( - &pending.tool_call_id, - serde_json::json!(result_content), + match sess.threads.get_mut(&thread_id) { + Some(thread) => { + if let Some(turn) = thread.last_turn_mut() { + if is_tool_error { + turn.record_tool_error_for( + &pending.tool_call_id, + result_content.clone(), + ); + } else { + turn.record_tool_result_for( + &pending.tool_call_id, + serde_json::json!(result_content), + ); + } + } + } + None => { + tracing::error!( + %thread_id, + "Thread disappeared while recording tool result during approval" ); } } @@ -1232,16 +1293,23 @@ impl Agent { if let Some((ext_name, instructions)) = check_auth_required(&pending.tool_name, &tool_result) { - self.handle_auth_intercept( - &session, - thread_id, - message, - &tool_result, - ext_name, - instructions.clone(), - ) - .await; - return Ok(SubmissionResult::response(instructions)); + let intercepted = self + .handle_auth_intercept( + &session, + thread_id, + message, + &tool_result, + ext_name, + instructions.clone(), + ) + .await; + return if intercepted { + Ok(SubmissionResult::response(instructions)) + } else { + Ok(SubmissionResult::error( + "Internal error: thread no longer exists", + )) + }; } context_messages.push(ChatMessage::tool_result( @@ -1439,22 +1507,23 @@ impl Agent { // Process all results before any conditional return so every // tool result is recorded in the session audit trail. let mut deferred_auth: Option = None; + let mut deferred_auth_failed = false; for (tc, deferred_result) in exec_results { - if let Ok(ref output) = deferred_result - && !output.is_empty() - { - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::ToolResult { - name: tc.name.clone(), - preview: output.clone(), - }, - &message.metadata, - ) - .await; + if let Ok(ref output) = deferred_result { + if !output.is_empty() { + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::ToolResult { + name: tc.name.clone(), + preview: output.clone(), + }, + &message.metadata, + ) + .await; + } } // Sanitize first, then record the cleaned version in thread. @@ -1470,41 +1539,64 @@ impl Agent { // Record sanitized result in thread { let mut sess = session.lock().await; - if let Some(thread) = sess.threads.get_mut(&thread_id) - && let Some(turn) = thread.last_turn_mut() - { - if is_deferred_error { - turn.record_tool_error_for(&tc.id, deferred_content.clone()); - } else { - turn.record_tool_result_for( - &tc.id, - serde_json::json!(deferred_content), + match sess.threads.get_mut(&thread_id) { + Some(thread) => { + if let Some(turn) = thread.last_turn_mut() { + if is_deferred_error { + turn.record_tool_error_for(&tc.id, deferred_content.clone()); + } else { + turn.record_tool_result_for( + &tc.id, + serde_json::json!(deferred_content), + ); + } + } + } + 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 - if deferred_auth.is_none() - && let Some((ext_name, instructions)) = + if deferred_auth.is_none() { + if let Some((ext_name, instructions)) = check_auth_required(&tc.name, &deferred_result) - { - self.handle_auth_intercept( - &session, - thread_id, - message, - &deferred_result, - ext_name, - instructions.clone(), - ) - .await; - deferred_auth = Some(instructions); + { + let intercepted = self + .handle_auth_intercept( + &session, + thread_id, + message, + &deferred_result, + ext_name, + instructions.clone(), + ) + .await; + if intercepted { + deferred_auth = Some(instructions); + } else { + // Thread disappeared — record the error so the + // caller below returns an error instead of auth + // instructions for a dead thread. + deferred_auth_failed = true; + } + } } context_messages.push(ChatMessage::tool_result(&tc.id, &tc.name, deferred_content)); } // Return auth response after all results are recorded + if deferred_auth_failed { + return Ok(SubmissionResult::error( + "Internal error: thread no longer exists", + )); + } if let Some(instructions) = deferred_auth { return Ok(SubmissionResult::response(instructions)); } @@ -1532,8 +1624,19 @@ impl Agent { { let mut sess = session.lock().await; - if let Some(thread) = sess.threads.get_mut(&thread_id) { - thread.await_approval(new_pending); + match sess.threads.get_mut(&thread_id) { + Some(thread) => { + 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", + )); + } } } @@ -1678,27 +1781,33 @@ impl Agent { } } } else { - // Rejected - complete the turn with a rejection message and persist - let rejection = format!( - "Tool '{}' was rejected. The agent will not execute this tool.\n\n\ - You can continue the conversation or try a different approach.", - pending.tool_name - ); - { - let mut sess = session.lock().await; - if let Some(thread) = sess.threads.get_mut(&thread_id) { - thread.clear_pending_approval(); - thread.complete_turn(&rejection); - // User message already persisted at turn start; save rejection response - self.persist_assistant_response( - thread_id, - &message.channel, - &message.user_id, - &rejection, - ) - .await; + // Rejection bookkeeping (clear_pending_approval + complete_turn) + // already happened inside the single lock scope above. Only + // persistence and status notification remain. + // Safety: `rejection` is always `Some` when `!approved` — set + // inside the lock scope above. Use a fallback rather than + // `.expect()` to satisfy the no-panic-in-production rule. + let rejection = match rejection { + Some(r) => r, + None => { + tracing::error!( + %thread_id, + "BUG: rejection message missing for rejected approval" + ); + return Ok(SubmissionResult::error( + "Internal error: rejection state inconsistency", + )); } - } + }; + + // User message already persisted at turn start; save rejection response + self.persist_assistant_response( + thread_id, + &message.channel, + &message.user_id, + &rejection, + ) + .await; let _ = self .channels @@ -1718,6 +1827,9 @@ impl Agent { /// Enters auth mode on the thread, completes + persists the turn, /// and sends the AuthRequired status to the channel. /// Returns the instructions string for the caller to wrap in a response. + /// Returns `true` if the thread was found and auth mode was entered, + /// `false` if the thread had already disappeared (caller should not + /// return auth instructions to the user in that case). async fn handle_auth_intercept( &self, session: &Arc>, @@ -1726,36 +1838,49 @@ impl Agent { tool_result: &Result, ext_name: String, instructions: String, - ) { + ) -> bool { let auth_data = parse_auth_result(tool_result); - { + let thread_exists = { let mut sess = session.lock().await; - if let Some(thread) = sess.threads.get_mut(&thread_id) { - thread.enter_auth_mode(ext_name.clone()); - thread.complete_turn(&instructions); - // User message already persisted at turn start; save auth instructions - self.persist_assistant_response( - thread_id, + match sess.threads.get_mut(&thread_id) { + Some(thread) => { + thread.enter_auth_mode(ext_name.clone()); + thread.complete_turn(&instructions); + // User message already persisted at turn start; save auth instructions + self.persist_assistant_response( + thread_id, + &message.channel, + &message.user_id, + &instructions, + ) + .await; + true + } + None => { + tracing::error!( + %thread_id, + "Thread disappeared during auth intercept — skipping AuthRequired event" + ); + false + } + } + }; + if thread_exists { + let _ = self + .channels + .send_status( &message.channel, - &message.user_id, - &instructions, + StatusUpdate::AuthRequired { + extension_name: ext_name, + instructions: Some(instructions.clone()), + auth_url: auth_data.auth_url, + setup_url: auth_data.setup_url, + }, + &message.metadata, ) .await; - } } - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::AuthRequired { - extension_name: ext_name, - instructions: Some(instructions.clone()), - auth_url: auth_data.auth_url, - setup_url: auth_data.setup_url, - }, - &message.metadata, - ) - .await; + thread_exists } async fn send_turn_cost_status( @@ -1801,8 +1926,16 @@ impl Agent { // Clear auth mode regardless of outcome { let mut sess = session.lock().await; - if let Some(thread) = sess.threads.get_mut(&thread_id) { - thread.pending_auth = None; + match sess.threads.get_mut(&thread_id) { + Some(thread) => { + thread.pending_auth = None; + } + None => { + tracing::error!( + %thread_id, + "Thread disappeared while clearing auth mode" + ); + } } } @@ -1837,50 +1970,74 @@ impl Agent { Ok(Some(result.message)) } Ok(result) => { - { + let thread_exists = { let mut sess = session.lock().await; - if let Some(thread) = sess.threads.get_mut(&thread_id) { - thread.enter_auth_mode(pending.extension_name.clone()); - } - } - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::AuthRequired { - extension_name: pending.extension_name.clone(), - instructions: Some(result.message.clone()), - auth_url: None, - setup_url: None, - }, - &message.metadata, - ) - .await; - Ok(Some(result.message)) - } - Err(e) => { - let msg = e.to_string(); - // Token validation errors: re-enter auth mode and re-prompt - if matches!(e, crate::extensions::ExtensionError::ValidationFailed(_)) { - { - let mut sess = session.lock().await; - if let Some(thread) = sess.threads.get_mut(&thread_id) { + match sess.threads.get_mut(&thread_id) { + Some(thread) => { thread.enter_auth_mode(pending.extension_name.clone()); + true + } + None => { + tracing::error!( + %thread_id, + "Thread disappeared while re-entering auth mode" + ); + false } } + }; + if thread_exists { let _ = self .channels .send_status( &message.channel, StatusUpdate::AuthRequired { extension_name: pending.extension_name.clone(), - instructions: Some(msg.clone()), + instructions: Some(result.message.clone()), auth_url: None, setup_url: None, }, &message.metadata, ) .await; + } + Ok(Some(result.message)) + } + Err(e) => { + let msg = e.to_string(); + // Token validation errors: re-enter auth mode and re-prompt + if matches!(e, crate::extensions::ExtensionError::ValidationFailed(_)) { + let thread_exists = { + let mut sess = session.lock().await; + match sess.threads.get_mut(&thread_id) { + Some(thread) => { + thread.enter_auth_mode(pending.extension_name.clone()); + true + } + None => { + tracing::error!( + %thread_id, + "Thread disappeared while re-entering auth mode after validation failure" + ); + false + } + } + }; + if thread_exists { + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::AuthRequired { + extension_name: pending.extension_name.clone(), + instructions: Some(msg.clone()), + auth_url: None, + setup_url: None, + }, + &message.metadata, + ) + .await; + } return Ok(Some(msg)); } // Infrastructure errors @@ -2354,6 +2511,70 @@ 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, None); + 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] fn test_queue_cap_rejects_at_capacity() { use crate::agent::session::{MAX_PENDING_MESSAGES, Thread, ThreadState}; @@ -2462,6 +2683,216 @@ mod tests { // Approval persistence is tested via e2e_builtin_tool_coverage integration tests. + #[test] + fn test_approval_request_id_mismatch_restores_pending() { + // Regression test for #1486: after a request_id mismatch, the pending + // approval must still be intact (take + verify + restore is atomic). + use crate::agent::session::{PendingApproval, Thread, ThreadState}; + use uuid::Uuid; + + let session_id = Uuid::new_v4(); + let thread_id = Uuid::new_v4(); + let mut thread = Thread::with_id(thread_id, session_id, None); + + let correct_request_id = Uuid::new_v4(); + let pending = PendingApproval { + request_id: correct_request_id, + tool_name: "shell".to_string(), + parameters: serde_json::json!({}), + display_parameters: serde_json::json!({}), + description: "test".to_string(), + tool_call_id: "call_0".to_string(), + context_messages: vec![], + deferred_tool_calls: vec![], + user_timezone: None, + allow_always: true, + }; + thread.await_approval(pending); + assert_eq!(thread.state, ThreadState::AwaitingApproval); + + // Simulate: take, verify mismatch, restore -- all must be atomic + let taken = thread.take_pending_approval().unwrap(); + assert_eq!(taken.request_id, correct_request_id); + // On mismatch, restore + thread.await_approval(taken); + // Must still be in AwaitingApproval with pending intact + assert_eq!(thread.state, ThreadState::AwaitingApproval); + assert!(thread.pending_approval.is_some()); + assert_eq!( + thread.pending_approval.as_ref().unwrap().request_id, + correct_request_id + ); + } + + /// Mirrors the single-lock pattern in `process_approval()`: take the + /// pending approval from the thread, optionally auto-approve, and + /// transition to `Processing` — all within a single borrow scope. + /// + /// Because the thread lookup, auto-approve, and state transition all + /// happen under one `&mut Session` borrow, the thread cannot be pruned + /// between steps. Returns the taken `PendingApproval` on success, or + /// `None` if the thread does not exist or has no pending approval. + fn take_and_approve( + session: &mut crate::agent::session::Session, + thread_id: uuid::Uuid, + always: bool, + ) -> Option { + // Take the pending approval — early return if thread or approval missing + let taken = { + let thread = session.threads.get_mut(&thread_id)?; + thread.take_pending_approval()? + // thread borrow ends here + }; + + if always { + session.auto_approve_tool(&taken.tool_name); + } + + // Re-borrow after session-level mutation + let thread = session.threads.get_mut(&thread_id).expect("thread exists"); + thread.state = ThreadState::Processing; + + Some(taken) + } + + /// Helper to build a minimal PendingApproval for tests. + fn make_pending(tool_name: &str) -> crate::agent::session::PendingApproval { + crate::agent::session::PendingApproval { + request_id: uuid::Uuid::new_v4(), + tool_name: tool_name.to_string(), + tool_call_id: "call-1".to_string(), + description: "test".to_string(), + parameters: serde_json::json!({}), + display_parameters: serde_json::json!({}), + context_messages: vec![], + deferred_tool_calls: vec![], + user_timezone: None, + allow_always: true, + } + } + + #[test] + fn test_auto_approve_with_thread_disappearance_never_commits() { + // Regression test: when always=true and the thread does not exist, + // the single-lock pattern in process_approval() prevents the + // auto-approve from ever being committed — the thread lookup fails + // before auto_approve_tool is called, so no rollback is needed. + use crate::agent::session::Session; + + let mut session = Session::new("test-user"); + let thread_id = uuid::Uuid::new_v4(); + + // Thread does not exist — take_and_approve must return None + // without adding to auto_approved_tools + assert!(!session.threads.contains_key(&thread_id)); + let result = take_and_approve(&mut session, thread_id, true); + assert!(result.is_none(), "must fail when thread is missing"); + assert!( + session.auto_approved_tools.is_empty(), + "auto-approve must never be committed when the thread is missing" + ); + } + + #[test] + fn test_auto_approve_with_present_thread_succeeds() { + // Happy-path counterpart: when always=true and the thread exists, the + // tool should remain auto-approved and the thread should transition + // to Processing. + use crate::agent::session::Session; + + let mut session = Session::new("test-user"); + let thread = session.create_thread(Some("test")); + let thread_id = thread.id; + + // Set up a pending approval on the thread + thread.state = ThreadState::AwaitingApproval; + thread.await_approval(make_pending("safe_tool")); + + let result = take_and_approve(&mut session, thread_id, true); + assert!( + result.is_some(), + "transition must succeed when thread is present" + ); + assert!( + session.is_tool_auto_approved("safe_tool"), + "tool must remain auto-approved after successful transition" + ); + assert_eq!( + session.threads.get(&thread_id).map(|t| &t.state), + Some(&ThreadState::Processing), + "thread must be in Processing state after approval" + ); + } + + #[test] + fn test_non_always_approve_does_not_add_to_auto_approved() { + // When always=false, the tool should not be added to auto-approved + // regardless of whether the thread exists. + use crate::agent::session::Session; + + let mut session = Session::new("test-user"); + let thread = session.create_thread(Some("test")); + let thread_id = thread.id; + + // Set up a pending approval on the thread + thread.state = ThreadState::AwaitingApproval; + thread.await_approval(make_pending("one_time_tool")); + + let result = take_and_approve(&mut session, thread_id, false); + assert!( + result.is_some(), + "transition must succeed when thread is present" + ); + assert!( + !session.is_tool_auto_approved("one_time_tool"), + "tool must not be auto-approved when always=false" + ); + } + + #[test] + fn test_auto_approve_thread_pruned_between_old_lock_scopes_impossible() { + // This test verifies the fix for the TOCTOU described in PR #1591: + // previously, `process_approval()` used two lock scopes — one to + // take the pending approval, and another to auto-approve + transition. + // If the thread was pruned between the two locks, the tool would be + // permanently auto-approved despite execution never starting. + // + // With the single-lock fix, auto-approve only happens when the thread + // is confirmed to exist (same borrow scope as the pending-approval + // take). We verify this by showing that removing the thread before + // the combined operation prevents auto-approve from ever committing. + use crate::agent::session::Session; + + let mut session = Session::new("test-user"); + let thread = session.create_thread(Some("test")); + let thread_id = thread.id; + + // Set up a pending approval + thread.state = ThreadState::AwaitingApproval; + thread.await_approval(make_pending("risky_tool")); + + // Simulate the old bug: remove the thread before the approval logic + // runs. Under the old two-lock pattern, the first lock would have + // already taken the pending approval, and the second lock would find + // the thread missing — leaving auto-approve committed. + session.threads.remove(&thread_id); + + // Under the single-lock pattern, the entire operation fails atomically + let result = take_and_approve(&mut session, thread_id, true); + assert!( + result.is_none(), + "must fail when thread is pruned before approval" + ); + assert!( + !session.is_tool_auto_approved("risky_tool"), + "auto-approve must never be committed when thread is pruned" + ); + assert!( + session.auto_approved_tools.is_empty(), + "no tools should be in the auto-approved set" + ); + } + // Helper function to extract the approval message without needing a full Agent instance fn extract_approval_message( session: &crate::agent::session::Session,