1use serde_json::Value;
2pub use soma_mcp_proxy::{
3 McpPromptRoute as GatewayPromptRoute, McpResourceRoute as GatewayResourceRoute,
4 McpToolRoute as GatewayToolRoute, parse_upstream_resource_uri, upstream_resource_uri,
5};
6
7use crate::upstream::{
8 McpRequestOutcome, McpRoundTrip, PromptDescriptor, ResourceDescriptor, ToolDescriptor,
9 UpstreamHealth,
10};
11
12use super::{GatewayManager, GatewayManagerError};
13
14impl GatewayManager {
15 pub async fn tool_routes(&self) -> Result<Vec<GatewayToolRoute>, GatewayManagerError> {
16 self.tool_routes_for_subject(None).await
17 }
18
19 pub async fn tool_routes_for_subject(
20 &self,
21 subject: Option<&str>,
22 ) -> Result<Vec<GatewayToolRoute>, GatewayManagerError> {
23 let pool = self.ready_pool()?;
24 let snapshots = discover_snapshots(&pool, subject).await?;
25 let mut candidates = Vec::new();
26 for snapshot in snapshots {
27 if !matches!(snapshot.health, UpstreamHealth::Connected) {
28 continue;
29 }
30 for descriptor in exposed_tools(&pool, &snapshot.name, subject).await? {
31 candidates.push((snapshot.name.clone(), descriptor));
32 }
33 }
34 Ok(soma_mcp_proxy::tool_routes_from_candidates(
35 candidates,
36 std::iter::empty::<&str>(),
37 ))
38 }
39
40 pub async fn call_mcp_tool(
41 &self,
42 name: &str,
43 params: Value,
44 ) -> Result<Option<Value>, GatewayManagerError> {
45 self.call_mcp_tool_for_subject(name, params, None).await
46 }
47
48 pub async fn call_mcp_tool_for_subject(
49 &self,
50 name: &str,
51 params: Value,
52 subject: Option<&str>,
53 ) -> Result<Option<Value>, GatewayManagerError> {
54 let Some(route) = self
55 .tool_routes_for_subject(subject)
56 .await?
57 .into_iter()
58 .find(|route| route.name == name)
59 else {
60 return Ok(None);
61 };
62 let pool = self.ready_pool()?;
63 call_tool(&pool, route, params, subject)
64 .await
65 .map(Some)
66 .map_err(Into::into)
67 }
68
69 pub async fn call_mcp_tool_once_for_subject(
70 &self,
71 name: &str,
72 params: Value,
73 round_trip: McpRoundTrip,
74 subject: Option<&str>,
75 ) -> Result<Option<McpRequestOutcome>, GatewayManagerError> {
76 let Some(route) = self
77 .tool_routes_for_subject(subject)
78 .await?
79 .into_iter()
80 .find(|route| route.name == name)
81 else {
82 return Ok(None);
83 };
84 let upstream = route.upstream.clone();
85 let pool = self.ready_pool()?;
86 let outcome = call_tool_once(&pool, route, params, round_trip, subject).await?;
87 Ok(Some(
88 self.register_task_outcome(outcome, &upstream, subject)?,
89 ))
90 }
91
92 pub async fn resource_routes(&self) -> Result<Vec<GatewayResourceRoute>, GatewayManagerError> {
93 self.resource_routes_for_subject(None).await
94 }
95
96 pub async fn resource_routes_for_subject(
97 &self,
98 subject: Option<&str>,
99 ) -> Result<Vec<GatewayResourceRoute>, GatewayManagerError> {
100 let pool = self.ready_pool()?;
101 let snapshots = discover_snapshots(&pool, subject).await?;
102 let mut routes = Vec::new();
103 for snapshot in snapshots {
104 if !matches!(snapshot.health, UpstreamHealth::Connected) {
105 continue;
106 }
107 for descriptor in list_resources(&pool, &snapshot.name, subject).await? {
108 routes.push(soma_mcp_proxy::resource_route(&snapshot.name, descriptor));
109 }
110 }
111 Ok(routes)
112 }
113
114 pub async fn read_mcp_resource(&self, uri: &str) -> Result<Option<Value>, GatewayManagerError> {
115 self.read_mcp_resource_for_subject(uri, None).await
116 }
117
118 pub async fn read_mcp_resource_for_subject(
119 &self,
120 uri: &str,
121 subject: Option<&str>,
122 ) -> Result<Option<Value>, GatewayManagerError> {
123 let Some((upstream, native_uri)) = parse_upstream_resource_uri(uri) else {
124 return Ok(None);
125 };
126 let pool = self.ready_pool()?;
127 read_resource(&pool, &upstream, &native_uri, subject)
128 .await
129 .map(Some)
130 .map_err(Into::into)
131 }
132
133 pub async fn read_mcp_resource_once_for_subject(
134 &self,
135 uri: &str,
136 round_trip: McpRoundTrip,
137 subject: Option<&str>,
138 ) -> Result<Option<McpRequestOutcome>, GatewayManagerError> {
139 let Some((upstream, native_uri)) = parse_upstream_resource_uri(uri) else {
140 return Ok(None);
141 };
142 let pool = self.ready_pool()?;
143 read_resource_once(&pool, &upstream, &native_uri, round_trip, subject)
144 .await
145 .map(Some)
146 .map_err(Into::into)
147 }
148
149 pub async fn prompt_routes(&self) -> Result<Vec<GatewayPromptRoute>, GatewayManagerError> {
150 self.prompt_routes_for_subject(None).await
151 }
152
153 pub async fn prompt_routes_for_subject(
154 &self,
155 subject: Option<&str>,
156 ) -> Result<Vec<GatewayPromptRoute>, GatewayManagerError> {
157 let pool = self.ready_pool()?;
158 let snapshots = discover_snapshots(&pool, subject).await?;
159 let mut candidates = Vec::new();
160 for snapshot in snapshots {
161 if !matches!(snapshot.health, UpstreamHealth::Connected) {
162 continue;
163 }
164 for descriptor in list_prompts(&pool, &snapshot.name, subject).await? {
165 candidates.push((snapshot.name.clone(), descriptor));
166 }
167 }
168 Ok(soma_mcp_proxy::prompt_routes_from_candidates(candidates))
169 }
170
171 pub async fn get_mcp_prompt(
172 &self,
173 name: &str,
174 arguments: Option<serde_json::Map<String, Value>>,
175 ) -> Result<Option<Value>, GatewayManagerError> {
176 self.get_mcp_prompt_for_subject(name, arguments, None).await
177 }
178
179 pub async fn get_mcp_prompt_for_subject(
180 &self,
181 name: &str,
182 arguments: Option<serde_json::Map<String, Value>>,
183 subject: Option<&str>,
184 ) -> Result<Option<Value>, GatewayManagerError> {
185 let Some(route) = self
186 .prompt_routes_for_subject(subject)
187 .await?
188 .into_iter()
189 .find(|route| route.name == name)
190 else {
191 return Ok(None);
192 };
193 let pool = self.ready_pool()?;
194 get_prompt(
195 &pool,
196 &route.upstream,
197 &route.native_name,
198 arguments,
199 subject,
200 )
201 .await
202 .map(Some)
203 .map_err(Into::into)
204 }
205
206 pub async fn get_mcp_prompt_once_for_subject(
207 &self,
208 name: &str,
209 arguments: Option<serde_json::Map<String, Value>>,
210 round_trip: McpRoundTrip,
211 subject: Option<&str>,
212 ) -> Result<Option<McpRequestOutcome>, GatewayManagerError> {
213 let Some(route) = self
214 .prompt_routes_for_subject(subject)
215 .await?
216 .into_iter()
217 .find(|route| route.name == name)
218 else {
219 return Ok(None);
220 };
221 let pool = self.ready_pool()?;
222 get_prompt_once(
223 &pool,
224 &route.upstream,
225 &route.native_name,
226 arguments,
227 round_trip,
228 subject,
229 )
230 .await
231 .map(Some)
232 .map_err(Into::into)
233 }
234
235 fn ready_pool(
236 &self,
237 ) -> Result<std::sync::Arc<crate::upstream::pool::UpstreamPool>, GatewayManagerError> {
238 self.ensure_ready()?;
239 Ok(self.pool.read().expect("gateway pool poisoned").clone())
240 }
241}
242
243async fn discover_snapshots(
244 pool: &crate::upstream::pool::UpstreamPool,
245 subject: Option<&str>,
246) -> Result<Vec<crate::upstream::UpstreamSnapshot>, crate::upstream::UpstreamError> {
247 let _ = subject;
248 #[cfg(feature = "oauth")]
249 if subject.is_some() {
250 return pool.discover_for_subject(subject).await;
251 }
252 pool.discover().await
253}
254
255async fn exposed_tools(
256 pool: &crate::upstream::pool::UpstreamPool,
257 upstream: &str,
258 subject: Option<&str>,
259) -> Result<Vec<ToolDescriptor>, crate::upstream::UpstreamError> {
260 let _ = subject;
261 #[cfg(feature = "oauth")]
262 if subject.is_some() {
263 return pool.exposed_tools_for_subject(upstream, subject).await;
264 }
265 pool.exposed_tools(upstream)
266}
267
268async fn call_tool(
269 pool: &crate::upstream::pool::UpstreamPool,
270 route: GatewayToolRoute,
271 params: Value,
272 subject: Option<&str>,
273) -> Result<Value, crate::upstream::UpstreamError> {
274 let _ = subject;
275 let call = crate::upstream::pool::ToolCall {
276 upstream: route.upstream,
277 tool: route.native_name,
278 params,
279 };
280 #[cfg(feature = "oauth")]
281 if subject.is_some() {
282 return pool.call_tool_for_subject(call, subject).await;
283 }
284 pool.call_tool(call).await
285}
286
287async fn call_tool_once(
288 pool: &crate::upstream::pool::UpstreamPool,
289 route: GatewayToolRoute,
290 params: Value,
291 round_trip: McpRoundTrip,
292 subject: Option<&str>,
293) -> Result<McpRequestOutcome, crate::upstream::UpstreamError> {
294 let _ = subject;
295 let call = crate::upstream::pool::ToolCall {
296 upstream: route.upstream,
297 tool: route.native_name,
298 params,
299 };
300 #[cfg(feature = "oauth")]
301 if subject.is_some() {
302 return pool
303 .call_tool_once_for_subject(call, round_trip, subject)
304 .await;
305 }
306 pool.call_tool_once(call, round_trip).await
307}
308
309async fn list_resources(
310 pool: &crate::upstream::pool::UpstreamPool,
311 upstream: &str,
312 subject: Option<&str>,
313) -> Result<Vec<ResourceDescriptor>, crate::upstream::UpstreamError> {
314 let _ = subject;
315 #[cfg(feature = "oauth")]
316 if subject.is_some() {
317 return pool.list_resources_for_subject(upstream, subject).await;
318 }
319 pool.list_resources(upstream).await
320}
321
322async fn read_resource(
323 pool: &crate::upstream::pool::UpstreamPool,
324 upstream: &str,
325 uri: &str,
326 subject: Option<&str>,
327) -> Result<Value, crate::upstream::UpstreamError> {
328 let _ = subject;
329 #[cfg(feature = "oauth")]
330 if subject.is_some() {
331 return pool.read_resource_for_subject(upstream, uri, subject).await;
332 }
333 pool.read_resource(upstream, uri).await
334}
335
336async fn read_resource_once(
337 pool: &crate::upstream::pool::UpstreamPool,
338 upstream: &str,
339 uri: &str,
340 round_trip: McpRoundTrip,
341 subject: Option<&str>,
342) -> Result<McpRequestOutcome, crate::upstream::UpstreamError> {
343 let _ = subject;
344 #[cfg(feature = "oauth")]
345 if subject.is_some() {
346 return pool
347 .read_resource_once_for_subject(upstream, uri, round_trip, subject)
348 .await;
349 }
350 pool.read_resource_once(upstream, uri, round_trip).await
351}
352
353async fn list_prompts(
354 pool: &crate::upstream::pool::UpstreamPool,
355 upstream: &str,
356 subject: Option<&str>,
357) -> Result<Vec<PromptDescriptor>, crate::upstream::UpstreamError> {
358 let _ = subject;
359 #[cfg(feature = "oauth")]
360 if subject.is_some() {
361 return pool.list_prompts_for_subject(upstream, subject).await;
362 }
363 pool.list_prompts(upstream).await
364}
365
366async fn get_prompt(
367 pool: &crate::upstream::pool::UpstreamPool,
368 upstream: &str,
369 name: &str,
370 arguments: Option<serde_json::Map<String, Value>>,
371 subject: Option<&str>,
372) -> Result<Value, crate::upstream::UpstreamError> {
373 let _ = subject;
374 #[cfg(feature = "oauth")]
375 if subject.is_some() {
376 return pool
377 .get_prompt_for_subject(upstream, name, arguments, subject)
378 .await;
379 }
380 pool.get_prompt(upstream, name, arguments).await
381}
382
383async fn get_prompt_once(
384 pool: &crate::upstream::pool::UpstreamPool,
385 upstream: &str,
386 name: &str,
387 arguments: Option<serde_json::Map<String, Value>>,
388 round_trip: McpRoundTrip,
389 subject: Option<&str>,
390) -> Result<McpRequestOutcome, crate::upstream::UpstreamError> {
391 let _ = subject;
392 #[cfg(feature = "oauth")]
393 if subject.is_some() {
394 return pool
395 .get_prompt_once_for_subject(upstream, name, arguments, round_trip, subject)
396 .await;
397 }
398 pool.get_prompt_once(upstream, name, arguments, round_trip)
399 .await
400}
401
402#[cfg(test)]
403#[path = "mcp_routes_tests.rs"]
404mod tests;