soma_gateway/gateway/manager/
mcp_routes.rs

1use serde_json::Value;
2pub use soma_mcp_proxy::{
3    parse_upstream_resource_uri, upstream_resource_uri, McpPromptRoute as GatewayPromptRoute,
4    McpResourceRoute as GatewayResourceRoute, McpToolRoute as GatewayToolRoute,
5};
6
7use crate::upstream::{PromptDescriptor, ResourceDescriptor, ToolDescriptor, UpstreamHealth};
8
9use super::{GatewayManager, GatewayManagerError};
10
11impl GatewayManager {
12    pub async fn tool_routes(&self) -> Result<Vec<GatewayToolRoute>, GatewayManagerError> {
13        self.tool_routes_for_subject(None).await
14    }
15
16    pub async fn tool_routes_for_subject(
17        &self,
18        subject: Option<&str>,
19    ) -> Result<Vec<GatewayToolRoute>, GatewayManagerError> {
20        let pool = self.ready_pool()?;
21        let snapshots = discover_snapshots(&pool, subject).await?;
22        let mut candidates = Vec::new();
23        for snapshot in snapshots {
24            if !matches!(snapshot.health, UpstreamHealth::Connected) {
25                continue;
26            }
27            for descriptor in exposed_tools(&pool, &snapshot.name, subject).await? {
28                candidates.push((snapshot.name.clone(), descriptor));
29            }
30        }
31        Ok(soma_mcp_proxy::tool_routes_from_candidates(
32            candidates,
33            std::iter::empty::<&str>(),
34        ))
35    }
36
37    pub async fn call_mcp_tool(
38        &self,
39        name: &str,
40        params: Value,
41    ) -> Result<Option<Value>, GatewayManagerError> {
42        self.call_mcp_tool_for_subject(name, params, None).await
43    }
44
45    pub async fn call_mcp_tool_for_subject(
46        &self,
47        name: &str,
48        params: Value,
49        subject: Option<&str>,
50    ) -> Result<Option<Value>, GatewayManagerError> {
51        let Some(route) = self
52            .tool_routes_for_subject(subject)
53            .await?
54            .into_iter()
55            .find(|route| route.name == name)
56        else {
57            return Ok(None);
58        };
59        let pool = self.ready_pool()?;
60        call_tool(&pool, route, params, subject)
61            .await
62            .map(Some)
63            .map_err(Into::into)
64    }
65
66    pub async fn resource_routes(&self) -> Result<Vec<GatewayResourceRoute>, GatewayManagerError> {
67        self.resource_routes_for_subject(None).await
68    }
69
70    pub async fn resource_routes_for_subject(
71        &self,
72        subject: Option<&str>,
73    ) -> Result<Vec<GatewayResourceRoute>, GatewayManagerError> {
74        let pool = self.ready_pool()?;
75        let snapshots = discover_snapshots(&pool, subject).await?;
76        let mut routes = Vec::new();
77        for snapshot in snapshots {
78            if !matches!(snapshot.health, UpstreamHealth::Connected) {
79                continue;
80            }
81            for descriptor in list_resources(&pool, &snapshot.name, subject).await? {
82                routes.push(soma_mcp_proxy::resource_route(&snapshot.name, descriptor));
83            }
84        }
85        Ok(routes)
86    }
87
88    pub async fn read_mcp_resource(&self, uri: &str) -> Result<Option<Value>, GatewayManagerError> {
89        self.read_mcp_resource_for_subject(uri, None).await
90    }
91
92    pub async fn read_mcp_resource_for_subject(
93        &self,
94        uri: &str,
95        subject: Option<&str>,
96    ) -> Result<Option<Value>, GatewayManagerError> {
97        let Some((upstream, native_uri)) = parse_upstream_resource_uri(uri) else {
98            return Ok(None);
99        };
100        let pool = self.ready_pool()?;
101        read_resource(&pool, &upstream, &native_uri, subject)
102            .await
103            .map(Some)
104            .map_err(Into::into)
105    }
106
107    pub async fn prompt_routes(&self) -> Result<Vec<GatewayPromptRoute>, GatewayManagerError> {
108        self.prompt_routes_for_subject(None).await
109    }
110
111    pub async fn prompt_routes_for_subject(
112        &self,
113        subject: Option<&str>,
114    ) -> Result<Vec<GatewayPromptRoute>, GatewayManagerError> {
115        let pool = self.ready_pool()?;
116        let snapshots = discover_snapshots(&pool, subject).await?;
117        let mut candidates = Vec::new();
118        for snapshot in snapshots {
119            if !matches!(snapshot.health, UpstreamHealth::Connected) {
120                continue;
121            }
122            for descriptor in list_prompts(&pool, &snapshot.name, subject).await? {
123                candidates.push((snapshot.name.clone(), descriptor));
124            }
125        }
126        Ok(soma_mcp_proxy::prompt_routes_from_candidates(candidates))
127    }
128
129    pub async fn get_mcp_prompt(
130        &self,
131        name: &str,
132        arguments: Option<serde_json::Map<String, Value>>,
133    ) -> Result<Option<Value>, GatewayManagerError> {
134        self.get_mcp_prompt_for_subject(name, arguments, None).await
135    }
136
137    pub async fn get_mcp_prompt_for_subject(
138        &self,
139        name: &str,
140        arguments: Option<serde_json::Map<String, Value>>,
141        subject: Option<&str>,
142    ) -> Result<Option<Value>, GatewayManagerError> {
143        let Some(route) = self
144            .prompt_routes_for_subject(subject)
145            .await?
146            .into_iter()
147            .find(|route| route.name == name)
148        else {
149            return Ok(None);
150        };
151        let pool = self.ready_pool()?;
152        get_prompt(
153            &pool,
154            &route.upstream,
155            &route.native_name,
156            arguments,
157            subject,
158        )
159        .await
160        .map(Some)
161        .map_err(Into::into)
162    }
163
164    fn ready_pool(
165        &self,
166    ) -> Result<std::sync::Arc<crate::upstream::pool::UpstreamPool>, GatewayManagerError> {
167        self.ensure_ready()?;
168        Ok(self.pool.read().expect("gateway pool poisoned").clone())
169    }
170}
171
172async fn discover_snapshots(
173    pool: &crate::upstream::pool::UpstreamPool,
174    subject: Option<&str>,
175) -> Result<Vec<crate::upstream::UpstreamSnapshot>, crate::upstream::UpstreamError> {
176    let _ = subject;
177    #[cfg(feature = "oauth")]
178    if subject.is_some() {
179        return pool.discover_for_subject(subject).await;
180    }
181    pool.discover().await
182}
183
184async fn exposed_tools(
185    pool: &crate::upstream::pool::UpstreamPool,
186    upstream: &str,
187    subject: Option<&str>,
188) -> Result<Vec<ToolDescriptor>, crate::upstream::UpstreamError> {
189    let _ = subject;
190    #[cfg(feature = "oauth")]
191    if subject.is_some() {
192        return pool.exposed_tools_for_subject(upstream, subject).await;
193    }
194    pool.exposed_tools(upstream)
195}
196
197async fn call_tool(
198    pool: &crate::upstream::pool::UpstreamPool,
199    route: GatewayToolRoute,
200    params: Value,
201    subject: Option<&str>,
202) -> Result<Value, crate::upstream::UpstreamError> {
203    let _ = subject;
204    let call = crate::upstream::pool::ToolCall {
205        upstream: route.upstream,
206        tool: route.native_name,
207        params,
208    };
209    #[cfg(feature = "oauth")]
210    if subject.is_some() {
211        return pool.call_tool_for_subject(call, subject).await;
212    }
213    pool.call_tool(call).await
214}
215
216async fn list_resources(
217    pool: &crate::upstream::pool::UpstreamPool,
218    upstream: &str,
219    subject: Option<&str>,
220) -> Result<Vec<ResourceDescriptor>, crate::upstream::UpstreamError> {
221    let _ = subject;
222    #[cfg(feature = "oauth")]
223    if subject.is_some() {
224        return pool.list_resources_for_subject(upstream, subject).await;
225    }
226    pool.list_resources(upstream).await
227}
228
229async fn read_resource(
230    pool: &crate::upstream::pool::UpstreamPool,
231    upstream: &str,
232    uri: &str,
233    subject: Option<&str>,
234) -> Result<Value, crate::upstream::UpstreamError> {
235    let _ = subject;
236    #[cfg(feature = "oauth")]
237    if subject.is_some() {
238        return pool.read_resource_for_subject(upstream, uri, subject).await;
239    }
240    pool.read_resource(upstream, uri).await
241}
242
243async fn list_prompts(
244    pool: &crate::upstream::pool::UpstreamPool,
245    upstream: &str,
246    subject: Option<&str>,
247) -> Result<Vec<PromptDescriptor>, crate::upstream::UpstreamError> {
248    let _ = subject;
249    #[cfg(feature = "oauth")]
250    if subject.is_some() {
251        return pool.list_prompts_for_subject(upstream, subject).await;
252    }
253    pool.list_prompts(upstream).await
254}
255
256async fn get_prompt(
257    pool: &crate::upstream::pool::UpstreamPool,
258    upstream: &str,
259    name: &str,
260    arguments: Option<serde_json::Map<String, Value>>,
261    subject: Option<&str>,
262) -> Result<Value, crate::upstream::UpstreamError> {
263    let _ = subject;
264    #[cfg(feature = "oauth")]
265    if subject.is_some() {
266        return pool
267            .get_prompt_for_subject(upstream, name, arguments, subject)
268            .await;
269    }
270    pool.get_prompt(upstream, name, arguments).await
271}
272
273#[cfg(test)]
274#[path = "mcp_routes_tests.rs"]
275mod tests;
soma_gateway/gateway/manager/mcp_routes.rs

soma_gateway/gateway/manager/
mcp_routes.rs
