Skip to main content

meerkat_mobkit/memory/
guards.rs

1//! Background-work resource guards (§8.1).
2//!
3//! Every judgment stage burns LLM calls, and MobKit multiplies stages by
4//! identities × interactions × mobs. Stage-level throttles scale *with*
5//! activity; these guards are the deterministic containment on top: hard
6//! per-window caps on background runs and a concurrency ceiling, consulted
7//! before every run. A skipped run is loud: a tracing warn always, plus a
8//! `memory.budget.denied` timeline event when a sink is wired (Principle 6).
9//!
10//! The load-*inverse* control Codex ships (skip background work below
11//! provider rate-limit headroom) needs a provider-quota surface MobKit does
12//! not have; until then the per-realm window budget is the stand-in, and
13//! the §16 open question on default budgets is answered by measurement.
14
15use std::collections::HashMap;
16use std::sync::{Arc, Mutex};
17use std::time::{Duration, Instant};
18
19use crate::memory::events::{MemoryEventSink, MemoryTimelineEvent};
20
21/// Default hard cap on background runs per realm per window (§8.1; the
22/// concrete number is §16 open question 5, this is the measured starting
23/// point).
24pub const DEFAULT_RUNS_PER_HOUR: u32 = 12;
25/// Default background-run concurrency per realm.
26pub const DEFAULT_MAX_CONCURRENT: u32 = 1;
27
28#[derive(Debug, Clone, PartialEq, Eq)]
29pub struct BackgroundBudgetConfig {
30    pub runs_per_window: u32,
31    pub window: Duration,
32    pub max_concurrent: u32,
33}
34
35impl Default for BackgroundBudgetConfig {
36    fn default() -> Self {
37        Self {
38            runs_per_window: DEFAULT_RUNS_PER_HOUR,
39            window: Duration::from_hours(1),
40            max_concurrent: DEFAULT_MAX_CONCURRENT,
41        }
42    }
43}
44
45#[derive(Debug, Default)]
46struct RealmBudgetState {
47    /// Start times of runs admitted within the sliding window.
48    starts: Vec<Instant>,
49    concurrent: u32,
50}
51
52struct BudgetInner {
53    config: BackgroundBudgetConfig,
54    realms: HashMap<String, RealmBudgetState>,
55    event_sink: Option<Arc<dyn MemoryEventSink>>,
56}
57
58/// Per-realm background budget (§8.1): a sliding-window run cap plus a
59/// concurrency ceiling. Cheap to clone; clones share state.
60#[derive(Clone)]
61pub struct BackgroundBudget {
62    inner: Arc<Mutex<BudgetInner>>,
63}
64
65/// RAII permit for one admitted background run; dropping it releases the
66/// concurrency slot (the window slot is consumed permanently).
67pub struct BudgetPermit {
68    inner: Arc<Mutex<BudgetInner>>,
69    realm: String,
70}
71
72impl std::fmt::Debug for BudgetPermit {
73    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74        f.debug_struct("BudgetPermit")
75            .field("realm", &self.realm)
76            .finish_non_exhaustive()
77    }
78}
79
80impl Drop for BudgetPermit {
81    fn drop(&mut self) {
82        let mut inner = self
83            .inner
84            .lock()
85            .unwrap_or_else(std::sync::PoisonError::into_inner);
86        if let Some(state) = inner.realms.get_mut(&self.realm) {
87            state.concurrent = state.concurrent.saturating_sub(1);
88        }
89    }
90}
91
92/// Why a run was denied. Carried in the warn log and (P3b) timeline event.
93#[derive(Debug, Clone, PartialEq, Eq)]
94pub enum BudgetDenied {
95    WindowExhausted { used: u32, cap: u32 },
96    ConcurrencyCeiling { cap: u32 },
97}
98
99impl std::fmt::Display for BudgetDenied {
100    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
101        match self {
102            Self::WindowExhausted { used, cap } => {
103                write!(f, "window budget exhausted ({used}/{cap} runs)")
104            }
105            Self::ConcurrencyCeiling { cap } => {
106                write!(f, "concurrency ceiling reached ({cap} in flight)")
107            }
108        }
109    }
110}
111
112impl BackgroundBudget {
113    pub fn new(config: BackgroundBudgetConfig) -> Self {
114        Self {
115            inner: Arc::new(Mutex::new(BudgetInner {
116                config,
117                realms: HashMap::new(),
118                event_sink: None,
119            })),
120        }
121    }
122
123    /// Wire the §9.3 timeline sink so guard denials surface on the console
124    /// alongside the tracing warn. Shared across clones.
125    pub fn set_event_sink(&self, sink: Arc<dyn MemoryEventSink>) {
126        self.inner
127            .lock()
128            .unwrap_or_else(std::sync::PoisonError::into_inner)
129            .event_sink = Some(sink);
130    }
131
132    /// Admit one background run for `realm`, or say loudly why not.
133    /// `stage` is only for the log line.
134    pub fn try_acquire(&self, realm: &str, stage: &str) -> Result<BudgetPermit, BudgetDenied> {
135        let mut inner = self
136            .inner
137            .lock()
138            .unwrap_or_else(std::sync::PoisonError::into_inner);
139        let window = inner.config.window;
140        let runs_cap = inner.config.runs_per_window;
141        let concurrent_cap = inner.config.max_concurrent;
142        let state = inner.realms.entry(realm.to_string()).or_default();
143        let now = Instant::now();
144        state
145            .starts
146            .retain(|start| now.duration_since(*start) < window);
147        let denied = if state.concurrent >= concurrent_cap {
148            Some(BudgetDenied::ConcurrencyCeiling {
149                cap: concurrent_cap,
150            })
151        } else if state.starts.len() as u32 >= runs_cap {
152            Some(BudgetDenied::WindowExhausted {
153                used: state.starts.len() as u32,
154                cap: runs_cap,
155            })
156        } else {
157            None
158        };
159        if let Some(denied) = denied {
160            tracing::warn!(
161                realm,
162                stage,
163                reason = %denied,
164                "agent memory background budget: run skipped"
165            );
166            if let Some(sink) = inner.event_sink.as_ref() {
167                sink.emit(MemoryTimelineEvent::BudgetDenied {
168                    realm: realm.to_string(),
169                    stage: stage.to_string(),
170                    reason: denied.to_string(),
171                });
172            }
173            return Err(denied);
174        }
175        state.starts.push(now);
176        state.concurrent += 1;
177        Ok(BudgetPermit {
178            inner: self.inner.clone(),
179            realm: realm.to_string(),
180        })
181    }
182}
183
184#[cfg(test)]
185#[allow(clippy::expect_used, clippy::unwrap_used)]
186mod tests {
187    use super::*;
188
189    fn config(runs: u32, concurrent: u32) -> BackgroundBudgetConfig {
190        BackgroundBudgetConfig {
191            runs_per_window: runs,
192            window: Duration::from_hours(1),
193            max_concurrent: concurrent,
194        }
195    }
196
197    #[test]
198    fn window_cap_denies_after_budget_spent() {
199        let budget = BackgroundBudget::new(config(2, 10));
200        let p1 = budget.try_acquire("realm-a", "distiller").expect("run 1");
201        drop(p1);
202        let p2 = budget.try_acquire("realm-a", "distiller").expect("run 2");
203        drop(p2);
204        let denied = budget
205            .try_acquire("realm-a", "distiller")
206            .expect_err("third run in window must deny");
207        assert!(
208            matches!(denied, BudgetDenied::WindowExhausted { used: 2, cap: 2 }),
209            "{denied:?}"
210        );
211        // Budgets are per realm: another realm is unaffected.
212        budget
213            .try_acquire("realm-b", "distiller")
214            .expect("other realm has its own window");
215    }
216
217    #[test]
218    fn denial_emits_timeline_event_when_sink_wired() {
219        let budget = BackgroundBudget::new(config(1, 10));
220        let sink = Arc::new(crate::memory::events::CollectingEventSink::new());
221        budget.set_event_sink(sink.clone());
222        let _permit = budget.try_acquire("realm-a", "steward").expect("first");
223        let _ = budget
224            .try_acquire("realm-a", "steward")
225            .expect_err("window spent");
226        assert_eq!(sink.types(), vec!["memory.budget.denied"]);
227        let events = sink.events.lock().unwrap();
228        assert!(matches!(
229            &events[0],
230            MemoryTimelineEvent::BudgetDenied { realm, stage, .. }
231                if realm == "realm-a" && stage == "steward"
232        ));
233    }
234
235    #[test]
236    fn concurrency_ceiling_releases_on_drop() {
237        let budget = BackgroundBudget::new(config(10, 1));
238        let permit = budget.try_acquire("realm-a", "distiller").expect("first");
239        let denied = budget
240            .try_acquire("realm-a", "distiller")
241            .expect_err("second concurrent run must deny");
242        assert!(
243            matches!(denied, BudgetDenied::ConcurrencyCeiling { cap: 1 }),
244            "{denied:?}"
245        );
246        drop(permit);
247        budget
248            .try_acquire("realm-a", "distiller")
249            .expect("slot released on drop");
250    }
251}