Skip to main content

litellm_rs/server/
config_sync.rs

1//! Authoritative encrypted runtime snapshots. Redis carries only revision IDs.
2use super::{runtime::build_runtime_revision, state::AppState};
3use crate::config::{
4    Config,
5    models::{cache::CacheConfig, provider::ProviderConfig, router::GatewayRouterConfig},
6};
7use crate::core::guardrails::GuardrailConfig;
8use crate::utils::{
9    auth::crypto::encryption::{decrypt_data, encrypt_data},
10    error::gateway_error::{GatewayError, Result},
11};
12use serde::{Deserialize, Serialize};
13use std::{collections::HashMap, sync::Arc, time::Duration};
14
15pub(crate) const REVISION_CHANNEL: &str = "litellm-rs:config:revisions";
16const RECONCILE_INTERVAL: Duration = Duration::from_secs(1);
17
18/// Public diagnostics contain no configuration values or backend error bodies.
19#[derive(Clone, Default, Serialize)]
20pub struct ConfigSyncStatus {
21    /// Revision currently serving requests on this node.
22    pub active_revision: u64,
23    /// Most recent authoritative revision read by this node.
24    pub observed_revision: u64,
25    /// Failed apply stage, kept until an authoritative snapshot applies successfully.
26    pub last_apply_error: Option<String>,
27    /// Failed storage/notification stage, cleared on successful synchronization.
28    pub last_sync_error: Option<String>,
29}
30
31pub(crate) struct ConfigSync {
32    key: Vec<u8>,
33    pub(crate) status: parking_lot::RwLock<ConfigSyncStatus>,
34}
35
36// Node-local listener/auth/storage/telemetry settings are deliberately excluded.
37#[derive(Serialize, Deserialize)]
38#[serde(deny_unknown_fields)]
39struct StoredRuntime {
40    revision: u64,
41    providers: Vec<ProviderConfig>,
42    router: GatewayRouterConfig,
43    model_aliases: HashMap<String, String>,
44    guardrails: GuardrailConfig,
45    cache: CacheConfig,
46}
47
48impl ConfigSync {
49    pub(crate) fn from_config(config: &Config) -> Result<Option<Arc<Self>>> {
50        let Some(env_name) = &config.gateway.storage.config_sync_key_env else {
51            return Ok(None);
52        };
53        let key = std::env::var(env_name).map_err(|_| {
54            GatewayError::Config(
55                "Configuration sync encryption key environment variable is missing".into(),
56            )
57        })?;
58        if key.len() < 32 {
59            return Err(GatewayError::Config(
60                "Configuration sync requires a dedicated encryption key of at least 32 bytes"
61                    .into(),
62            ));
63        }
64        Ok(Some(Arc::new(Self {
65            key: key.into_bytes(),
66            status: Default::default(),
67        })))
68    }
69
70    pub(crate) fn encrypt(&self, config: &Config, revision: u64) -> Result<String> {
71        let payload = StoredRuntime {
72            revision,
73            providers: config.gateway.providers.clone(),
74            router: config.gateway.router.clone(),
75            model_aliases: config.gateway.model_aliases.clone(),
76            guardrails: config.gateway.guardrails.clone(),
77            cache: config.gateway.cache.clone(),
78        };
79        let json = serde_json::to_string(&payload)
80            .map_err(|_| GatewayError::Config("Runtime snapshot serialization failed".into()))?;
81        encrypt_data(&self.key, &json)
82    }
83
84    fn candidate(&self, base: &Config, revision: u64, ciphertext: &str) -> Result<Config> {
85        let json = decrypt_data(&self.key, ciphertext).map_err(|_| {
86            GatewayError::Config(
87                "Runtime snapshot decryption failed; check the shared encryption key".into(),
88            )
89        })?;
90        let payload: StoredRuntime = serde_json::from_str(&json)
91            .map_err(|_| GatewayError::Config("Runtime snapshot decoding failed".into()))?;
92        if payload.revision != revision {
93            return Err(GatewayError::Config(
94                "Runtime snapshot revision does not match its authenticated payload".into(),
95            ));
96        }
97        let mut candidate = base.clone();
98        candidate.gateway.providers = payload.providers;
99        candidate.gateway.router = payload.router;
100        candidate.gateway.model_aliases = payload.model_aliases;
101        candidate.gateway.guardrails = payload.guardrails;
102        candidate.gateway.cache = payload.cache;
103        Ok(candidate)
104    }
105}
106
107impl AppState {
108    /// Current node's revision and sanitized synchronization diagnostics.
109    pub fn config_sync_status(&self) -> ConfigSyncStatus {
110        let mut status = self
111            .config_sync
112            .as_ref()
113            .map(|sync| sync.status.read().clone())
114            .unwrap_or_default();
115        status.active_revision = self.pin_runtime().generation;
116        status
117    }
118
119    pub(crate) async fn reconcile_runtime_config(&self) -> Result<()> {
120        let Some(sync) = &self.config_sync else {
121            return Ok(());
122        };
123        let _guard = self.apply_lock.lock().await;
124        let row = self.storage.database.runtime_config().await.map_err(|_| {
125            sync.status.write().last_sync_error =
126                Some("Authoritative configuration read failed".into());
127            GatewayError::Storage("Authoritative configuration read failed".into())
128        })?;
129        let revision = u64::try_from(row.revision).map_err(|_| {
130            GatewayError::Storage("Invalid authoritative configuration revision".into())
131        })?;
132        sync.status.write().observed_revision = revision;
133        if revision <= self.pin_runtime().generation {
134            return Ok(());
135        }
136        let apply = async {
137            let ciphertext = row.encrypted_payload.as_deref().ok_or_else(|| {
138                GatewayError::Storage("Authoritative configuration payload is missing".into())
139            })?;
140            let candidate = sync.candidate(&self.config(), revision, ciphertext)?;
141            let built = build_runtime_revision(
142                candidate,
143                revision,
144                self.pricing.clone(),
145                self.storage.redis.clone(),
146            )
147            .await
148            .map_err(|_| {
149                GatewayError::Config("Authoritative runtime validation/build failed".into())
150            })?;
151            self.publish_runtime(built);
152            Ok::<_, GatewayError>(())
153        }
154        .await;
155        let mut status = sync.status.write();
156        match apply {
157            Ok(()) => {
158                status.last_apply_error = None;
159                Ok(())
160            }
161            Err(error) => {
162                status.last_apply_error = Some(error.to_string());
163                Err(error)
164            }
165        }
166    }
167
168    async fn observe_revision(&self, message: redis::Msg) -> Result<()> {
169        let revision = message.get_payload::<u64>().map_err(|_| {
170            GatewayError::Storage("Invalid configuration revision notification".into())
171        })?;
172        if revision > self.pin_runtime().generation {
173            self.reconcile_runtime_config().await?;
174        }
175        Ok(())
176    }
177}
178
179/// Owns the subscriber so dropping/stopping a server also stops its worker.
180pub(super) struct ConfigSyncTask(tokio::task::JoinHandle<()>);
181impl Drop for ConfigSyncTask {
182    fn drop(&mut self) {
183        self.0.abort();
184    }
185}
186
187pub(super) fn start(state: AppState) -> Option<ConfigSyncTask> {
188    let sync = state.config_sync.clone()?;
189    Some(ConfigSyncTask(tokio::spawn(async move {
190        loop {
191            let result: Result<()> = async {
192                let mut subscription = state
193                    .storage
194                    .redis
195                    .subscribe(&[REVISION_CHANNEL.to_string()])
196                    .await?;
197                // Subscribe before loading, closing the reconnect read/subscribe race.
198                state.reconcile_runtime_config().await?;
199                loop {
200                    match tokio::time::timeout(RECONCILE_INTERVAL, subscription.next_message())
201                        .await
202                    {
203                        Ok(Ok(message)) => state.observe_revision(message).await?,
204                        Ok(Err(error)) => return Err(error),
205                        Err(_) => state.reconcile_runtime_config().await?,
206                    }
207                    sync.status.write().last_sync_error = None;
208                }
209            }
210            .await;
211            if let Err(error) = result {
212                sync.status.write().last_sync_error = Some(error.to_string());
213                tracing::error!(error = %error, "Configuration synchronization failed");
214            }
215            tokio::time::sleep(RECONCILE_INTERVAL).await;
216        }
217    })))
218}
219
220#[cfg(all(test, feature = "sqlite"))]
221#[path = "config_sync_tests.rs"]
222mod tests;