soma_gateway/gateway/
manager.rs

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

soma_gateway/gateway/
manager.rs
