magi-code 0.77.1

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

#[derive(Debug, Clone)]
pub(crate) struct AgentRunCanceled;

impl std::fmt::Display for AgentRunCanceled {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str("prompt canceled")
    }
}

impl std::error::Error for AgentRunCanceled {}

pub(crate) fn is_run_canceled(error: &anyhow::Error) -> bool {
    error.downcast_ref::<AgentRunCanceled>().is_some()
}

#[derive(Debug, Clone)]
pub(crate) struct AgentCancellationHandle {
    flag: Arc<AtomicBool>,
}

impl AgentCancellationHandle {
    pub(crate) fn cancel(&self) {
        self.flag.store(true, Ordering::SeqCst);
    }
}

#[derive(Debug, Clone, Default)]
pub(crate) struct AgentCancellation {
    flags: Vec<Arc<AtomicBool>>,
}

impl AgentCancellation {
    pub(crate) fn new(flag: Arc<AtomicBool>) -> Self {
        Self { flags: vec![flag] }
    }

    pub(crate) fn child_token(&self) -> (Self, AgentCancellationHandle) {
        let child_flag = Arc::new(AtomicBool::new(false));
        let mut flags = self.flags.clone();
        flags.push(Arc::clone(&child_flag));
        (Self { flags }, AgentCancellationHandle { flag: child_flag })
    }

    pub(crate) fn is_canceled(&self) -> bool {
        self.flags.iter().any(|flag| flag.load(Ordering::SeqCst))
    }

    pub(crate) fn check(&self) -> anyhow::Result<()> {
        if self.is_canceled() {
            Err(AgentRunCanceled.into())
        } else {
            Ok(())
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn cancellation_error_preserves_display_and_downcast() {
        let error: anyhow::Error = AgentRunCanceled.into();

        assert_eq!(error.to_string(), "prompt canceled");
        assert!(is_run_canceled(&error));
    }

    #[test]
    fn parent_cancellation_propagates_to_child() {
        let parent_flag = Arc::new(AtomicBool::new(false));
        let parent = AgentCancellation::new(Arc::clone(&parent_flag));
        let (child, _child_handle) = parent.child_token();

        parent_flag.store(true, Ordering::SeqCst);

        assert!(parent.is_canceled());
        assert!(child.is_canceled());
        assert!(child.check().is_err());
    }

    #[test]
    fn child_cancellation_does_not_cancel_parent() {
        let parent_flag = Arc::new(AtomicBool::new(false));
        let parent = AgentCancellation::new(parent_flag);
        let (child, child_handle) = parent.child_token();

        child_handle.cancel();

        assert!(child.is_canceled());
        assert!(child.check().is_err());
        assert!(!parent.is_canceled());
        assert!(parent.check().is_ok());
    }
}