1use serde::{Deserialize, Serialize};
9use std::collections::BTreeMap;
10
11use crate::core::addons::capabilities::AddonCapabilities;
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default, Hash)]
15#[serde(rename_all = "snake_case")]
16pub enum TransportKind {
17 #[default]
19 Stdio,
20 Http,
22}
23
24impl TransportKind {
25 pub fn as_str(self) -> &'static str {
26 match self {
27 TransportKind::Stdio => "stdio",
28 TransportKind::Http => "http",
29 }
30 }
31}
32
33#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
39#[serde(default)]
40pub struct GatewayServer {
41 pub name: String,
43 pub transport: TransportKind,
45 pub enabled: bool,
47
48 pub command: String,
51 pub args: Vec<String>,
53 pub env: BTreeMap<String, String>,
55 #[serde(default, skip_serializing_if = "String::is_empty")]
61 pub binary_sha256: String,
62
63 pub url: String,
66 pub headers: BTreeMap<String, String>,
68
69 #[serde(default, skip_serializing_if = "Option::is_none")]
75 pub capabilities: Option<AddonCapabilities>,
76
77 #[serde(default, skip_serializing_if = "String::is_empty")]
83 pub integration: String,
84}
85
86impl Default for GatewayServer {
87 fn default() -> Self {
88 Self {
89 name: String::new(),
90 transport: TransportKind::Stdio,
91 enabled: true,
92 command: String::new(),
93 args: Vec::new(),
94 env: BTreeMap::new(),
95 binary_sha256: String::new(),
96 url: String::new(),
97 headers: BTreeMap::new(),
98 capabilities: None,
99 integration: String::new(),
100 }
101 }
102}
103
104#[derive(Debug, Clone, PartialEq, Eq)]
106pub enum ResolvedTransport {
107 Stdio {
108 command: String,
109 args: Vec<String>,
110 env: BTreeMap<String, String>,
111 binary_sha256: String,
113 capabilities: Option<AddonCapabilities>,
115 },
116 Http {
117 url: String,
118 headers: BTreeMap<String, String>,
119 },
120}
121
122impl GatewayServer {
123 pub fn resolve(&self) -> Result<ResolvedTransport, String> {
126 if self.name.trim().is_empty() {
127 return Err("gateway server is missing a `name`".to_string());
128 }
129 match self.transport {
130 TransportKind::Stdio => {
131 if self.command.trim().is_empty() {
132 return Err(format!(
133 "gateway server `{}` uses stdio transport but has no `command`",
134 self.name
135 ));
136 }
137 Ok(ResolvedTransport::Stdio {
138 command: self.command.clone(),
139 args: self.args.clone(),
140 env: self.env.clone(),
141 binary_sha256: self.binary_sha256.clone(),
142 capabilities: self.capabilities.clone(),
143 })
144 }
145 TransportKind::Http => {
146 let url = self.url.trim();
147 if !(url.starts_with("http://") || url.starts_with("https://")) {
148 return Err(format!(
149 "gateway server `{}` uses http transport but `url` is not http(s)",
150 self.name
151 ));
152 }
153 Ok(ResolvedTransport::Http {
154 url: url.to_string(),
155 headers: self.headers.clone(),
156 })
157 }
158 }
159 }
160}
161
162#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
164#[serde(default)]
165pub struct GatewayConfig {
166 pub enabled: bool,
168 pub top_n: usize,
170 pub cache_ttl_secs: u64,
172 pub call_timeout_secs: u64,
174 pub servers: Vec<GatewayServer>,
176
177 pub compress_output: bool,
183 pub handle_spill: bool,
187 pub index_output: bool,
191 pub output_budget_tokens: usize,
194}
195
196impl Default for GatewayConfig {
197 fn default() -> Self {
198 Self {
199 enabled: false,
200 top_n: 5,
201 cache_ttl_secs: 300,
202 call_timeout_secs: 30,
203 servers: Vec::new(),
204 compress_output: false,
205 handle_spill: false,
206 index_output: false,
207 output_budget_tokens: 2000,
208 }
209 }
210}
211
212impl GatewayConfig {
213 pub fn enabled_effective(&self) -> bool {
216 if let Ok(v) = std::env::var("LEAN_CTX_GATEWAY") {
217 return !matches!(v.trim(), "0" | "false" | "off");
218 }
219 self.enabled
220 }
221
222 pub fn active_servers(&self) -> impl Iterator<Item = &GatewayServer> {
224 self.servers.iter().filter(|s| s.enabled)
225 }
226
227 pub fn effective_top_n(&self) -> usize {
229 self.top_n.clamp(1, 50)
230 }
231
232 pub fn postprocess_active(&self) -> bool {
236 self.compress_output || self.handle_spill || self.index_output
237 }
238
239 pub fn effective_output_budget(&self) -> usize {
243 self.output_budget_tokens.max(256)
244 }
245}
246
247#[cfg(test)]
248mod tests {
249 use super::*;
250
251 #[test]
252 fn default_is_disabled_noop() {
253 let cfg = GatewayConfig::default();
254 assert!(!cfg.enabled);
255 assert!(!cfg.enabled_effective());
256 assert_eq!(cfg.effective_top_n(), 5);
257 assert!(cfg.servers.is_empty());
258 assert!(!cfg.compress_output);
260 assert!(!cfg.handle_spill);
261 assert!(!cfg.index_output);
262 assert!(!cfg.postprocess_active());
263 assert_eq!(cfg.effective_output_budget(), 2000);
264 }
265
266 #[test]
267 fn zero_budget_floors_to_sane_minimum() {
268 let cfg = GatewayConfig {
269 output_budget_tokens: 0,
270 ..Default::default()
271 };
272 assert_eq!(cfg.effective_output_budget(), 256);
273 }
274
275 #[test]
276 fn server_integration_field_round_trips() {
277 let toml_src = r#"
278enabled = true
279compress_output = true
280index_output = true
281output_budget_tokens = 1500
282
283[[servers]]
284name = "repomix"
285command = "npx"
286args = ["-y", "repomix", "--mcp"]
287integration = "codebase-pack"
288"#;
289 let cfg: GatewayConfig = toml::from_str(toml_src).expect("parse");
290 assert!(cfg.compress_output);
291 assert!(cfg.index_output);
292 assert!(cfg.postprocess_active());
293 assert_eq!(cfg.effective_output_budget(), 1500);
294 assert_eq!(cfg.servers[0].integration, "codebase-pack");
295 let back = toml::to_string(&cfg).expect("serialize");
297 assert!(back.contains("integration = \"codebase-pack\""));
298 }
299
300 #[test]
301 fn stdio_server_resolves_with_command() {
302 let s = GatewayServer {
303 name: "fs".into(),
304 transport: TransportKind::Stdio,
305 command: "mcp-fs".into(),
306 args: vec!["/tmp".into()],
307 ..Default::default()
308 };
309 let r = s.resolve().expect("resolve");
310 assert_eq!(
311 r,
312 ResolvedTransport::Stdio {
313 command: "mcp-fs".into(),
314 args: vec!["/tmp".into()],
315 env: BTreeMap::new(),
316 binary_sha256: String::new(),
317 capabilities: None,
318 }
319 );
320 }
321
322 #[test]
323 fn stdio_without_command_is_error() {
324 let s = GatewayServer {
325 name: "broken".into(),
326 transport: TransportKind::Stdio,
327 ..Default::default()
328 };
329 assert!(s.resolve().is_err());
330 }
331
332 #[test]
333 fn http_requires_http_scheme() {
334 let ok = GatewayServer {
335 name: "remote".into(),
336 transport: TransportKind::Http,
337 url: "https://example.com/mcp".into(),
338 ..Default::default()
339 };
340 assert!(ok.resolve().is_ok());
341
342 let bad = GatewayServer {
343 name: "remote".into(),
344 transport: TransportKind::Http,
345 url: "ftp://example.com".into(),
346 ..Default::default()
347 };
348 assert!(bad.resolve().is_err());
349 }
350
351 #[test]
352 fn unnamed_server_is_error() {
353 let s = GatewayServer {
354 transport: TransportKind::Stdio,
355 command: "x".into(),
356 ..Default::default()
357 };
358 assert!(s.resolve().is_err());
359 }
360
361 #[test]
362 fn active_servers_skips_disabled() {
363 let cfg = GatewayConfig {
364 enabled: true,
365 servers: vec![
366 GatewayServer {
367 name: "a".into(),
368 command: "a".into(),
369 enabled: true,
370 ..Default::default()
371 },
372 GatewayServer {
373 name: "b".into(),
374 command: "b".into(),
375 enabled: false,
376 ..Default::default()
377 },
378 ],
379 ..Default::default()
380 };
381 let active: Vec<_> = cfg.active_servers().map(|s| s.name.as_str()).collect();
382 assert_eq!(active, vec!["a"]);
383 }
384
385 #[test]
386 fn parses_array_of_tables_toml() {
387 let toml_src = r#"
388enabled = true
389top_n = 8
390
391[[servers]]
392name = "fs"
393transport = "stdio"
394command = "mcp-server-filesystem"
395args = ["/tmp"]
396
397[[servers]]
398name = "remote"
399transport = "http"
400url = "https://example.com/mcp"
401enabled = false
402"#;
403 let cfg: GatewayConfig = toml::from_str(toml_src).expect("parse");
404 assert!(cfg.enabled);
405 assert_eq!(cfg.top_n, 8);
406 assert_eq!(cfg.servers.len(), 2);
407 assert_eq!(cfg.servers[0].transport, TransportKind::Stdio);
408 assert_eq!(cfg.servers[0].command, "mcp-server-filesystem");
409 assert_eq!(cfg.servers[1].transport, TransportKind::Http);
410 assert!(!cfg.servers[1].enabled);
411 }
412}