codex/codex-rs/ext/extension-api/tests/state.rs

use std::panic::AssertUnwindSafe;
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;

use codex_extension_api::ExtensionData;
use pretty_assertions::assert_eq;

#[test]
fn typed_values_can_be_inserted_replaced_and_removed() {
    let data = ExtensionData::new("thread-1");

    assert_eq!(data.insert(/*value*/ 41_u64), None);
    assert_eq!(data.insert("alpha".to_string()), None);
    assert_eq!(data.get::<u64>().as_deref(), Some(&41));
    assert_eq!(
        data.get::<String>().map(|value| value.as_str().to_string()),
        Some("alpha".to_string())
    );

    assert_eq!(data.insert(/*value*/ 42_u64).as_deref(), Some(&41));
    assert_eq!(data.get::<u64>().as_deref(), Some(&42));
    assert_eq!(
        data.remove::<String>()
            .map(|value| value.as_str().to_string()),
        Some("alpha".to_string())
    );
    assert_eq!(data.get::<String>(), None);
    assert_eq!(data.get::<u64>().as_deref(), Some(&42));
}

#[test]
fn get_or_init_initializes_once_and_returns_shared_value() {
    const CALLER_COUNT: usize = 8;

    #[derive(Debug, PartialEq, Eq)]
    struct SharedValue(usize);

    let data = Arc::new(ExtensionData::new("session"));
    let callers_started = Arc::new(AtomicUsize::new(0));
    let initialization_count = Arc::new(AtomicUsize::new(0));

    let handles: [_; CALLER_COUNT] = std::array::from_fn(|_| {
        let data = Arc::clone(&data);
        let callers_started = Arc::clone(&callers_started);
        let initialization_count = Arc::clone(&initialization_count);
        std::thread::spawn(move || {
            callers_started.fetch_add(1, Ordering::SeqCst);
            data.get_or_init(|| {
                initialization_count.fetch_add(1, Ordering::SeqCst);
                // Keep the first initializer active until every worker has attempted
                // get_or_init, forcing callers to overlap on the same missing entry.
                while callers_started.load(Ordering::SeqCst) < CALLER_COUNT {
                    std::thread::yield_now();
                }
                SharedValue(7)
            })
        })
    });
    let values = handles
        .into_iter()
        .map(|handle| handle.join().expect("initializer thread should succeed"))
        .collect::<Vec<_>>();

    assert_eq!(initialization_count.load(Ordering::SeqCst), 1);
    assert_eq!(
        values.iter().map(Arc::as_ref).collect::<Vec<_>>(),
        vec![&SharedValue(7); CALLER_COUNT]
    );
    assert!(
        values
            .iter()
            .skip(1)
            .all(|value| Arc::ptr_eq(&values[0], value))
    );
}

#[test]
fn stores_are_isolated_and_preserve_level_id() {
    let session_data = ExtensionData::new("root-1");
    let thread_data = ExtensionData::new("root-1");

    session_data.insert(/*value*/ 17_u32);
    thread_data.insert("thread value".to_string());

    assert_eq!(session_data.level_id(), "root-1");
    assert_eq!(thread_data.level_id(), "root-1");
    assert_eq!(session_data.get::<u32>().as_deref(), Some(&17));
    assert_eq!(session_data.get::<String>(), None);
    assert_eq!(thread_data.get::<u32>(), None);
    assert_eq!(
        thread_data
            .get::<String>()
            .map(|value| value.as_str().to_string()),
        Some("thread value".to_string())
    );
}

#[test]
fn store_remains_usable_after_panicking_initializer() {
    let data = ExtensionData::new("turn-1");

    let result = std::panic::catch_unwind(AssertUnwindSafe(|| {
        data.get_or_init::<u64>(|| panic!("initializer failed"));
    }));

    assert!(result.is_err());
    assert_eq!(*data.get_or_init(|| 99_u64), 99);
}