ai-agent 0.13.4

//! Combined AbortSignal implementation
//!
//! Creates a combined AbortSignal that aborts when the input signal aborts,
//! an optional second signal aborts, or an optional timeout elapses.

use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::thread;
use std::time::Duration;

/// Result of create_combined_abort_signal
pub struct CombinedAbortSignal {
    /// The combined abort signal
    pub signal: Arc<AtomicBool>,
    /// Cleanup function to remove event listeners and clear the internal timeout timer
    pub cleanup: Box<dyn Fn() + Send + Sync>,
}

/// Creates a combined AbortSignal that aborts when the input signal aborts,
/// an optional second signal aborts, or an optional timeout elapses.
///
/// Use `timeout_ms` instead of passing `AbortSignal.timeout(ms)` as a signal —
/// this implementation uses `std::thread::sleep` + atomic flag so the timer
/// is freed immediately on cleanup.
pub fn create_combined_abort_signal(
    signal: Option<&Arc<AtomicBool>>,
    opts: Option<CombinedAbortSignalOpts>,
) -> CombinedAbortSignal {
    let signal_b = opts.as_ref().and_then(|o| o.signal_b.as_ref());
    let timeout_ms = opts.as_ref().and_then(|o| o.timeout_ms);

    let combined = Arc::new(AtomicBool::new(false));

    // Check if already aborted
    let is_aborted = signal.map(|s| s.load(Ordering::SeqCst)).unwrap_or(false)
        || signal_b.map(|s| s.load(Ordering::SeqCst)).unwrap_or(false);

    if is_aborted {
        combined.store(true, Ordering::SeqCst);
        return CombinedAbortSignal {
            signal: combined,
            cleanup: Box::new(|| {}),
        };
    }

    let combined_clone = combined.clone();
    let signal_clone = signal.map(|s| s.clone());
    let signal_b_clone = signal_b.cloned();

    // Set up abort handler
    let abort_closure = Box::new(move || {
        combined_clone.store(true, Ordering::SeqCst);
    });

    let mut abort_handles: Vec<Box<dyn Fn() + Send + Sync>> = Vec::new();

    // For simplicity in this implementation, we use a polling approach
    // The actual signal watching would be done by the caller
    // This is a simplified version that handles timeout

    let timer_handle: Option<thread::JoinHandle<()>> = if let Some(ms) = timeout_ms {
        let combined_timer = combined.clone();
        Some(thread::spawn(move || {
            thread::sleep(Duration::from_millis(ms));
            combined_timer.store(true, Ordering::SeqCst);
        }))
    } else {
        None
    };

    let timer_handle = std::sync::Mutex::new(timer_handle);
    let cleanup = Box::new(move || {
        if let Ok(mut handle) = timer_handle.lock() {
            if let Some(h) = handle.take() {
                let _ = h.join();
            }
        }
    });

    CombinedAbortSignal {
        signal: combined,
        cleanup,
    }
}

/// Options for create_combined_abort_signal
#[derive(Debug, Clone, Default)]
pub struct CombinedAbortSignalOpts {
    /// Optional second abort signal
    pub signal_b: Option<Arc<AtomicBool>>,
    /// Optional timeout in milliseconds
    pub timeout_ms: Option<u64>,
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::AtomicBool;
    use std::sync::Arc;
    use std::thread;
    use std::time::Duration;

    #[test]
    fn test_combined_abort_signal_no_opts() {
        let result = create_combined_abort_signal(None, None);
        assert!(!result.signal.load(Ordering::SeqCst));
    }

    #[test]
    fn test_combined_abort_signal_timeout() {
        let result = create_combined_abort_signal(
            None,
            Some(CombinedAbortSignalOpts {
                signal_b: None,
                timeout_ms: Some(10),
            }),
        );

        assert!(!result.signal.load(Ordering::SeqCst));

        // Wait for timeout
        thread::sleep(Duration::from_millis(20));

        assert!(result.signal.load(Ordering::SeqCst));
    }

    #[test]
    fn test_combined_abort_signal_already_aborted() {
        let aborted = Arc::new(AtomicBool::new(true));
        let result = create_combined_abort_signal(Some(&aborted), None);

        assert!(result.signal.load(Ordering::SeqCst));
    }

    #[test]
    fn test_combined_abort_signal_cleanup() {
        let result = create_combined_abort_signal(
            None,
            Some(CombinedAbortSignalOpts {
                signal_b: None,
                timeout_ms: Some(100),
            }),
        );

        // Call cleanup - should not panic
        (result.cleanup)();
    }
}