litellm_rs/server/
state.rs

1//! Application state shared across HTTP handlers
2//!
3//! This module provides the AppState struct and its implementations.
4
5use crate::config::Config;
6use crate::core::audit::AuditLogger;
7use crate::core::budget::{BudgetManager, UnifiedBudgetLimits};
8use crate::core::cache::LLMCache;
9use crate::core::guardrails::GuardrailEngine;
10use crate::core::ip_access::IpAccessControl;
11use crate::core::keys::{DatabaseKeyRepository, KeyManager};
12use crate::core::observability::RuntimeObservability;
13use crate::core::pricing_service::PricingService;
14use crate::core::router::UnifiedRouter;
15use crate::core::teams::TeamManager;
16use crate::core::virtual_keys::RuntimeVirtualKeyManager;
17use crate::server::routes::ai::budgeted::BudgetedExecutor;
18use crate::storage::database::SeaOrmTeamRepository;
19use crate::utils::error::gateway_error::{GatewayError, Result};
20use crate::utils::sync::AtomicValue;
21use std::sync::Arc;
22use tokio::sync::Mutex;
23
24#[cfg(test)]
25#[path = "runtime_reload_tests.rs"]
26mod runtime_reload_tests;
27
28#[cfg(test)]
29#[path = "runtime_reload_profile.rs"]
30mod runtime_reload_profile;
31
32pub use super::runtime::RuntimeRevision;
33use super::runtime::{build_response_cache, build_runtime_revision};
34
35/// HTTP server state shared across handlers
36///
37/// This struct contains shared resources that need to be accessed across
38/// multiple request handlers. All fields are wrapped in Arc for efficient
39/// sharing across threads.
40///
41/// `config` uses [`AtomicValue`] so callers can explicitly swap the entire
42/// configuration at runtime while readers obtain lock-free `Arc<Config>`
43/// snapshots. This type does not start a file watcher by itself.
44///
45/// Router, guardrails, and response cache are published together through
46/// [`Self::apply_runtime`]. [`Self::pin_runtime`] holds one generation for a
47/// request lifecycle. Live accessors load from that same bundle so HTTP traffic
48/// never stays on construction-time Arcs after a successful apply.
49#[derive(Clone)]
50pub struct AppState {
51    /// Gateway configuration (atomically swappable by explicit callers)
52    pub config: Arc<AtomicValue<Config>>,
53    /// Authentication system
54    pub auth: Arc<crate::auth::AuthSystem>,
55    /// Storage layer
56    pub storage: Arc<crate::storage::StorageLayer>,
57    /// Unified pricing service
58    pub pricing: Arc<PricingService>,
59    /// Budget limits for provider and model cost tracking
60    pub budget_limits: Arc<UnifiedBudgetLimits>,
61    /// General budget manager for keyed budget scopes such as API key budgets
62    pub budget_manager: Arc<BudgetManager>,
63    /// Team manager for team lifecycle operations (shared, in-memory by default)
64    pub team_manager: Arc<TeamManager>,
65    /// API key manager for `/v1/keys` route handlers (shared across requests)
66    pub key_manager: RuntimeVirtualKeyManager,
67    /// Budget orchestration service for AI route reserve/call/settle lifecycles
68    pub(crate) budgeted: BudgetedExecutor,
69    /// Non-blocking external request lifecycle callback dispatcher
70    pub callbacks: RuntimeObservability,
71    /// Explicitly configured request audit logger.
72    pub audit_logger: Arc<AuditLogger>,
73    /// IP policy consumed by the outer HTTP middleware.
74    pub ip_access: Arc<IpAccessControl>,
75    #[cfg(feature = "mcp")]
76    pub(crate) mcp_inflight: Arc<super::routes::mcp::Inflight>,
77    pub(super) runtime: Arc<AtomicValue<RuntimeRevision>>,
78    #[cfg(feature = "a2a")]
79    pub(crate) a2a_tasks: Arc<crate::server::routes::a2a::TaskOwners>,
80    pub(super) apply_lock: Arc<Mutex<()>>,
81    pub(super) config_sync: Option<Arc<super::config_sync::ConfigSync>>,
82}
83
84impl AppState {
85    /// Create a new AppState with unified router
86    pub fn new_with_unified_router(
87        config: Config,
88        auth: crate::auth::AuthSystem,
89        unified_router: crate::core::router::UnifiedRouter,
90        storage: crate::storage::StorageLayer,
91        pricing: Arc<PricingService>,
92        budget_limits: Arc<UnifiedBudgetLimits>,
93    ) -> Self {
94        let redis = storage.redis.clone();
95        let config = Arc::new(config);
96        let revision = RuntimeRevision {
97            generation: 0,
98            unified_router: Arc::new(unified_router),
99            guardrails: Arc::new(GuardrailEngine::disabled()),
100            response_cache: build_response_cache(&config, redis),
101            config: Arc::clone(&config),
102        };
103        Self::new_with_runtime(revision, auth, storage, pricing, budget_limits)
104    }
105
106    /// Create AppState from a fully built runtime revision.
107    pub fn new_with_runtime(
108        revision: RuntimeRevision,
109        auth: crate::auth::AuthSystem,
110        storage: crate::storage::StorageLayer,
111        pricing: Arc<PricingService>,
112        budget_limits: Arc<UnifiedBudgetLimits>,
113    ) -> Self {
114        let storage = Arc::new(storage);
115        let key_manager = KeyManager::new(DatabaseKeyRepository::new(storage.clone()))
116            .with_hmac_secret(revision.config.gateway.auth.api_key_hmac_secret.clone());
117        let budget_manager = Arc::new(BudgetManager::new());
118        let budgeted = BudgetedExecutor::new(
119            budget_limits.clone(),
120            budget_manager.clone(),
121            pricing.clone(),
122            key_manager.clone(),
123        );
124        let team_manager = Arc::new(TeamManager::new(Arc::new(SeaOrmTeamRepository::new(
125            storage.database.clone(),
126        ))));
127        Self {
128            config: Arc::new(AtomicValue::from(Arc::clone(&revision.config))),
129            auth: Arc::new(auth),
130            storage,
131            pricing,
132            budget_limits,
133            budget_manager,
134            team_manager,
135            key_manager,
136            budgeted,
137            callbacks: RuntimeObservability::disabled(),
138            audit_logger: Arc::new(AuditLogger::disabled()),
139            ip_access: Arc::new(IpAccessControl::disabled()),
140            #[cfg(feature = "mcp")]
141            mcp_inflight: Arc::new(super::routes::mcp::Inflight::default()),
142            runtime: Arc::new(AtomicValue::new(revision)),
143            #[cfg(feature = "a2a")]
144            a2a_tasks: Arc::default(),
145            apply_lock: Arc::new(Mutex::new(())),
146            config_sync: None,
147        }
148    }
149
150    /// Attach a configured callback dispatcher.
151    pub fn with_callbacks(mut self, callbacks: RuntimeObservability) -> Self {
152        self.callbacks = callbacks;
153        self
154    }
155
156    /// Attach the request audit logger built during startup.
157    pub fn with_audit_logger(mut self, audit_logger: Arc<AuditLogger>) -> Self {
158        self.audit_logger = audit_logger;
159        self
160    }
161
162    /// Attach validated content and network policy engines.
163    pub fn with_request_policies(
164        self,
165        guardrails: Arc<GuardrailEngine>,
166        ip_access: Arc<IpAccessControl>,
167    ) -> Self {
168        let current = self.runtime.load();
169        self.runtime.store(RuntimeRevision {
170            generation: current.generation,
171            config: Arc::clone(&current.config),
172            unified_router: Arc::clone(&current.unified_router),
173            guardrails,
174            response_cache: current.response_cache.clone(),
175        });
176        let mut this = self;
177        this.ip_access = ip_access;
178        this
179    }
180
181    /// Lock-free pin of the active runtime generation.
182    ///
183    /// Clone the returned `Arc` (or clone this pin) to keep router, guardrails,
184    /// cache, config, and generation stable for the rest of a request even if
185    /// [`Self::apply_runtime`] publishes a newer revision.
186    pub fn pin_runtime(&self) -> Arc<RuntimeRevision> {
187        self.runtime.load()
188    }
189
190    /// Live router for the current generation.
191    pub fn unified_router(&self) -> Arc<UnifiedRouter> {
192        Arc::clone(&self.pin_runtime().unified_router)
193    }
194
195    /// Live guardrail engine for the current generation.
196    pub fn guardrails(&self) -> Arc<GuardrailEngine> {
197        Arc::clone(&self.pin_runtime().guardrails)
198    }
199
200    /// Live response cache for the current generation.
201    pub fn response_cache(&self) -> Option<Arc<LLMCache>> {
202        self.pin_runtime().response_cache.clone()
203    }
204
205    /// Validate and build a complete candidate revision, then publish it atomically.
206    ///
207    /// Any build failure leaves the previous revision fully active. On success,
208    /// router, provider health, guardrails, response cache, generation, and
209    /// [`Self::config`] observe one consistent generation. Pricing is reused
210    /// from [`Self::pricing`] rather than rebuilt.
211    pub async fn apply_runtime(&self, candidate: Config) -> Result<u64> {
212        let generation = self
213            .apply_runtime_at(candidate, self.pin_runtime().generation)
214            .await?;
215        self.notify_runtime_revision(generation).await?;
216        Ok(generation)
217    }
218
219    pub(crate) async fn apply_runtime_at(&self, candidate: Config, expected: u64) -> Result<u64> {
220        let _guard = self.apply_lock.lock().await;
221        if self.pin_runtime().generation != expected {
222            return Err(GatewayError::Conflict(
223                "Runtime revision changed; reload configuration before retrying".into(),
224            ));
225        }
226        let generation = expected
227            .checked_add(1)
228            .ok_or_else(|| GatewayError::Config("runtime generation overflow".into()))?;
229        let revision = build_runtime_revision(
230            candidate,
231            generation,
232            Arc::clone(&self.pricing),
233            Arc::clone(&self.storage.redis),
234        )
235        .await
236        .inspect_err(|_| {
237            if let Some(sync) = &self.config_sync {
238                sync.status.write().last_apply_error =
239                    Some("Runtime candidate validation/build failed".into());
240            }
241        })?;
242        if let Some(sync) = &self.config_sync {
243            let ciphertext = sync.encrypt(&revision.config, generation)?;
244            self.storage
245                .database
246                .commit_runtime_config(expected, ciphertext)
247                .await
248                .inspect_err(|_| {
249                    sync.status.write().last_sync_error =
250                        Some("Authoritative configuration commit failed".into());
251                })?;
252            sync.status.write().observed_revision = generation;
253        }
254        self.publish_runtime(revision);
255        if let Some(sync) = &self.config_sync {
256            sync.status.write().last_apply_error = None;
257        }
258        Ok(generation)
259    }
260
261    /// Notify only after callers have recorded the committed mutation's audit data.
262    pub(crate) async fn notify_runtime_revision(&self, generation: u64) -> Result<()> {
263        if let Some(sync) = &self.config_sync {
264            let published = tokio::time::timeout(
265                std::time::Duration::from_secs(5),
266                self.storage.redis.publish(
267                    super::config_sync::REVISION_CHANNEL,
268                    &generation.to_string(),
269                ),
270            )
271            .await;
272            if !matches!(published, Ok(Ok(()))) {
273                sync.status.write().last_sync_error =
274                    Some("Revision committed but notification failed".into());
275                return Err(GatewayError::Storage(format!(
276                    "Runtime revision {generation} committed and applied, but notification failed; peers reconcile from the database"
277                )));
278            }
279            sync.status.write().last_sync_error = None;
280        }
281        Ok(())
282    }
283
284    pub(super) fn publish_runtime(&self, revision: RuntimeRevision) {
285        let previous = self.runtime.load();
286        revision
287            .unified_router
288            .inherit_runtime_state(&previous.unified_router);
289        // Stop cleanup tasks on the obsolete response cache before publishing.
290        // DualCache also shuts down on Drop, but in-flight Arc pins can keep the
291        // old revision alive; explicit shutdown avoids leaking barrier loops
292        // across every configuration reload.
293        if let Some(cache) = previous.response_cache.as_ref() {
294            cache.shutdown();
295        }
296        let config = Arc::clone(&revision.config);
297        self.runtime.store(revision);
298        self.config.store_arc(config);
299    }
300
301    /// Load a snapshot of the current gateway configuration.
302    ///
303    /// Returns an `Arc<Config>` that is valid for the lifetime of the
304    /// caller — subsequent explicit swaps will not affect already-loaded
305    /// snapshots.
306    pub fn config(&self) -> Arc<Config> {
307        self.config.load()
308    }
309}
310
311#[cfg(test)]
312mod tests {
313    use super::*;
314    use crate::config::models::provider::ProviderConfig;
315    use crate::core::guardrails::PromptInjectionConfig;
316    use crate::server::HttpServer;
317
318    async fn test_app_state() -> AppState {
319        let mut config = crate::server::valid_test_config();
320        config.gateway.storage.database.enabled = false;
321        config.gateway.storage.redis.enabled = false;
322        config.gateway.pricing.source = None;
323        match HttpServer::new(&config).await {
324            Ok(server) => server.state().clone(),
325            Err(error) => panic!("test AppState startup failed: {error}"),
326        }
327    }
328
329    fn extra_provider(name: &str) -> ProviderConfig {
330        ProviderConfig {
331            name: name.to_string(),
332            provider_type: "openai".to_string(),
333            api_key: "sk-test".to_string(),
334            ..ProviderConfig::default()
335        }
336    }
337
338    #[tokio::test]
339    async fn apply_runtime_successful_swap_publishes_one_generation() {
340        let state = test_app_state().await;
341        let pinned = state.pin_runtime();
342        assert_eq!(pinned.generation, 0);
343        assert!(pinned.response_cache.is_none());
344
345        let mut next = (*state.config()).clone();
346        next.gateway.cache.enabled = true;
347        next.gateway.guardrails.enabled = false;
348        next.gateway.providers.push(extra_provider("test-openai-2"));
349
350        let generation = state
351            .apply_runtime(next)
352            .await
353            .expect("valid candidate should apply");
354        assert_eq!(generation, 1);
355
356        let live = state.pin_runtime();
357        assert_eq!(live.generation, 1);
358        assert!(live.response_cache.is_some());
359        assert!(!live.guardrails.is_enabled());
360        assert!(
361            live.unified_router.list_deployments().len()
362                > pinned.unified_router.list_deployments().len()
363        );
364        assert!(Arc::ptr_eq(&live.config, &state.config.load()));
365        assert!(!Arc::ptr_eq(&live.unified_router, &pinned.unified_router));
366        assert!(!Arc::ptr_eq(&live.guardrails, &pinned.guardrails));
367
368        assert_eq!(pinned.generation, 0);
369        assert!(pinned.response_cache.is_none());
370        assert!(pinned.guardrails.is_enabled());
371    }
372
373    #[tokio::test]
374    async fn apply_runtime_failed_build_leaves_previous_revision_active() {
375        let state = test_app_state().await;
376        let before = state.pin_runtime();
377        let router = Arc::clone(&before.unified_router);
378
379        let mut invalid = (*state.config()).clone();
380        let mut injection = invalid
381            .gateway
382            .guardrails
383            .prompt_injection
384            .clone()
385            .unwrap_or_else(PromptInjectionConfig::new);
386        injection.enabled = true;
387        injection.custom_patterns.push("(".to_string());
388        invalid.gateway.guardrails.prompt_injection = Some(injection);
389
390        let error = state
391            .apply_runtime(invalid)
392            .await
393            .expect_err("invalid guardrails must fail apply");
394        match error {
395            GatewayError::Config(message) => {
396                assert!(
397                    message.to_lowercase().contains("guardrail")
398                        || message.to_lowercase().contains("pattern"),
399                    "unexpected apply error: {message}"
400                );
401            }
402            other => panic!("expected config error, got {other:?}"),
403        }
404
405        let after = state.pin_runtime();
406        assert_eq!(after.generation, before.generation);
407        assert!(Arc::ptr_eq(&after.unified_router, &router));
408        assert!(Arc::ptr_eq(&after.guardrails, &before.guardrails));
409        assert_eq!(after.generation, 0);
410    }
411
412    #[tokio::test]
413    async fn in_flight_pin_stays_on_old_generation_after_apply() {
414        let state = test_app_state().await;
415        let in_flight = Arc::clone(&state.pin_runtime());
416
417        let mut next = (*state.config()).clone();
418        next.gateway.cache.enabled = true;
419        next.gateway.guardrails.enabled = false;
420        next.gateway
421            .providers
422            .push(extra_provider("test-openai-inflight"));
423
424        state
425            .apply_runtime(next)
426            .await
427            .expect("valid candidate should apply");
428
429        let live = state.pin_runtime();
430        assert_eq!(live.generation, in_flight.generation + 1);
431        assert_eq!(in_flight.generation, 0);
432        assert!(!Arc::ptr_eq(
433            &in_flight.unified_router,
434            &live.unified_router
435        ));
436        assert!(!Arc::ptr_eq(&in_flight.guardrails, &live.guardrails));
437        assert!(in_flight.response_cache.is_none());
438        assert!(live.response_cache.is_some());
439    }
440}
litellm_rs/server/state.rs

litellm_rs/server/
state.rs
