soma_mcp_client/upstream/
pool.rs

1use std::collections::BTreeMap;
2use std::sync::{Arc, RwLock};
3
4use serde_json::Value;
5
6use crate::config::UpstreamConfig;
7use crate::process::guard::SpawnGuard;
8use crate::upstream::http_client::{decide_http_transport, transport_kind_for_decision};
9use crate::upstream::{
10    McpRequestOutcome, McpRoundTrip, ResponseCaps, ToolDescriptor, TransportKind, UpstreamError,
11    UpstreamHealth, UpstreamSnapshot,
12};
13
14pub mod connect_stdio;
15pub mod discovery;
16pub mod health;
17mod lifecycle_compat;
18pub mod live;
19pub mod prompts;
20pub mod resources;
21#[cfg(feature = "oauth")]
22pub mod subject;
23pub mod tasks;
24pub mod tools;
25
26#[derive(Debug, Clone, PartialEq, Eq)]
27pub struct PoolOptions {
28    pub response_caps: ResponseCaps,
29    pub discovery_concurrency: usize,
30}
31
32impl Default for PoolOptions {
33    fn default() -> Self {
34        Self {
35            response_caps: ResponseCaps::default(),
36            discovery_concurrency: 8,
37        }
38    }
39}
40
41impl PoolOptions {
42    #[must_use]
43    pub fn normalized(mut self) -> Self {
44        self.discovery_concurrency = self.discovery_concurrency.max(1);
45        self
46    }
47}
48
49#[derive(Debug, Clone, PartialEq, Eq)]
50pub struct ToolCall {
51    pub upstream: String,
52    pub tool: String,
53    pub params: Value,
54}
55
56#[derive(Debug, Clone)]
57pub struct InProcessUpstream {
58    snapshot: UpstreamSnapshot,
59    tool_results: BTreeMap<String, Value>,
60}
61
62impl InProcessUpstream {
63    #[must_use]
64    pub fn new(name: impl Into<String>) -> Self {
65        Self {
66            snapshot: UpstreamSnapshot::empty(name, TransportKind::InProcess),
67            tool_results: BTreeMap::new(),
68        }
69    }
70
71    #[must_use]
72    pub fn with_tool(mut self, tool: ToolDescriptor, result: Value) -> Self {
73        self.tool_results.insert(tool.name.clone(), result);
74        self.snapshot.tools.push(tool);
75        self
76    }
77
78    #[must_use]
79    pub fn with_snapshot(mut self, snapshot: UpstreamSnapshot) -> Self {
80        self.snapshot = snapshot;
81        self
82    }
83
84    fn call_tool(&self, call: &ToolCall) -> Result<Value, UpstreamError> {
85        if !call.params.is_object() {
86            return Err(UpstreamError::ParamsMustBeObject);
87        }
88        self.tool_results
89            .get(&call.tool)
90            .cloned()
91            .ok_or_else(|| UpstreamError::NotExposed {
92                upstream: call.upstream.clone(),
93                item: call.tool.clone(),
94            })
95    }
96}
97
98struct PoolEntry {
99    config: UpstreamConfig,
100    snapshot: UpstreamSnapshot,
101    in_process: Option<InProcessUpstream>,
102    live: Option<Arc<live::LiveUpstream>>,
103}
104
105#[cfg(feature = "oauth")]
106struct SubjectPoolEntry {
107    snapshot: UpstreamSnapshot,
108    live: Arc<live::LiveUpstream>,
109}
110
111#[derive(Clone)]
112pub struct UpstreamPool {
113    entries: Arc<RwLock<BTreeMap<String, PoolEntry>>>,
114    #[cfg(feature = "oauth")]
115    subject_entries: Arc<RwLock<BTreeMap<(String, String), SubjectPoolEntry>>>,
116    #[cfg(feature = "oauth")]
117    oauth_provider: Arc<RwLock<Option<Arc<dyn crate::oauth::UpstreamOAuthProvider>>>>,
118    options: PoolOptions,
119}
120
121impl Default for UpstreamPool {
122    fn default() -> Self {
123        Self::new(PoolOptions::default())
124    }
125}
126
127impl UpstreamPool {
128    #[must_use]
129    pub fn new(options: PoolOptions) -> Self {
130        Self {
131            entries: Arc::new(RwLock::new(BTreeMap::new())),
132            #[cfg(feature = "oauth")]
133            subject_entries: Arc::new(RwLock::new(BTreeMap::new())),
134            #[cfg(feature = "oauth")]
135            oauth_provider: Arc::new(RwLock::new(None)),
136            options: options.normalized(),
137        }
138    }
139
140    #[must_use]
141    pub fn response_caps(&self) -> &ResponseCaps {
142        &self.options.response_caps
143    }
144
145    #[must_use]
146    pub fn discovery_concurrency(&self) -> usize {
147        self.options.discovery_concurrency
148    }
149
150    pub fn register_config(&self, config: UpstreamConfig) -> Result<(), UpstreamError> {
151        let transport = transport_for_config(&config);
152        let health = if config.enabled {
153            health_for_config(config.name.as_str(), transport)
154        } else {
155            UpstreamHealth::Disabled
156        };
157        let mut snapshot = UpstreamSnapshot::empty(config.name.clone(), transport);
158        snapshot.health = health;
159        let entry = PoolEntry {
160            config: config.clone(),
161            snapshot,
162            in_process: None,
163            live: None,
164        };
165        self.entries
166            .write()
167            .expect("upstream pool lock poisoned")
168            .insert(config.name.clone(), entry);
169        Ok(())
170    }
171
172    pub fn register_in_process(
173        &self,
174        config: UpstreamConfig,
175        upstream: InProcessUpstream,
176    ) -> Result<(), UpstreamError> {
177        let mut snapshot = upstream.snapshot.clone();
178        snapshot.name = config.name.clone();
179        snapshot.transport = TransportKind::InProcess;
180        snapshot.health = if config.enabled {
181            UpstreamHealth::Connected
182        } else {
183            UpstreamHealth::Disabled
184        };
185        let entry = PoolEntry {
186            config: config.clone(),
187            snapshot,
188            in_process: Some(upstream),
189            live: None,
190        };
191        self.entries
192            .write()
193            .expect("upstream pool lock poisoned")
194            .insert(config.name.clone(), entry);
195        Ok(())
196    }
197
198    pub async fn call_tool(&self, call: ToolCall) -> Result<Value, UpstreamError> {
199        self.ensure_connected(&call.upstream).await?;
200        let live_peer = {
201            let entries = self.entries.read().expect("upstream pool lock poisoned");
202            let entry =
203                entries
204                    .get(&call.upstream)
205                    .ok_or_else(|| UpstreamError::UnknownUpstream {
206                        upstream: call.upstream.clone(),
207                    })?;
208            ensure_routable(entry)?;
209            tools::ensure_tool_exposed(entry, &call.tool)?;
210            if let Some(in_process) = &entry.in_process {
211                let result = in_process.call_tool(&call)?;
212                let bytes = serde_json::to_vec(&result).map_or(usize::MAX, |bytes| bytes.len());
213                self.response_caps()
214                    .enforce(crate::upstream::CapScope::ToolsCall, bytes)?;
215                return Ok(result);
216            }
217            entry.live.as_ref().map(|live| live.peer())
218        };
219        let Some(peer) = live_peer else {
220            return Err(UpstreamError::Unsupported {
221                upstream: call.upstream,
222                capability: "tools/call",
223            });
224        };
225        let upstream = call.upstream.clone();
226        let result = live::call_live_tool(&upstream, peer, call.tool, call.params).await?;
227        let bytes = serde_json::to_vec(&result).map_or(usize::MAX, |bytes| bytes.len());
228        self.response_caps()
229            .enforce(crate::upstream::CapScope::ToolsCall, bytes)?;
230        Ok(result)
231    }
232
233    pub async fn call_tool_once(
234        &self,
235        call: ToolCall,
236        round_trip: McpRoundTrip,
237    ) -> Result<McpRequestOutcome, UpstreamError> {
238        self.ensure_connected(&call.upstream).await?;
239        let (live_peer, config) = {
240            let entries = self.entries.read().expect("upstream pool lock poisoned");
241            let entry =
242                entries
243                    .get(&call.upstream)
244                    .ok_or_else(|| UpstreamError::UnknownUpstream {
245                        upstream: call.upstream.clone(),
246                    })?;
247            ensure_routable(entry)?;
248            tools::ensure_tool_exposed(entry, &call.tool)?;
249            if let Some(in_process) = &entry.in_process {
250                let value = in_process.call_tool(&call)?;
251                let result = rmcp::model::CallToolResult::structured(value);
252                let value =
253                    serde_json::to_value(result).map_err(|error| UpstreamError::LiveCall {
254                        upstream: call.upstream.clone(),
255                        operation: "tools/call",
256                        message: error.to_string(),
257                    })?;
258                let outcome = McpRequestOutcome::Complete(value);
259                let bytes =
260                    serde_json::to_vec(outcome.payload()).map_or(usize::MAX, |bytes| bytes.len());
261                self.response_caps()
262                    .enforce(crate::upstream::CapScope::ToolsCall, bytes)?;
263                return Ok(outcome);
264            }
265            (
266                entry.live.as_ref().map(|live| live.peer()),
267                entry.config.clone(),
268            )
269        };
270        let upstream = call.upstream.clone();
271        let outcome = if round_trip.request_meta.is_some() {
272            live::call_live_tool_once_scoped(
273                &config,
274                live::LiveConnectContext::shared(self.response_caps()),
275                call.tool,
276                call.params,
277                round_trip,
278            )
279            .await?
280        } else {
281            let Some(peer) = live_peer else {
282                return Err(UpstreamError::Unsupported {
283                    upstream: call.upstream,
284                    capability: "tools/call",
285                });
286            };
287            live::call_live_tool_once(&upstream, peer, call.tool, call.params, round_trip).await?
288        };
289        let bytes = serde_json::to_vec(outcome.payload()).map_or(usize::MAX, |bytes| bytes.len());
290        self.response_caps()
291            .enforce(crate::upstream::CapScope::ToolsCall, bytes)?;
292        Ok(outcome)
293    }
294
295    pub async fn ensure_connected(&self, upstream: &str) -> Result<(), UpstreamError> {
296        let config = {
297            let entries = self.entries.read().expect("upstream pool lock poisoned");
298            let entry = entries
299                .get(upstream)
300                .ok_or_else(|| UpstreamError::UnknownUpstream {
301                    upstream: upstream.to_owned(),
302                })?;
303            if entry.in_process.is_some() || entry.live.is_some() || !entry.config.enabled {
304                return Ok(());
305            }
306            entry.config.clone()
307        };
308        let context = live::LiveConnectContext::shared(self.response_caps());
309        let (live, snapshot) = live::connect_live(&config, &SpawnGuard::default(), context).await?;
310        let mut entries = self.entries.write().expect("upstream pool lock poisoned");
311        let entry = entries
312            .get_mut(upstream)
313            .ok_or_else(|| UpstreamError::UnknownUpstream {
314                upstream: upstream.to_owned(),
315            })?;
316        entry.snapshot = snapshot;
317        entry.live = Some(Arc::new(live));
318        Ok(())
319    }
320
321    pub async fn refresh_all(&self) {
322        let names = self
323            .entries
324            .read()
325            .expect("upstream pool lock poisoned")
326            .keys()
327            .cloned()
328            .collect::<Vec<_>>();
329        for name in names {
330            if let Err(error) = self.ensure_connected(&name).await {
331                let _ = self.record_discovery_error(&name, error);
332            }
333        }
334    }
335
336    pub(super) fn record_discovery_error(
337        &self,
338        upstream: &str,
339        error: UpstreamError,
340    ) -> Result<(), UpstreamError> {
341        let mut entries = self.entries.write().expect("upstream pool lock poisoned");
342        let entry = entries
343            .get_mut(upstream)
344            .ok_or_else(|| UpstreamError::UnknownUpstream {
345                upstream: upstream.to_owned(),
346            })?;
347        match error {
348            UpstreamError::Unsupported {
349                upstream,
350                capability,
351            } => {
352                entry.snapshot.health = UpstreamHealth::Unsupported {
353                    reason: format!("upstream `{upstream}` does not support `{capability}`"),
354                };
355            }
356            other => {
357                entry.snapshot.health = UpstreamHealth::Degraded {
358                    consecutive_failures: 1,
359                    error: Some(other.to_string()),
360                };
361                entry.snapshot.stale = true;
362            }
363        }
364        Ok(())
365    }
366
367    fn snapshots(&self) -> Vec<UpstreamSnapshot> {
368        self.entries
369            .read()
370            .expect("upstream pool lock poisoned")
371            .values()
372            .map(|entry| entry.snapshot.clone())
373            .collect()
374    }
375
376    fn with_entry<T>(
377        &self,
378        upstream: &str,
379        f: impl FnOnce(&PoolEntry) -> Result<T, UpstreamError>,
380    ) -> Result<T, UpstreamError> {
381        let entries = self.entries.read().expect("upstream pool lock poisoned");
382        let entry = entries
383            .get(upstream)
384            .ok_or_else(|| UpstreamError::UnknownUpstream {
385                upstream: upstream.to_owned(),
386            })?;
387        f(entry)
388    }
389}
390
391fn ensure_routable(entry: &PoolEntry) -> Result<(), UpstreamError> {
392    if entry.snapshot.health.is_routable() {
393        return Ok(());
394    }
395    Err(UpstreamError::NotRoutable {
396        upstream: entry.snapshot.name.clone(),
397        reason: health_reason(&entry.snapshot.health),
398    })
399}
400
401fn health_reason(health: &UpstreamHealth) -> String {
402    match health {
403        UpstreamHealth::Connected => "connected".to_owned(),
404        UpstreamHealth::Disabled => "disabled".to_owned(),
405        UpstreamHealth::Degraded { error, .. } => error
406            .clone()
407            .unwrap_or_else(|| "capability degraded".to_owned()),
408        UpstreamHealth::Unsupported { reason } => reason.clone(),
409    }
410}
411
412fn transport_for_config(config: &UpstreamConfig) -> TransportKind {
413    if let Some(url) = config.url.as_deref() {
414        return transport_kind_for_decision(&decide_http_transport(url));
415    }
416    if config.command.is_some() {
417        return TransportKind::Stdio;
418    }
419    TransportKind::InProcess
420}
421
422fn health_for_config(name: &str, transport: TransportKind) -> UpstreamHealth {
423    match transport {
424        TransportKind::InProcess => UpstreamHealth::Unsupported {
425            reason: format!("configured upstream `{name}` has no live in-process connector"),
426        },
427        TransportKind::HttpJson | TransportKind::HttpSse | TransportKind::WebSocket => {
428            UpstreamHealth::Unsupported {
429                reason: format!("live upstream `{name}` is not connected yet"),
430            }
431        }
432        TransportKind::Stdio => UpstreamHealth::Unsupported {
433            reason: format!("live stdio upstream `{name}` is not connected yet"),
434        },
435    }
436}
437
438#[cfg(test)]
439#[path = "pool_tests.rs"]
440mod tests;
soma_mcp_client/upstream/pool.rs

soma_mcp_client/upstream/
pool.rs
