otherone_agent/multi_agent/
supervisor.rs

1use std::collections::{HashMap, HashSet};
2use std::sync::atomic::{AtomicUsize, Ordering};
3use std::sync::{Arc, Mutex, OnceLock};
4use std::time::Instant;
5
6use tokio::sync::{mpsc, RwLock, Semaphore};
7use tokio_util::sync::CancellationToken;
8
9use crate::error::AgentError;
10
11use super::event::EventBus;
12use super::registry::RuntimeSnapshot;
13use super::types::{
14    AgentCallId, AgentDefinition, AgentId, AgentRunCommand, EffectiveLimits, ModelProfile,
15    ResultContract, RunId, RunRecord, RunStatus, RunUsage, SessionTarget,
16};
17
18#[derive(Debug, Default)]
19struct BudgetState {
20    total_runs: usize,
21    model_calls: u32,
22    tool_calls: u32,
23    agent_calls: u32,
24    total_tokens: u64,
25    output_bytes: usize,
26}
27
28/// 一个根任务共享的原子预算账本。
29pub struct BudgetLedger {
30    limits: EffectiveLimits,
31    state: Mutex<BudgetState>,
32}
33
34impl BudgetLedger {
35    pub(crate) fn new(limits: EffectiveLimits) -> Self {
36        Self {
37            limits,
38            state: Mutex::new(BudgetState::default()),
39        }
40    }
41
42    pub(crate) fn limits(&self) -> &EffectiveLimits {
43        &self.limits
44    }
45
46    pub(crate) fn reserve_run(&self) -> Result<(), AgentError> {
47        let mut state = self.lock_state()?;
48        if state.total_runs >= self.limits.runtime.max_total_runs {
49            return Err(AgentError::BudgetExceeded("max_total_runs".to_string()));
50        }
51        state.total_runs += 1;
52        Ok(())
53    }
54
55    pub(crate) fn reserve_model_call(&self) -> Result<(), AgentError> {
56        let mut state = self.lock_state()?;
57        if state.model_calls >= self.limits.runtime.max_model_calls {
58            return Err(AgentError::BudgetExceeded("max_model_calls".to_string()));
59        }
60        state.model_calls += 1;
61        Ok(())
62    }
63
64    pub(crate) fn reserve_tool_call(&self) -> Result<(), AgentError> {
65        let mut state = self.lock_state()?;
66        if state.tool_calls >= self.limits.runtime.max_tool_calls {
67            return Err(AgentError::BudgetExceeded("max_tool_calls".to_string()));
68        }
69        state.tool_calls += 1;
70        Ok(())
71    }
72
73    pub(crate) fn reserve_agent_call(&self) -> Result<(), AgentError> {
74        let mut state = self.lock_state()?;
75        if state.agent_calls >= self.limits.runtime.max_agent_calls {
76            return Err(AgentError::BudgetExceeded("max_agent_calls".to_string()));
77        }
78        state.agent_calls += 1;
79        Ok(())
80    }
81
82    pub(crate) fn add_tokens(&self, tokens: u64) -> Result<(), AgentError> {
83        let mut state = self.lock_state()?;
84        state.total_tokens = state.total_tokens.saturating_add(tokens);
85        if self
86            .limits
87            .runtime
88            .token_budget
89            .is_some_and(|limit| state.total_tokens > limit)
90        {
91            return Err(AgentError::BudgetExceeded("token_budget".to_string()));
92        }
93        Ok(())
94    }
95
96    pub(crate) fn add_output_bytes(&self, bytes: usize) -> Result<(), AgentError> {
97        let mut state = self.lock_state()?;
98        state.output_bytes = state.output_bytes.saturating_add(bytes);
99        if state.output_bytes > self.limits.runtime.max_result_bytes {
100            return Err(AgentError::BudgetExceeded("max_result_bytes".to_string()));
101        }
102        Ok(())
103    }
104
105    fn lock_state(&self) -> Result<std::sync::MutexGuard<'_, BudgetState>, AgentError> {
106        self.state
107            .lock()
108            .map_err(|_| AgentError::Internal("budget ledger lock poisoned".to_string()))
109    }
110}
111
112#[derive(Clone, Default)]
113pub struct InMemoryRunStore {
114    records: Arc<RwLock<HashMap<RunId, RunRecord>>>,
115}
116
117impl InMemoryRunStore {
118    pub async fn create(&self, record: RunRecord) -> Result<(), AgentError> {
119        let mut records = self.records.write().await;
120        if records.contains_key(&record.run_id) {
121            return Err(AgentError::Internal(format!(
122                "run '{}' already exists",
123                record.run_id
124            )));
125        }
126        records.insert(record.run_id.clone(), record);
127        Ok(())
128    }
129
130    pub async fn update_status(&self, run_id: &RunId, status: RunStatus) -> Result<(), AgentError> {
131        let mut records = self.records.write().await;
132        let record = records
133            .get_mut(run_id)
134            .ok_or_else(|| AgentError::Internal(format!("run '{run_id}' not found")))?;
135        record.status = status;
136        Ok(())
137    }
138
139    pub async fn finish(
140        &self,
141        run_id: &RunId,
142        status: RunStatus,
143        usage: RunUsage,
144        failure: Option<super::types::AgentFailure>,
145    ) -> Result<(), AgentError> {
146        let mut records = self.records.write().await;
147        let record = records
148            .get_mut(run_id)
149            .ok_or_else(|| AgentError::Internal(format!("run '{run_id}' not found")))?;
150        record.status = status;
151        record.usage = usage;
152        record.failure = failure;
153        record.finished_at = Some(chrono::Utc::now().to_rfc3339());
154        Ok(())
155    }
156
157    pub async fn get(&self, run_id: &RunId) -> Option<RunRecord> {
158        self.records.read().await.get(run_id).cloned()
159    }
160
161    pub async fn list_root(&self, root_run_id: &RunId) -> Vec<RunRecord> {
162        let mut records = self
163            .records
164            .read()
165            .await
166            .values()
167            .filter(|record| &record.root_run_id == root_run_id)
168            .cloned()
169            .collect::<Vec<_>>();
170        records.sort_by(|left, right| left.started_at.cmp(&right.started_at));
171        records
172    }
173}
174
175#[derive(Clone)]
176pub(crate) struct RunContext {
177    pub run_id: RunId,
178    pub root_run_id: RunId,
179    pub parent_run_id: Option<RunId>,
180    pub agent_call_id: Option<AgentCallId>,
181    pub agent: Arc<AgentDefinition>,
182    pub model: Arc<ModelProfile>,
183    pub session_id: String,
184    pub parent_session_id: Option<String>,
185    pub session: SessionTarget,
186    pub runtime_context: Option<otherone_storage::types::RuntimeContext>,
187    pub depth: usize,
188    pub ancestry: Vec<AgentId>,
189    pub deadline: Instant,
190    pub max_iterations_override: Option<u32>,
191    pub result_contract: Option<ResultContract>,
192    pub cancellation: CancellationToken,
193    pub budget: Arc<BudgetLedger>,
194    pub snapshot: Arc<RuntimeSnapshot>,
195    pub events: EventBus,
196    pub child_count: Arc<AtomicUsize>,
197    pub active_children: Arc<AtomicUsize>,
198    pub usage: Arc<Mutex<RunUsage>>,
199    pub metadata: otherone_storage::types::AttributeBag,
200}
201
202impl RunContext {
203    pub fn reserve_child(&self) -> Result<(), AgentError> {
204        let previous = self.child_count.fetch_add(1, Ordering::SeqCst);
205        if previous >= self.budget.limits().runtime.max_children_per_run {
206            self.child_count.fetch_sub(1, Ordering::SeqCst);
207            return Err(AgentError::BudgetExceeded(
208                "max_children_per_run".to_string(),
209            ));
210        }
211        Ok(())
212    }
213
214    pub fn view(&self) -> super::types::AgentRunContextView {
215        super::types::AgentRunContextView {
216            run_id: self.run_id.clone(),
217            root_run_id: self.root_run_id.clone(),
218            parent_run_id: self.parent_run_id.clone(),
219            agent_id: self.agent.id.clone(),
220            session_id: self.session_id.clone(),
221            depth: self.depth,
222            ancestry: self.ancestry.clone(),
223            runtime_context: self.runtime_context.clone(),
224        }
225    }
226
227    pub fn record_usage(&self, usage: &RunUsage) {
228        if let Ok(mut current) = self.usage.lock() {
229            *current = usage.clone();
230        }
231    }
232
233    pub fn usage_snapshot(&self) -> RunUsage {
234        self.usage
235            .lock()
236            .map(|usage| usage.clone())
237            .unwrap_or_default()
238    }
239}
240
241#[derive(Clone)]
242pub(crate) struct ActiveRun {
243    pub context: RunContext,
244    pub commands: mpsc::Sender<AgentRunCommand>,
245}
246
247pub(crate) struct RunSupervisor {
248    pub run_store: InMemoryRunStore,
249    pub model_semaphore: Arc<Semaphore>,
250    pub tool_semaphore: Arc<Semaphore>,
251    active_runs: RwLock<HashMap<RunId, ActiveRun>>,
252    session_writers: Arc<Mutex<HashSet<String>>>,
253}
254
255impl RunSupervisor {
256    pub fn new(limits: &super::types::RuntimeLimits) -> Self {
257        Self {
258            run_store: InMemoryRunStore::default(),
259            model_semaphore: Arc::new(Semaphore::new(limits.max_concurrent_model_calls)),
260            tool_semaphore: Arc::new(Semaphore::new(limits.max_concurrent_tool_calls)),
261            active_runs: RwLock::new(HashMap::new()),
262            session_writers: global_session_writers(),
263        }
264    }
265
266    pub async fn register(
267        &self,
268        context: RunContext,
269        commands: mpsc::Sender<AgentRunCommand>,
270    ) -> Result<SessionLease, AgentError> {
271        let lease = SessionLease::acquire(
272            Arc::clone(&self.session_writers),
273            session_key(&context),
274            context.session_id.clone(),
275        )?;
276        self.active_runs
277            .write()
278            .await
279            .insert(context.run_id.clone(), ActiveRun { context, commands });
280        Ok(lease)
281    }
282
283    pub async fn unregister(&self, run_id: &RunId) {
284        self.active_runs.write().await.remove(run_id);
285    }
286
287    pub async fn active(&self, run_id: &RunId) -> Option<ActiveRun> {
288        self.active_runs.read().await.get(run_id).cloned()
289    }
290
291    pub async fn cancel(&self, run_id: &RunId) -> Result<(), AgentError> {
292        let active = self.active(run_id).await.ok_or_else(|| {
293            AgentError::InvalidConfiguration(format!("run '{run_id}' is not active"))
294        })?;
295        active.context.cancellation.cancel();
296        Ok(())
297    }
298
299    pub async fn send_message(
300        &self,
301        run_id: &RunId,
302        messages: Vec<String>,
303    ) -> Result<(), AgentError> {
304        let active = self.active(run_id).await.ok_or_else(|| {
305            AgentError::InvalidConfiguration(format!("run '{run_id}' is not active"))
306        })?;
307        active
308            .commands
309            .send(AgentRunCommand::EnqueueUserMessages(messages))
310            .await
311            .map_err(|_| AgentError::InvalidConfiguration(format!("run '{run_id}' is closed")))
312    }
313}
314
315fn session_key(context: &RunContext) -> String {
316    let partition = context
317        .runtime_context
318        .as_ref()
319        .map(|context| context.partition_key.as_str())
320        .unwrap_or("local");
321    match context.session.storage_type {
322        crate::types::StorageType::LocalFile => format!(
323            "local:{}:{partition}:{}",
324            otherone_storage::localfile::get_storage_path().display(),
325            context.session_id
326        ),
327        crate::types::StorageType::Database => {
328            let database = context
329                .session
330                .database_config
331                .as_ref()
332                .map(|config| format!("{}:{}/{}", config.host, config.port, config.database))
333                .unwrap_or_else(|| "database".to_string());
334            format!("database:{database}:{partition}:{}", context.session_id)
335        }
336    }
337}
338
339fn global_session_writers() -> Arc<Mutex<HashSet<String>>> {
340    static WRITERS: OnceLock<Arc<Mutex<HashSet<String>>>> = OnceLock::new();
341    Arc::clone(WRITERS.get_or_init(|| Arc::new(Mutex::new(HashSet::new()))))
342}
343
344pub(crate) struct SessionLease {
345    writers: Arc<Mutex<HashSet<String>>>,
346    key: String,
347}
348
349impl SessionLease {
350    fn acquire(
351        writers: Arc<Mutex<HashSet<String>>>,
352        key: String,
353        session_id: String,
354    ) -> Result<Self, AgentError> {
355        {
356            let mut active = writers
357                .lock()
358                .map_err(|_| AgentError::Internal("session lock poisoned".to_string()))?;
359            if !active.insert(key.clone()) {
360                return Err(AgentError::SessionBusy(session_id));
361            }
362        }
363        Ok(Self { writers, key })
364    }
365}
366
367impl Drop for SessionLease {
368    fn drop(&mut self) {
369        if let Ok(mut writers) = self.writers.lock() {
370            writers.remove(&self.key);
371        }
372    }
373}
374
375#[cfg(test)]
376mod tests {
377    use super::*;
378
379    #[test]
380    fn budget_enforces_total_runs() {
381        let limits = super::super::types::RuntimeLimits {
382            max_total_runs: 1,
383            ..Default::default()
384        };
385        let ledger = BudgetLedger::new(EffectiveLimits { runtime: limits });
386        assert!(ledger.reserve_run().is_ok());
387        assert!(matches!(
388            ledger.reserve_run(),
389            Err(AgentError::BudgetExceeded(_))
390        ));
391    }
392}
otherone_agent/multi_agent/supervisor.rs

otherone_agent/multi_agent/
supervisor.rs
