soma_gateway/gateway/
manager.rs

1use std::sync::{Arc, RwLock};
2
3use serde_json::Value;
4use thiserror::Error;
5
6use crate::config::{ConfigError, GatewayConfig, GatewayConfigView, UpstreamConfig};
7use crate::gateway::config_store::FsGatewayConfigStore;
8use crate::upstream::pool::{ToolCall, UpstreamPool};
9use crate::upstream::{UpstreamError, UpstreamSnapshot};
10use crate::usage::{NoopUsageSink, UsageEvent, UsageSink};
11
12pub mod core;
13pub mod mcp_projection;
14pub mod mcp_routes;
15#[cfg(feature = "protected-routes")]
16pub mod mcp_scoped_routes;
17#[cfg(feature = "oauth")]
18pub mod oauth_lifecycle;
19pub mod pool_lifecycle;
20#[cfg(feature = "protected-routes")]
21pub mod protected_routes;
22#[cfg(feature = "protected-routes")]
23pub mod virtual_servers;
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub enum GatewayLifecycle {
27    Ready,
28    Reloading,
29}
30
31#[derive(Debug, Error)]
32pub enum GatewayManagerError {
33    #[error(transparent)]
34    Config(#[from] ConfigError),
35    #[error(transparent)]
36    Upstream(#[from] UpstreamError),
37    #[error("gateway_reloading")]
38    GatewayReloading,
39    #[error("gateway config store is not mounted")]
40    StoreNotMounted,
41    #[error("upstream `{0}` is already configured")]
42    UpstreamExists(String),
43    #[error("upstream `{0}` is not configured")]
44    UpstreamMissing(String),
45    #[error("gateway oauth runtime error: {0}")]
46    OAuth(String),
47}
48
49impl From<soma_mcp_client::ConfigError> for GatewayManagerError {
50    fn from(error: soma_mcp_client::ConfigError) -> Self {
51        Self::Config(error.into())
52    }
53}
54
55pub struct GatewayManager {
56    config: RwLock<GatewayConfig>,
57    pool: RwLock<Arc<UpstreamPool>>,
58    lifecycle: RwLock<GatewayLifecycle>,
59    usage: Arc<dyn UsageSink>,
60    store: Option<FsGatewayConfigStore>,
61    #[cfg(feature = "oauth")]
62    oauth_runtime: RwLock<Option<Arc<soma_mcp_client::oauth::UpstreamOAuthRuntime>>>,
63}
64
65impl GatewayManager {
66    pub fn new(config: GatewayConfig) -> Result<Self, GatewayManagerError> {
67        Self::with_usage(config, Arc::new(NoopUsageSink))
68    }
69
70    pub fn with_usage(
71        config: GatewayConfig,
72        usage: Arc<dyn UsageSink>,
73    ) -> Result<Self, GatewayManagerError> {
74        Self::build(config, usage, None)
75    }
76
77    pub fn from_store(store: FsGatewayConfigStore) -> Result<Self, GatewayManagerError> {
78        let config = store.load_or_install_default()?;
79        Self::build(config, Arc::new(NoopUsageSink), Some(store))
80    }
81
82    fn build(
83        config: GatewayConfig,
84        usage: Arc<dyn UsageSink>,
85        store: Option<FsGatewayConfigStore>,
86    ) -> Result<Self, GatewayManagerError> {
87        config.validate()?;
88        let pool = pool_lifecycle::build_pool_from_config(&config)?;
89        Ok(Self {
90            config: RwLock::new(config),
91            pool: RwLock::new(Arc::new(pool)),
92            lifecycle: RwLock::new(GatewayLifecycle::Ready),
93            usage,
94            store,
95            #[cfg(feature = "oauth")]
96            oauth_runtime: RwLock::new(None),
97        })
98    }
99
100    #[must_use]
101    pub fn lifecycle(&self) -> GatewayLifecycle {
102        *self.lifecycle.read().expect("gateway lifecycle poisoned")
103    }
104
105    #[must_use]
106    pub fn config_view(&self) -> GatewayConfigView {
107        self.config
108            .read()
109            .expect("gateway config poisoned")
110            .redacted_view()
111    }
112
113    pub async fn discover(&self) -> Result<Vec<UpstreamSnapshot>, GatewayManagerError> {
114        self.ensure_ready()?;
115        let pool = self.pool.read().expect("gateway pool poisoned").clone();
116        Ok(pool.discover().await?)
117    }
118
119    pub fn exposed_tool_count(&self) -> Result<usize, GatewayManagerError> {
120        self.ensure_ready()?;
121        Ok(self
122            .pool
123            .read()
124            .expect("gateway pool poisoned")
125            .exposed_tool_count())
126    }
127
128    pub async fn call_tool(
129        &self,
130        upstream: impl Into<String>,
131        tool: impl Into<String>,
132        params: Value,
133    ) -> Result<Value, GatewayManagerError> {
134        self.ensure_ready()?;
135        let upstream = upstream.into();
136        let tool = tool.into();
137        let pool = self.pool.read().expect("gateway pool poisoned").clone();
138        let result = pool
139            .call_tool(ToolCall {
140                upstream: upstream.clone(),
141                tool,
142                params,
143            })
144            .await;
145        let success = result.is_ok();
146        let bytes = result
147            .as_ref()
148            .ok()
149            .and_then(|value| serde_json::to_vec(value).ok())
150            .map_or(0, |bytes| bytes.len());
151        self.usage.record(UsageEvent {
152            action: "call_tool".to_owned(),
153            upstream: Some(upstream),
154            success,
155            bytes,
156        });
157        Ok(result?)
158    }
159
160    pub fn add_upstream(
161        &self,
162        upstream: UpstreamConfig,
163    ) -> Result<GatewayConfigView, GatewayManagerError> {
164        upstream.validate()?;
165        self.mutate_config(|config| {
166            if config
167                .upstream
168                .iter()
169                .any(|item| item.name == upstream.name)
170            {
171                return Err(GatewayManagerError::UpstreamExists(upstream.name.clone()));
172            }
173            config.upstream.push(upstream);
174            Ok(())
175        })
176    }
177
178    pub fn update_upstream(
179        &self,
180        upstream: UpstreamConfig,
181    ) -> Result<GatewayConfigView, GatewayManagerError> {
182        upstream.validate()?;
183        self.mutate_config(|config| {
184            let Some(slot) = config
185                .upstream
186                .iter_mut()
187                .find(|item| item.name == upstream.name)
188            else {
189                return Err(GatewayManagerError::UpstreamMissing(upstream.name.clone()));
190            };
191            *slot = upstream;
192            Ok(())
193        })
194    }
195
196    pub fn remove_upstream(&self, name: &str) -> Result<GatewayConfigView, GatewayManagerError> {
197        let name = name.trim();
198        if name.is_empty() {
199            return Err(GatewayManagerError::Config(ConfigError::invalid(
200                "name",
201                "must not be empty",
202            )));
203        }
204        self.mutate_config(|config| {
205            let before = config.upstream.len();
206            config.upstream.retain(|item| item.name != name);
207            if config.upstream.len() == before {
208                return Err(GatewayManagerError::UpstreamMissing(name.to_owned()));
209            }
210            Ok(())
211        })
212    }
213
214    pub fn reload_from_store(&self) -> Result<GatewayConfigView, GatewayManagerError> {
215        let Some(store) = &self.store else {
216            return Err(GatewayManagerError::StoreNotMounted);
217        };
218        self.replace_config(store.load()?)
219    }
220
221    fn ensure_ready(&self) -> Result<(), GatewayManagerError> {
222        if self.lifecycle() == GatewayLifecycle::Ready {
223            return Ok(());
224        }
225        Err(GatewayManagerError::GatewayReloading)
226    }
227
228    fn mutate_config(
229        &self,
230        mutate: impl FnOnce(&mut GatewayConfig) -> Result<(), GatewayManagerError>,
231    ) -> Result<GatewayConfigView, GatewayManagerError> {
232        let mut next = self.config.read().expect("gateway config poisoned").clone();
233        mutate(&mut next)?;
234        if let Some(store) = &self.store {
235            store.save(&next)?;
236        }
237        self.replace_config(next)
238    }
239
240    fn replace_config(
241        &self,
242        next: GatewayConfig,
243    ) -> Result<GatewayConfigView, GatewayManagerError> {
244        next.validate()?;
245        self.with_reloading(|| {
246            self.replace_config_and_pool(next)?;
247            Ok(self.config_view())
248        })
249    }
250
251    pub(super) fn reload_config(&self, next: GatewayConfig) -> Result<(), GatewayManagerError> {
252        next.validate()?;
253        self.with_reloading(|| self.replace_config_and_pool(next))
254    }
255
256    fn replace_config_and_pool(&self, next: GatewayConfig) -> Result<(), GatewayManagerError> {
257        let pool = pool_lifecycle::build_pool_from_config(&next)?;
258        #[cfg(feature = "oauth")]
259        if let Some(runtime) = self
260            .oauth_runtime
261            .read()
262            .expect("gateway oauth runtime poisoned")
263            .as_ref()
264        {
265            pool.install_oauth_provider(runtime.provider());
266        }
267        *self.config.write().expect("gateway config poisoned") = next;
268        *self.pool.write().expect("gateway pool poisoned") = Arc::new(pool);
269        Ok(())
270    }
271
272    #[cfg(feature = "oauth")]
273    pub fn install_upstream_oauth_runtime(
274        &self,
275        runtime: soma_mcp_client::oauth::UpstreamOAuthRuntime,
276    ) {
277        let runtime = Arc::new(runtime);
278        self.pool
279            .read()
280            .expect("gateway pool poisoned")
281            .install_oauth_provider(runtime.provider());
282        *self
283            .oauth_runtime
284            .write()
285            .expect("gateway oauth runtime poisoned") = Some(runtime);
286    }
287
288    fn with_reloading<T>(
289        &self,
290        operation: impl FnOnce() -> Result<T, GatewayManagerError>,
291    ) -> Result<T, GatewayManagerError> {
292        *self.lifecycle.write().expect("gateway lifecycle poisoned") = GatewayLifecycle::Reloading;
293        let result = operation();
294        *self.lifecycle.write().expect("gateway lifecycle poisoned") = GatewayLifecycle::Ready;
295        result
296    }
297
298    #[cfg(test)]
299    pub(crate) fn install_pool_for_tests(&self, pool: UpstreamPool) {
300        *self.pool.write().expect("gateway pool poisoned") = Arc::new(pool);
301    }
302
303    #[cfg(test)]
304    pub(crate) fn set_lifecycle_for_tests(&self, lifecycle: GatewayLifecycle) {
305        *self.lifecycle.write().expect("gateway lifecycle poisoned") = lifecycle;
306    }
307}
308
309#[cfg(test)]
310#[path = "manager_tests.rs"]
311mod tests;
soma_gateway/gateway/manager.rs

soma_gateway/gateway/
manager.rs
