Skip to main content

a3s_code_core/
hitl.rs

1//! Human-in-the-Loop (HITL) confirmation mechanism
2//!
3//! Provides the runtime confirmation flow for tool execution. Works with
4//! `PermissionPolicy` (permissions.rs) which decides Allow/Deny/Ask.
5//! When the permission decision is `Ask`, this module handles:
6//! - Interactive confirmation request/response flow
7//! - Timeout handling with configurable actions
8//! - YOLO mode for lane-based auto-approval (skips confirmation for entire lanes)
9
10use crate::agent::AgentEvent;
11use serde::{Deserialize, Serialize};
12use std::collections::{HashMap, HashSet};
13use std::sync::Arc;
14use std::time::{Duration, Instant};
15use tokio::sync::{broadcast, oneshot, RwLock};
16
17// Re-export SessionLane for backward compatibility (canonical home: queue.rs)
18pub use crate::queue::SessionLane;
19
20/// Action to take when confirmation times out
21#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
22pub enum TimeoutAction {
23    /// Reject the tool execution on timeout
24    #[default]
25    Reject,
26    /// Auto-approve the tool execution on timeout
27    AutoApprove,
28}
29
30/// Confirmation policy configuration
31///
32/// Controls the runtime behavior of HITL confirmation flow.
33/// The *decision* of whether to ask is made by `PermissionPolicy` (permissions.rs).
34/// This policy controls *how* the confirmation works: timeouts, YOLO lanes.
35#[derive(Debug, Clone, Serialize, Deserialize)]
36pub struct ConfirmationPolicy {
37    /// Whether HITL is enabled (default: false, all tools auto-approved)
38    pub enabled: bool,
39
40    /// Default timeout in milliseconds (default: 30000 = 30s)
41    pub default_timeout_ms: u64,
42
43    /// Action to take on timeout (default: Reject)
44    pub timeout_action: TimeoutAction,
45
46    /// YOLO mode: lanes that auto-approve without confirmation.
47    /// When a lane is in this set, tools in that lane skip confirmation
48    /// even if `PermissionPolicy` returns `Ask`.
49    pub yolo_lanes: HashSet<SessionLane>,
50}
51
52impl Default for ConfirmationPolicy {
53    fn default() -> Self {
54        Self {
55            enabled: false,             // HITL disabled by default
56            default_timeout_ms: 30_000, // 30 seconds
57            timeout_action: TimeoutAction::Reject,
58            yolo_lanes: HashSet::new(), // No YOLO lanes by default
59        }
60    }
61}
62
63impl ConfirmationPolicy {
64    /// Create a new policy with HITL enabled
65    pub fn enabled() -> Self {
66        Self {
67            enabled: true,
68            ..Default::default()
69        }
70    }
71
72    /// Enable YOLO mode for specific lanes
73    pub fn with_yolo_lanes(mut self, lanes: impl IntoIterator<Item = SessionLane>) -> Self {
74        self.yolo_lanes = lanes.into_iter().collect();
75        self
76    }
77
78    /// Set timeout
79    pub fn with_timeout(mut self, timeout_ms: u64, action: TimeoutAction) -> Self {
80        self.default_timeout_ms = timeout_ms;
81        self.timeout_action = action;
82        self
83    }
84
85    /// Check if a tool should skip confirmation (YOLO lane check)
86    ///
87    /// Returns true if the tool's lane is in YOLO mode, meaning it should
88    /// be auto-approved even when `PermissionPolicy` returns `Ask`.
89    pub fn is_yolo(&self, tool_name: &str) -> bool {
90        if !self.enabled {
91            return true; // HITL disabled = everything auto-approved
92        }
93        let lane = SessionLane::from_tool_name(tool_name);
94        self.yolo_lanes.contains(&lane)
95    }
96
97    /// Check if a tool requires confirmation
98    ///
99    /// This is the inverse of `is_yolo()` — returns true when HITL is enabled
100    /// and the tool's lane is NOT in YOLO mode.
101    pub fn requires_confirmation(&self, tool_name: &str) -> bool {
102        !self.is_yolo(tool_name)
103    }
104}
105
106/// Confirmation response from user
107#[derive(Debug, Clone, Serialize, Deserialize)]
108pub struct ConfirmationResponse {
109    /// Whether the tool execution was approved
110    pub approved: bool,
111    /// Optional reason for rejection
112    pub reason: Option<String>,
113}
114
115/// Snapshot of a pending confirmation request.
116#[derive(Debug, Clone, Serialize, Deserialize)]
117pub struct PendingConfirmationInfo {
118    pub tool_id: String,
119    pub tool_name: String,
120    pub args: serde_json::Value,
121    pub remaining_ms: u64,
122}
123
124/// Trait for confirmation providers (HITL runtime behavior)
125///
126/// This trait abstracts the confirmation flow, allowing different implementations
127/// (e.g., interactive, auto-approve, test mocks) while keeping the agent logic clean.
128#[async_trait::async_trait]
129pub trait ConfirmationProvider: Send + Sync {
130    /// Check if a tool requires confirmation
131    async fn requires_confirmation(&self, tool_name: &str) -> bool;
132
133    /// Request confirmation for a tool execution
134    ///
135    /// Returns a receiver that will receive the confirmation response.
136    async fn request_confirmation(
137        &self,
138        tool_id: &str,
139        tool_name: &str,
140        args: &serde_json::Value,
141    ) -> oneshot::Receiver<ConfirmationResponse>;
142
143    /// Handle a confirmation response from the user
144    ///
145    /// Returns Ok(true) if the confirmation was found and processed,
146    /// Ok(false) if no pending confirmation was found.
147    async fn confirm(
148        &self,
149        tool_id: &str,
150        approved: bool,
151        reason: Option<String>,
152    ) -> Result<bool, String>;
153
154    /// Get the current policy
155    async fn policy(&self) -> ConfirmationPolicy;
156
157    /// Update the confirmation policy
158    async fn set_policy(&self, policy: ConfirmationPolicy);
159
160    /// Check for and handle timed out confirmations
161    async fn check_timeouts(&self) -> usize;
162
163    /// Cancel all pending confirmations
164    async fn cancel_all(&self) -> usize;
165
166    /// Snapshot pending confirmations for status inspection.
167    async fn pending_confirmations(&self) -> Vec<PendingConfirmationInfo> {
168        Vec::new()
169    }
170}
171
172/// A pending confirmation request
173pub struct PendingConfirmation {
174    /// Tool call ID
175    pub tool_id: String,
176    /// Tool name
177    pub tool_name: String,
178    /// Tool arguments
179    pub args: serde_json::Value,
180    /// When the confirmation was requested
181    pub created_at: Instant,
182    /// Timeout in milliseconds
183    pub timeout_ms: u64,
184    /// Channel to send the response
185    response_tx: oneshot::Sender<ConfirmationResponse>,
186}
187
188impl PendingConfirmation {
189    /// Check if this confirmation has timed out
190    pub fn is_timed_out(&self) -> bool {
191        self.created_at.elapsed() > Duration::from_millis(self.timeout_ms)
192    }
193
194    /// Get remaining time until timeout in milliseconds
195    pub fn remaining_ms(&self) -> u64 {
196        let elapsed = self.created_at.elapsed().as_millis() as u64;
197        self.timeout_ms.saturating_sub(elapsed)
198    }
199}
200
201/// Manages confirmation requests for a session
202pub struct ConfirmationManager {
203    /// Confirmation policy
204    policy: RwLock<ConfirmationPolicy>,
205    /// Pending confirmations by tool_id
206    pending: Arc<RwLock<HashMap<String, PendingConfirmation>>>,
207    /// Event broadcaster
208    event_tx: broadcast::Sender<AgentEvent>,
209}
210
211impl ConfirmationManager {
212    /// Create a new confirmation manager
213    pub fn new(policy: ConfirmationPolicy, event_tx: broadcast::Sender<AgentEvent>) -> Self {
214        Self {
215            policy: RwLock::new(policy),
216            pending: Arc::new(RwLock::new(HashMap::new())),
217            event_tx,
218        }
219    }
220
221    /// Get the current policy
222    pub async fn policy(&self) -> ConfirmationPolicy {
223        self.policy.read().await.clone()
224    }
225
226    /// Update the confirmation policy
227    pub async fn set_policy(&self, policy: ConfirmationPolicy) {
228        *self.policy.write().await = policy;
229    }
230
231    /// Check if a tool requires confirmation
232    pub async fn requires_confirmation(&self, tool_name: &str) -> bool {
233        self.policy.read().await.requires_confirmation(tool_name)
234    }
235
236    /// Request confirmation for a tool execution
237    ///
238    /// Returns a receiver that will receive the confirmation response.
239    /// Emits a ConfirmationRequired event.
240    pub async fn request_confirmation(
241        &self,
242        tool_id: &str,
243        tool_name: &str,
244        args: &serde_json::Value,
245    ) -> oneshot::Receiver<ConfirmationResponse> {
246        let (tx, rx) = oneshot::channel();
247
248        let policy = self.policy.read().await;
249        let timeout_ms = policy.default_timeout_ms;
250        drop(policy);
251
252        let pending = PendingConfirmation {
253            tool_id: tool_id.to_string(),
254            tool_name: tool_name.to_string(),
255            args: args.clone(),
256            created_at: Instant::now(),
257            timeout_ms,
258            response_tx: tx,
259        };
260
261        // Store the pending confirmation
262        {
263            let mut pending_map = self.pending.write().await;
264            pending_map.insert(tool_id.to_string(), pending);
265        }
266
267        // Emit confirmation required event
268        let _ = self.event_tx.send(AgentEvent::ConfirmationRequired {
269            tool_id: tool_id.to_string(),
270            tool_name: tool_name.to_string(),
271            args: args.clone(),
272            timeout_ms,
273        });
274
275        rx
276    }
277
278    /// Handle a confirmation response from the user
279    ///
280    /// Returns Ok(true) if the confirmation was found and processed,
281    /// Ok(false) if no pending confirmation was found.
282    pub async fn confirm(
283        &self,
284        tool_id: &str,
285        approved: bool,
286        reason: Option<String>,
287    ) -> Result<bool, String> {
288        let pending = {
289            let mut pending_map = self.pending.write().await;
290            pending_map.remove(tool_id)
291        };
292
293        if let Some(confirmation) = pending {
294            // Emit confirmation received event
295            let _ = self.event_tx.send(AgentEvent::ConfirmationReceived {
296                tool_id: tool_id.to_string(),
297                approved,
298                reason: reason.clone(),
299            });
300
301            // Send the response
302            let response = ConfirmationResponse { approved, reason };
303            let _ = confirmation.response_tx.send(response);
304
305            Ok(true)
306        } else {
307            Ok(false)
308        }
309    }
310
311    /// Check for and handle timed out confirmations
312    ///
313    /// Returns the number of confirmations that timed out.
314    pub async fn check_timeouts(&self) -> usize {
315        let policy = self.policy.read().await;
316        let timeout_action = policy.timeout_action;
317        drop(policy);
318
319        let mut timed_out = Vec::new();
320
321        // Find timed out confirmations
322        {
323            let pending_map = self.pending.read().await;
324            for (tool_id, pending) in pending_map.iter() {
325                if pending.is_timed_out() {
326                    timed_out.push(tool_id.clone());
327                }
328            }
329        }
330
331        // Handle timed out confirmations
332        for tool_id in &timed_out {
333            let pending = {
334                let mut pending_map = self.pending.write().await;
335                pending_map.remove(tool_id)
336            };
337
338            if let Some(confirmation) = pending {
339                let (approved, action_taken) = match timeout_action {
340                    TimeoutAction::Reject => (false, "rejected"),
341                    TimeoutAction::AutoApprove => (true, "auto_approved"),
342                };
343
344                // Emit timeout event
345                let _ = self.event_tx.send(AgentEvent::ConfirmationTimeout {
346                    tool_id: tool_id.clone(),
347                    action_taken: action_taken.to_string(),
348                });
349
350                // Send the response
351                let response = ConfirmationResponse {
352                    approved,
353                    reason: Some(format!("Confirmation timed out, action: {}", action_taken)),
354                };
355                let _ = confirmation.response_tx.send(response);
356            }
357        }
358
359        timed_out.len()
360    }
361
362    /// Get the number of pending confirmations
363    pub async fn pending_count(&self) -> usize {
364        self.pending.read().await.len()
365    }
366
367    /// Get pending confirmation details (for debugging/status)
368    pub async fn pending_confirmations(&self) -> Vec<(String, String, u64)> {
369        let pending_map = self.pending.read().await;
370        pending_map
371            .values()
372            .map(|p| (p.tool_id.clone(), p.tool_name.clone(), p.remaining_ms()))
373            .collect()
374    }
375
376    /// Get detailed pending confirmation snapshots.
377    pub async fn pending_confirmation_details(&self) -> Vec<PendingConfirmationInfo> {
378        let pending_map = self.pending.read().await;
379        pending_map
380            .values()
381            .map(|p| PendingConfirmationInfo {
382                tool_id: p.tool_id.clone(),
383                tool_name: p.tool_name.clone(),
384                args: p.args.clone(),
385                remaining_ms: p.remaining_ms(),
386            })
387            .collect()
388    }
389
390    /// Cancel a pending confirmation
391    pub async fn cancel(&self, tool_id: &str) -> bool {
392        let pending = {
393            let mut pending_map = self.pending.write().await;
394            pending_map.remove(tool_id)
395        };
396
397        if let Some(confirmation) = pending {
398            let response = ConfirmationResponse {
399                approved: false,
400                reason: Some("Confirmation cancelled".to_string()),
401            };
402            let _ = confirmation.response_tx.send(response);
403            true
404        } else {
405            false
406        }
407    }
408
409    /// Cancel all pending confirmations
410    pub async fn cancel_all(&self) -> usize {
411        let pending_list: Vec<_> = {
412            let mut pending_map = self.pending.write().await;
413            pending_map.drain().collect()
414        };
415
416        let count = pending_list.len();
417
418        for (_, confirmation) in pending_list {
419            let response = ConfirmationResponse {
420                approved: false,
421                reason: Some("Confirmation cancelled".to_string()),
422            };
423            let _ = confirmation.response_tx.send(response);
424        }
425
426        count
427    }
428}
429
430// Implement ConfirmationProvider trait for ConfirmationManager
431#[async_trait::async_trait]
432impl ConfirmationProvider for ConfirmationManager {
433    async fn requires_confirmation(&self, tool_name: &str) -> bool {
434        self.requires_confirmation(tool_name).await
435    }
436
437    async fn request_confirmation(
438        &self,
439        tool_id: &str,
440        tool_name: &str,
441        args: &serde_json::Value,
442    ) -> oneshot::Receiver<ConfirmationResponse> {
443        self.request_confirmation(tool_id, tool_name, args).await
444    }
445
446    async fn confirm(
447        &self,
448        tool_id: &str,
449        approved: bool,
450        reason: Option<String>,
451    ) -> Result<bool, String> {
452        self.confirm(tool_id, approved, reason).await
453    }
454
455    async fn policy(&self) -> ConfirmationPolicy {
456        self.policy().await
457    }
458
459    async fn set_policy(&self, policy: ConfirmationPolicy) {
460        self.set_policy(policy).await
461    }
462
463    async fn check_timeouts(&self) -> usize {
464        self.check_timeouts().await
465    }
466
467    async fn cancel_all(&self) -> usize {
468        self.cancel_all().await
469    }
470
471    async fn pending_confirmations(&self) -> Vec<PendingConfirmationInfo> {
472        self.pending_confirmation_details().await
473    }
474}
475
476#[cfg(test)]
477mod tests {
478    use super::*;
479
480    // ========================================================================
481    // SessionLane Tests
482    // ========================================================================
483
484    #[test]
485    fn test_session_lane() {
486        assert_eq!(SessionLane::from_tool_name("read"), SessionLane::Query);
487        assert_eq!(SessionLane::from_tool_name("grep"), SessionLane::Query);
488        assert_eq!(SessionLane::from_tool_name("bash"), SessionLane::Execute);
489        assert_eq!(SessionLane::from_tool_name("write"), SessionLane::Execute);
490    }
491
492    #[test]
493    fn test_session_lane_priority() {
494        assert_eq!(SessionLane::Control.priority(), 0);
495        assert_eq!(SessionLane::Query.priority(), 1);
496        assert_eq!(SessionLane::Execute.priority(), 2);
497        assert_eq!(SessionLane::Generate.priority(), 3);
498
499        // Control has highest priority (lowest number)
500        assert!(SessionLane::Control.priority() < SessionLane::Query.priority());
501        assert!(SessionLane::Query.priority() < SessionLane::Execute.priority());
502        assert!(SessionLane::Execute.priority() < SessionLane::Generate.priority());
503    }
504
505    #[test]
506    fn test_session_lane_all_query() {
507        let query_tools = ["read", "glob", "ls", "grep", "list_files", "search"];
508        for tool in query_tools {
509            assert_eq!(
510                SessionLane::from_tool_name(tool),
511                SessionLane::Query,
512                "Tool '{}' should be in Query lane",
513                tool
514            );
515        }
516    }
517
518    #[test]
519    fn test_session_lane_all_execute() {
520        let execute_tools = ["bash", "write", "edit", "delete", "move", "copy", "execute"];
521        for tool in execute_tools {
522            assert_eq!(
523                SessionLane::from_tool_name(tool),
524                SessionLane::Execute,
525                "Tool '{}' should be in Execute lane",
526                tool
527            );
528        }
529    }
530
531    // ========================================================================
532    // TimeoutAction Tests
533    // ========================================================================
534
535    // ========================================================================
536    // ConfirmationPolicy Tests
537    // ========================================================================
538
539    #[test]
540    fn test_confirmation_policy_default() {
541        let policy = ConfirmationPolicy::default();
542        assert!(!policy.enabled);
543        // HITL disabled = everything is YOLO (no confirmation needed)
544        assert!(!policy.requires_confirmation("bash"));
545        assert!(!policy.requires_confirmation("write"));
546        assert!(!policy.requires_confirmation("read"));
547    }
548
549    #[test]
550    fn test_confirmation_policy_enabled() {
551        let policy = ConfirmationPolicy::enabled();
552        assert!(policy.enabled);
553        // All tools require confirmation when enabled with no YOLO lanes
554        assert!(policy.requires_confirmation("bash"));
555        assert!(policy.requires_confirmation("write"));
556        assert!(policy.requires_confirmation("read"));
557        assert!(policy.requires_confirmation("grep"));
558    }
559
560    #[test]
561    fn test_confirmation_policy_yolo_mode() {
562        let policy = ConfirmationPolicy::enabled().with_yolo_lanes([SessionLane::Execute]);
563
564        assert!(!policy.requires_confirmation("bash")); // Execute lane in YOLO mode
565        assert!(!policy.requires_confirmation("write")); // Execute lane in YOLO mode
566        assert!(policy.requires_confirmation("read")); // Query lane NOT in YOLO
567    }
568
569    #[test]
570    fn test_confirmation_policy_yolo_multiple_lanes() {
571        let policy = ConfirmationPolicy::enabled()
572            .with_yolo_lanes([SessionLane::Query, SessionLane::Execute]);
573
574        // All tools in YOLO lanes should be auto-approved
575        assert!(!policy.requires_confirmation("bash")); // Execute
576        assert!(!policy.requires_confirmation("read")); // Query
577        assert!(!policy.requires_confirmation("grep")); // Query
578    }
579
580    #[test]
581    fn test_confirmation_policy_is_yolo() {
582        let policy = ConfirmationPolicy::enabled().with_yolo_lanes([SessionLane::Execute]);
583
584        assert!(policy.is_yolo("bash")); // Execute lane
585        assert!(policy.is_yolo("write")); // Execute lane
586        assert!(!policy.is_yolo("read")); // Query lane, not YOLO
587    }
588
589    #[test]
590    fn test_confirmation_policy_disabled_is_always_yolo() {
591        let policy = ConfirmationPolicy::default(); // disabled
592        assert!(policy.is_yolo("bash"));
593        assert!(policy.is_yolo("read"));
594        assert!(policy.is_yolo("unknown_tool"));
595    }
596
597    #[test]
598    fn test_confirmation_policy_with_timeout() {
599        let policy = ConfirmationPolicy::enabled().with_timeout(5000, TimeoutAction::AutoApprove);
600
601        assert_eq!(policy.default_timeout_ms, 5000);
602        assert_eq!(policy.timeout_action, TimeoutAction::AutoApprove);
603    }
604
605    // ========================================================================
606    // ConfirmationManager Basic Tests
607    // ========================================================================
608
609    #[tokio::test]
610    async fn test_confirmation_manager_no_hitl() {
611        let (event_tx, _) = broadcast::channel(100);
612        let manager = ConfirmationManager::new(ConfirmationPolicy::default(), event_tx);
613
614        assert!(!manager.requires_confirmation("bash").await);
615    }
616
617    #[tokio::test]
618    async fn test_confirmation_manager_with_hitl() {
619        let (event_tx, _) = broadcast::channel(100);
620        let manager = ConfirmationManager::new(ConfirmationPolicy::enabled(), event_tx);
621
622        // All tools require confirmation when HITL enabled with no YOLO lanes
623        assert!(manager.requires_confirmation("bash").await);
624        assert!(manager.requires_confirmation("read").await);
625    }
626
627    #[tokio::test]
628    async fn test_confirmation_manager_with_yolo() {
629        let (event_tx, _) = broadcast::channel(100);
630        let policy = ConfirmationPolicy::enabled().with_yolo_lanes([SessionLane::Query]);
631        let manager = ConfirmationManager::new(policy, event_tx);
632
633        assert!(manager.requires_confirmation("bash").await); // Execute lane, not YOLO
634        assert!(!manager.requires_confirmation("read").await); // Query lane, YOLO
635    }
636
637    #[tokio::test]
638    async fn test_confirmation_manager_policy_update() {
639        let (event_tx, _) = broadcast::channel(100);
640        let manager = ConfirmationManager::new(ConfirmationPolicy::default(), event_tx);
641
642        // Initially disabled
643        assert!(!manager.requires_confirmation("bash").await);
644
645        // Update policy to enabled
646        manager.set_policy(ConfirmationPolicy::enabled()).await;
647        assert!(manager.requires_confirmation("bash").await);
648
649        // Update policy with YOLO mode
650        manager
651            .set_policy(ConfirmationPolicy::enabled().with_yolo_lanes([SessionLane::Execute]))
652            .await;
653        assert!(!manager.requires_confirmation("bash").await);
654    }
655
656    // ========================================================================
657    // Confirmation Flow Tests
658    // ========================================================================
659
660    #[tokio::test]
661    async fn test_confirmation_flow_approve() {
662        let (event_tx, mut event_rx) = broadcast::channel(100);
663        let manager = ConfirmationManager::new(ConfirmationPolicy::enabled(), event_tx);
664
665        // Request confirmation
666        let rx = manager
667            .request_confirmation("tool-1", "bash", &serde_json::json!({"command": "ls"}))
668            .await;
669
670        // Check event was emitted
671        let event = event_rx.recv().await.unwrap();
672        match event {
673            AgentEvent::ConfirmationRequired {
674                tool_id,
675                tool_name,
676                timeout_ms,
677                ..
678            } => {
679                assert_eq!(tool_id, "tool-1");
680                assert_eq!(tool_name, "bash");
681                assert_eq!(timeout_ms, 30_000); // Default timeout
682            }
683            _ => panic!("Expected ConfirmationRequired event"),
684        }
685
686        // Approve the confirmation
687        let result = manager.confirm("tool-1", true, None).await;
688        assert!(result.is_ok());
689        assert!(result.unwrap());
690
691        // Check ConfirmationReceived event
692        let event = event_rx.recv().await.unwrap();
693        match event {
694            AgentEvent::ConfirmationReceived {
695                tool_id, approved, ..
696            } => {
697                assert_eq!(tool_id, "tool-1");
698                assert!(approved);
699            }
700            _ => panic!("Expected ConfirmationReceived event"),
701        }
702
703        // Check response
704        let response = rx.await.unwrap();
705        assert!(response.approved);
706        assert!(response.reason.is_none());
707    }
708
709    #[tokio::test]
710    async fn test_confirmation_flow_reject() {
711        let (event_tx, mut event_rx) = broadcast::channel(100);
712        let manager = ConfirmationManager::new(ConfirmationPolicy::enabled(), event_tx);
713
714        // Request confirmation
715        let rx = manager
716            .request_confirmation(
717                "tool-1",
718                "bash",
719                &serde_json::json!({"command": "rm -rf /"}),
720            )
721            .await;
722
723        // Skip ConfirmationRequired event
724        let _ = event_rx.recv().await.unwrap();
725
726        // Reject the confirmation with reason
727        let result = manager
728            .confirm("tool-1", false, Some("Dangerous command".to_string()))
729            .await;
730        assert!(result.is_ok());
731        assert!(result.unwrap());
732
733        // Check ConfirmationReceived event
734        let event = event_rx.recv().await.unwrap();
735        match event {
736            AgentEvent::ConfirmationReceived {
737                tool_id,
738                approved,
739                reason,
740            } => {
741                assert_eq!(tool_id, "tool-1");
742                assert!(!approved);
743                assert_eq!(reason, Some("Dangerous command".to_string()));
744            }
745            _ => panic!("Expected ConfirmationReceived event"),
746        }
747
748        // Check response
749        let response = rx.await.unwrap();
750        assert!(!response.approved);
751        assert_eq!(response.reason, Some("Dangerous command".to_string()));
752    }
753
754    #[tokio::test]
755    async fn test_confirmation_not_found() {
756        let (event_tx, _) = broadcast::channel(100);
757        let manager = ConfirmationManager::new(ConfirmationPolicy::enabled(), event_tx);
758
759        // Try to confirm non-existent confirmation
760        let result = manager.confirm("non-existent", true, None).await;
761        assert!(result.is_ok());
762        assert!(!result.unwrap()); // Returns false for not found
763    }
764
765    // ========================================================================
766    // Multiple Confirmations Tests
767    // ========================================================================
768
769    #[tokio::test]
770    async fn test_multiple_confirmations() {
771        let (event_tx, _) = broadcast::channel(100);
772        let manager = ConfirmationManager::new(ConfirmationPolicy::enabled(), event_tx);
773
774        // Request multiple confirmations
775        let rx1 = manager
776            .request_confirmation("tool-1", "bash", &serde_json::json!({"cmd": "1"}))
777            .await;
778        let rx2 = manager
779            .request_confirmation("tool-2", "write", &serde_json::json!({"cmd": "2"}))
780            .await;
781        let rx3 = manager
782            .request_confirmation("tool-3", "edit", &serde_json::json!({"cmd": "3"}))
783            .await;
784
785        // Check pending count
786        assert_eq!(manager.pending_count().await, 3);
787
788        // Approve tool-1
789        manager.confirm("tool-1", true, None).await.unwrap();
790        let response1 = rx1.await.unwrap();
791        assert!(response1.approved);
792
793        // Reject tool-2
794        manager.confirm("tool-2", false, None).await.unwrap();
795        let response2 = rx2.await.unwrap();
796        assert!(!response2.approved);
797
798        // Approve tool-3
799        manager.confirm("tool-3", true, None).await.unwrap();
800        let response3 = rx3.await.unwrap();
801        assert!(response3.approved);
802
803        // All confirmations processed
804        assert_eq!(manager.pending_count().await, 0);
805    }
806
807    #[tokio::test]
808    async fn test_pending_confirmations_info() {
809        let (event_tx, _) = broadcast::channel(100);
810        let manager = ConfirmationManager::new(ConfirmationPolicy::enabled(), event_tx);
811
812        // Request confirmations
813        let _rx1 = manager
814            .request_confirmation("tool-1", "bash", &serde_json::json!({}))
815            .await;
816        let _rx2 = manager
817            .request_confirmation("tool-2", "write", &serde_json::json!({}))
818            .await;
819
820        let pending = manager.pending_confirmations().await;
821        assert_eq!(pending.len(), 2);
822
823        // Check that both tools are in pending list
824        let tool_ids: Vec<&str> = pending.iter().map(|(id, _, _)| id.as_str()).collect();
825        assert!(tool_ids.contains(&"tool-1"));
826        assert!(tool_ids.contains(&"tool-2"));
827    }
828
829    // ========================================================================
830    // Cancel Tests
831    // ========================================================================
832
833    #[tokio::test]
834    async fn test_cancel_confirmation() {
835        let (event_tx, _) = broadcast::channel(100);
836        let manager = ConfirmationManager::new(ConfirmationPolicy::enabled(), event_tx);
837
838        // Request confirmation
839        let rx = manager
840            .request_confirmation("tool-1", "bash", &serde_json::json!({}))
841            .await;
842
843        assert_eq!(manager.pending_count().await, 1);
844
845        // Cancel confirmation
846        let cancelled = manager.cancel("tool-1").await;
847        assert!(cancelled);
848        assert_eq!(manager.pending_count().await, 0);
849
850        // Check response indicates cancellation
851        let response = rx.await.unwrap();
852        assert!(!response.approved);
853        assert_eq!(response.reason, Some("Confirmation cancelled".to_string()));
854    }
855
856    #[tokio::test]
857    async fn test_cancel_nonexistent() {
858        let (event_tx, _) = broadcast::channel(100);
859        let manager = ConfirmationManager::new(ConfirmationPolicy::enabled(), event_tx);
860
861        let cancelled = manager.cancel("non-existent").await;
862        assert!(!cancelled);
863    }
864
865    #[tokio::test]
866    async fn test_cancel_all() {
867        let (event_tx, _) = broadcast::channel(100);
868        let manager = ConfirmationManager::new(ConfirmationPolicy::enabled(), event_tx);
869
870        // Request multiple confirmations
871        let rx1 = manager
872            .request_confirmation("tool-1", "bash", &serde_json::json!({}))
873            .await;
874        let rx2 = manager
875            .request_confirmation("tool-2", "write", &serde_json::json!({}))
876            .await;
877        let rx3 = manager
878            .request_confirmation("tool-3", "edit", &serde_json::json!({}))
879            .await;
880
881        assert_eq!(manager.pending_count().await, 3);
882
883        // Cancel all
884        let cancelled_count = manager.cancel_all().await;
885        assert_eq!(cancelled_count, 3);
886        assert_eq!(manager.pending_count().await, 0);
887
888        // All responses should indicate cancellation
889        for rx in [rx1, rx2, rx3] {
890            let response = rx.await.unwrap();
891            assert!(!response.approved);
892            assert_eq!(response.reason, Some("Confirmation cancelled".to_string()));
893        }
894    }
895
896    // ========================================================================
897    // Timeout Tests
898    // ========================================================================
899
900    #[tokio::test]
901    async fn test_timeout_reject() {
902        let (event_tx, mut event_rx) = broadcast::channel(100);
903        let policy = ConfirmationPolicy {
904            enabled: true,
905            default_timeout_ms: 50, // Very short timeout for testing
906            timeout_action: TimeoutAction::Reject,
907            ..Default::default()
908        };
909        let manager = ConfirmationManager::new(policy, event_tx);
910
911        // Request confirmation
912        let rx = manager
913            .request_confirmation("tool-1", "bash", &serde_json::json!({}))
914            .await;
915
916        // Skip ConfirmationRequired event
917        let _ = event_rx.recv().await.unwrap();
918
919        // Wait for timeout
920        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
921
922        // Check timeouts
923        let timed_out = manager.check_timeouts().await;
924        assert_eq!(timed_out, 1);
925
926        // Check timeout event
927        let event = event_rx.recv().await.unwrap();
928        match event {
929            AgentEvent::ConfirmationTimeout {
930                tool_id,
931                action_taken,
932            } => {
933                assert_eq!(tool_id, "tool-1");
934                assert_eq!(action_taken, "rejected");
935            }
936            _ => panic!("Expected ConfirmationTimeout event"),
937        }
938
939        // Check response indicates timeout rejection
940        let response = rx.await.unwrap();
941        assert!(!response.approved);
942        assert!(response.reason.as_ref().unwrap().contains("timed out"));
943    }
944
945    #[tokio::test]
946    async fn test_timeout_auto_approve() {
947        let (event_tx, mut event_rx) = broadcast::channel(100);
948        let policy = ConfirmationPolicy {
949            enabled: true,
950            default_timeout_ms: 50, // Very short timeout for testing
951            timeout_action: TimeoutAction::AutoApprove,
952            ..Default::default()
953        };
954        let manager = ConfirmationManager::new(policy, event_tx);
955
956        // Request confirmation
957        let rx = manager
958            .request_confirmation("tool-1", "bash", &serde_json::json!({}))
959            .await;
960
961        // Skip ConfirmationRequired event
962        let _ = event_rx.recv().await.unwrap();
963
964        // Wait for timeout
965        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
966
967        // Check timeouts
968        let timed_out = manager.check_timeouts().await;
969        assert_eq!(timed_out, 1);
970
971        // Check timeout event
972        let event = event_rx.recv().await.unwrap();
973        match event {
974            AgentEvent::ConfirmationTimeout {
975                tool_id,
976                action_taken,
977            } => {
978                assert_eq!(tool_id, "tool-1");
979                assert_eq!(action_taken, "auto_approved");
980            }
981            _ => panic!("Expected ConfirmationTimeout event"),
982        }
983
984        // Check response indicates timeout auto-approval
985        let response = rx.await.unwrap();
986        assert!(response.approved);
987        assert!(response.reason.as_ref().unwrap().contains("auto_approved"));
988    }
989
990    #[tokio::test]
991    async fn test_no_timeout_when_confirmed() {
992        let (event_tx, _) = broadcast::channel(100);
993        let policy = ConfirmationPolicy {
994            enabled: true,
995            default_timeout_ms: 50,
996            timeout_action: TimeoutAction::Reject,
997            ..Default::default()
998        };
999        let manager = ConfirmationManager::new(policy, event_tx);
1000
1001        // Request confirmation
1002        let rx = manager
1003            .request_confirmation("tool-1", "bash", &serde_json::json!({}))
1004            .await;
1005
1006        // Confirm immediately
1007        manager.confirm("tool-1", true, None).await.unwrap();
1008
1009        // Wait past timeout
1010        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
1011
1012        // Check timeouts - should be 0 since already confirmed
1013        let timed_out = manager.check_timeouts().await;
1014        assert_eq!(timed_out, 0);
1015
1016        // Response should be approval (not timeout)
1017        let response = rx.await.unwrap();
1018        assert!(response.approved);
1019        assert!(response.reason.is_none());
1020    }
1021}