1pub 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;