Skip to main content

starweaver_model/adapter/
guard.rs

1use std::sync::atomic::{AtomicBool, Ordering};
2
3static ALLOW_REAL_MODEL_REQUESTS: AtomicBool = AtomicBool::new(true);
4
5/// Return whether production model requests are globally allowed.
6#[must_use]
7pub fn allow_real_model_requests() -> bool {
8    ALLOW_REAL_MODEL_REQUESTS.load(Ordering::SeqCst)
9}
10
11/// Set whether production model requests are globally allowed.
12pub fn set_allow_real_model_requests(allow: bool) {
13    ALLOW_REAL_MODEL_REQUESTS.store(allow, Ordering::SeqCst);
14}
15
16/// Scoped guard that restores the previous production-request setting when dropped.
17#[derive(Debug)]
18pub struct RealModelRequestGuard {
19    previous: bool,
20}
21
22impl RealModelRequestGuard {
23    /// Set the production-request setting for this scope.
24    #[must_use]
25    pub fn set(allow: bool) -> Self {
26        let previous = ALLOW_REAL_MODEL_REQUESTS.swap(allow, Ordering::SeqCst);
27        Self { previous }
28    }
29}
30
31impl Drop for RealModelRequestGuard {
32    fn drop(&mut self) {
33        ALLOW_REAL_MODEL_REQUESTS.store(self.previous, Ordering::SeqCst);
34    }
35}
36
37/// Block production model requests until the returned guard is dropped.
38#[must_use]
39pub fn block_real_model_requests() -> RealModelRequestGuard {
40    RealModelRequestGuard::set(false)
41}
42
43/// Allow production model requests until the returned guard is dropped.
44#[must_use]
45pub fn allow_real_model_requests_guard() -> RealModelRequestGuard {
46    RealModelRequestGuard::set(true)
47}