use crate::budget::{BudgetError, BudgetLimits, CoordinationBudget};
use crate::types::AgentOutput;
use car_engine::Runtime;
use car_eventlog::EventLog;
use car_policy::PolicyEngine;
use car_state::StateStore;
use std::sync::Arc;
use tokio::sync::{Mutex as TokioMutex, RwLock as TokioRwLock};
pub struct SharedInfra {
pub state: Arc<StateStore>,
pub log: Arc<TokioMutex<EventLog>>,
pub policies: Arc<TokioRwLock<PolicyEngine>>,
pub budget: Arc<CoordinationBudget>,
}
impl SharedInfra {
pub fn new() -> Self {
Self {
state: Arc::new(StateStore::new()),
log: Arc::new(TokioMutex::new(EventLog::new())),
policies: Arc::new(TokioRwLock::new(PolicyEngine::new())),
budget: Arc::new(CoordinationBudget::unbounded()),
}
}
pub fn with_budget(mut self, limits: BudgetLimits) -> Self {
self.budget = Arc::new(CoordinationBudget::new(limits));
self
}
pub fn with_shared_budget(mut self, budget: Arc<CoordinationBudget>) -> Self {
self.budget = budget;
self
}
pub fn begin_agent(&self) -> Result<(), BudgetError> {
self.budget.try_begin_agent()
}
pub fn record_output(&self, out: &AgentOutput) {
self.budget.record_output(out);
}
pub fn make_runtime(&self) -> Runtime {
Runtime::with_shared(
Arc::clone(&self.state),
Arc::clone(&self.log),
Arc::clone(&self.policies),
)
}
pub fn make_isolated_runtime(
&self,
agent_name: &str,
) -> (Runtime, crate::task_context::AgentContext) {
let ctx = crate::task_context::AgentContext::new(agent_name, Arc::clone(&self.state));
let rt = Runtime::with_shared(
Arc::clone(&ctx.local_state),
Arc::clone(&ctx.local_log),
Arc::clone(&self.policies),
);
(rt, ctx)
}
}
impl Default for SharedInfra {
fn default() -> Self {
Self::new()
}
}