Skip to main content

aiclient_api/server/
state.rs

1use arc_swap::ArcSwap;
2use axum::extract::FromRef;
3use std::collections::HashMap;
4use std::sync::Arc;
5use tokio::sync::RwLock;
6use tokio::time::Instant;
7
8use crate::config::types::Config;
9use crate::providers::Provider;
10use crate::server::middleware::RateLimitMap;
11use crate::usage::UsageTracker;
12
13#[derive(Clone)]
14pub struct AppState {
15    pub config: Arc<ArcSwap<Config>>,
16    pub providers: Arc<RwLock<HashMap<String, Arc<dyn Provider>>>>,
17    pub start_time: Instant,
18    pub rate_limiter: RateLimitMap,
19    pub usage_tracker: UsageTracker,
20}
21
22impl AppState {
23    pub fn new(config: Config) -> Self {
24        Self {
25            config: Arc::new(ArcSwap::from_pointee(config)),
26            providers: Arc::new(RwLock::new(HashMap::new())),
27            start_time: Instant::now(),
28            rate_limiter: super::middleware::new_rate_limit_map(),
29            usage_tracker: UsageTracker::new(),
30        }
31    }
32}
33
34impl FromRef<AppState> for RateLimitMap {
35    fn from_ref(state: &AppState) -> Self {
36        state.rate_limiter.clone()
37    }
38}