Skip to main content

soma_mcp_client/
upstream.rs

1//! Self-contained upstream runtime primitives.
2
3pub mod http_body_cap;
4pub mod http_client;
5pub mod pool;
6pub mod relay;
7pub mod transport;
8
9use serde::{Deserialize, Serialize};
10use serde_json::Value;
11use thiserror::Error;
12
13#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
14pub struct ToolDescriptor {
15    pub name: String,
16    #[serde(default, skip_serializing_if = "Option::is_none")]
17    pub description: Option<String>,
18    #[serde(default, skip_serializing_if = "Option::is_none")]
19    pub input_schema: Option<Value>,
20    #[serde(default, skip_serializing_if = "Option::is_none")]
21    pub output_schema: Option<Value>,
22    #[serde(default)]
23    pub destructive: bool,
24}
25
26impl ToolDescriptor {
27    #[must_use]
28    pub fn new(name: impl Into<String>) -> Self {
29        Self {
30            name: name.into(),
31            description: None,
32            input_schema: None,
33            output_schema: None,
34            destructive: true,
35        }
36    }
37}
38
39#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
40pub struct ResourceDescriptor {
41    pub uri: String,
42    #[serde(default, skip_serializing_if = "Option::is_none")]
43    pub name: Option<String>,
44}
45
46#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
47pub struct PromptDescriptor {
48    pub name: String,
49    #[serde(default, skip_serializing_if = "Option::is_none")]
50    pub description: Option<String>,
51}
52
53#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
54pub enum TransportKind {
55    InProcess,
56    HttpJson,
57    HttpSse,
58    Stdio,
59    WebSocket,
60}
61
62#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
63pub enum UpstreamHealth {
64    Connected,
65    Degraded {
66        consecutive_failures: u32,
67        error: Option<String>,
68    },
69    Disabled,
70    Unsupported {
71        reason: String,
72    },
73}
74
75impl UpstreamHealth {
76    #[must_use]
77    pub fn is_routable(&self) -> bool {
78        matches!(self, Self::Connected)
79    }
80}
81
82#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
83pub struct UpstreamSnapshot {
84    pub name: String,
85    pub transport: TransportKind,
86    pub health: UpstreamHealth,
87    pub tools: Vec<ToolDescriptor>,
88    pub resources: Vec<ResourceDescriptor>,
89    pub prompts: Vec<PromptDescriptor>,
90    #[serde(default)]
91    pub stale: bool,
92}
93
94impl UpstreamSnapshot {
95    #[must_use]
96    pub fn empty(name: impl Into<String>, transport: TransportKind) -> Self {
97        Self {
98            name: name.into(),
99            transport,
100            health: UpstreamHealth::Connected,
101            tools: Vec::new(),
102            resources: Vec::new(),
103            prompts: Vec::new(),
104            stale: false,
105        }
106    }
107}
108
109#[derive(Debug, Clone, Copy, PartialEq, Eq)]
110pub enum CapScope {
111    ToolsList,
112    ToolsCall,
113    ResourcesList,
114    ResourcesRead,
115    PromptsList,
116    PromptsGet,
117    RelayCall,
118    HttpJson,
119    HttpSseEvent,
120    WebSocketFrame,
121    StdioMessage,
122}
123
124#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
125pub struct ResponseCaps {
126    pub tools_list_bytes: usize,
127    pub tools_call_bytes: usize,
128    pub resources_list_bytes: usize,
129    pub resources_read_bytes: usize,
130    pub prompts_list_bytes: usize,
131    pub prompts_get_bytes: usize,
132    pub relay_call_bytes: usize,
133    pub http_json_bytes: usize,
134    pub http_sse_event_bytes: usize,
135    pub websocket_frame_bytes: usize,
136    pub stdio_message_bytes: usize,
137}
138
139impl Default for ResponseCaps {
140    fn default() -> Self {
141        Self {
142            tools_list_bytes: 512 * 1024,
143            tools_call_bytes: 2 * 1024 * 1024,
144            resources_list_bytes: 512 * 1024,
145            resources_read_bytes: 2 * 1024 * 1024,
146            prompts_list_bytes: 256 * 1024,
147            prompts_get_bytes: 512 * 1024,
148            relay_call_bytes: 4 * 1024 * 1024,
149            http_json_bytes: 2 * 1024 * 1024,
150            http_sse_event_bytes: 512 * 1024,
151            websocket_frame_bytes: 512 * 1024,
152            stdio_message_bytes: 2 * 1024 * 1024,
153        }
154    }
155}
156
157impl ResponseCaps {
158    #[must_use]
159    pub fn limit_for(&self, scope: CapScope) -> usize {
160        match scope {
161            CapScope::ToolsList => self.tools_list_bytes,
162            CapScope::ToolsCall => self.tools_call_bytes,
163            CapScope::ResourcesList => self.resources_list_bytes,
164            CapScope::ResourcesRead => self.resources_read_bytes,
165            CapScope::PromptsList => self.prompts_list_bytes,
166            CapScope::PromptsGet => self.prompts_get_bytes,
167            CapScope::RelayCall => self.relay_call_bytes,
168            CapScope::HttpJson => self.http_json_bytes,
169            CapScope::HttpSseEvent => self.http_sse_event_bytes,
170            CapScope::WebSocketFrame => self.websocket_frame_bytes,
171            CapScope::StdioMessage => self.stdio_message_bytes,
172        }
173    }
174
175    pub fn enforce(&self, scope: CapScope, observed_bytes: usize) -> Result<(), UpstreamError> {
176        let limit = self.limit_for(scope);
177        if observed_bytes > limit {
178            return Err(UpstreamError::ResponseTooLarge {
179                scope,
180                observed_bytes,
181                limit,
182            });
183        }
184        Ok(())
185    }
186}
187
188#[derive(Debug, Error, PartialEq, Eq)]
189pub enum UpstreamError {
190    #[error("upstream `{upstream}` is not configured")]
191    UnknownUpstream { upstream: String },
192    #[error("upstream `{upstream}` is not routable: {reason}")]
193    NotRoutable { upstream: String, reason: String },
194    #[error("upstream `{upstream}` does not expose `{item}`")]
195    NotExposed { upstream: String, item: String },
196    #[error("upstream `{upstream}` does not support `{capability}`")]
197    Unsupported {
198        upstream: String,
199        capability: &'static str,
200    },
201    #[error("upstream `{upstream}` failed to connect: {message}")]
202    LiveConnect { upstream: String, message: String },
203    #[error("upstream `{upstream}` {operation} failed: {message}")]
204    LiveCall {
205        upstream: String,
206        operation: &'static str,
207        message: String,
208    },
209    #[error("{scope:?} payload was {observed_bytes} bytes, exceeding {limit} bytes")]
210    ResponseTooLarge {
211        scope: CapScope,
212        observed_bytes: usize,
213        limit: usize,
214    },
215    #[error("tool params must be a JSON object")]
216    ParamsMustBeObject,
217}
218
219impl UpstreamError {
220    pub(crate) fn connect(
221        config: &crate::config::UpstreamConfig,
222        error: impl std::fmt::Display,
223    ) -> Self {
224        Self::LiveConnect {
225            upstream: config.name.clone(),
226            message: error.to_string(),
227        }
228    }
229}
230
231#[cfg(test)]
232#[path = "upstream_tests.rs"]
233mod tests;