Vakyartha documentation

1use std::collections::HashMap;
2use std::path::PathBuf;
3use std::sync::Arc;
4use std::time::{Duration, Instant};
5
6use tokio::sync::Mutex;
7
8use serde_json::Value;
9
10use crate::client::{McpClient, McpError, McpToolInfo, NotificationSink, ServerConfig};
11
12/// How long listing a server (connect + `tools/list`) may take before the
13/// attempt counts as a failure. Far below the 60s per-request ceiling: a
14/// server that cannot even list its tools should cost the turn seconds, not
15/// a minute.
16pub const LIST_TIMEOUT: Duration = Duration::from_secs(10);
17
18/// First wait after a failed attempt; doubles per consecutive failure up to
19/// [`FAILURE_BACKOFF_CAP`]. The pool never retries on its own — this only
20/// stops repeated demand from respawning a broken server on every call.
21const FAILURE_BACKOFF_BASE: Duration = Duration::from_secs(5);
22const FAILURE_BACKOFF_CAP: Duration = Duration::from_secs(5 * 60);
23
24/// How long a pooled connection may sit unused before it is shut down. A
25/// process that stays up for weeks must not hold a subprocess for every
26/// server it has ever touched; the next call respawns on demand.
27pub const IDLE_TTL: Duration = Duration::from_secs(15 * 60);
28
29/// One pooled connection plus the bookkeeping the pool needs to decide
30/// whether to keep it.
31struct Pooled {
32 client: Arc<McpClient>,
33 last_used: Instant,
34}
35
36/// The outcome of listing one server's catalog.
37///
38/// A failure is a `Err(reason)`, never a synthesised tool. The old shape
39/// returned a tool literally named `error`, which made a dead server look
40/// like a working one to the prompt, to the alias table, and to the retry
41/// guard that then refused to try again.
42pub type ListOutcome = Result<Vec<McpToolInfo>, String>;
43
44/// What demand has taught the pool about one server. Nothing here is ever
45/// learned by spawning a server for its own sake: a server nobody has asked
46/// for has an empty observation.
47#[derive(Debug, Clone, Default, PartialEq)]
48pub struct ServerObservation {
49 /// The catalog the last live connection returned, redacted and with
50 /// descriptions trimmed. `None` until the server is first used.
51 pub tools: Option<Vec<McpToolInfo>>,
52 /// Why the last attempt failed, until an attempt succeeds.
53 pub failure: Option<String>,
54}
55
56#[derive(Default)]
57struct ServerState {
58 observation: ServerObservation,
59 attempts: u32,
60 failed_at: Option<Instant>,
61}
62
63impl ServerState {
64 fn retry_after(&self, now: Instant) -> Option<Duration> {
65 let failed_at = self.failed_at?;
66 let shift = self.attempts.saturating_sub(1).min(16);
67 let wait = FAILURE_BACKOFF_BASE
68 .saturating_mul(1u32 << shift)
69 .min(FAILURE_BACKOFF_CAP);
70 wait.checked_sub(now.duration_since(failed_at))
71 }
72}
73
74/// The on-demand connection pool: one client per configured server, spawned
75/// only when a call needs it, shared by every session of the owning `Core`,
76/// and shut down after [`IDLE_TTL`] unused. It never connects on its own
77/// initiative — no warm-up, no background probe, no respawn after eviction.
78pub struct McpManager {
79 servers: HashMap<String, ServerConfig>,
80 clients: Mutex<HashMap<String, Pooled>>,
81 state: std::sync::Mutex<HashMap<String, ServerState>>,
82 cwd: PathBuf,
83 sandbox: Option<Arc<dyn vak_tools::sandbox::Sandbox>>,
84 notifications: Option<NotificationSink>,
85 observer: Option<tokio::sync::mpsc::UnboundedSender<String>>,
86}
87
88impl McpManager {
89 pub fn new(servers: HashMap<String, ServerConfig>, cwd: PathBuf) -> Self {
90 Self::new_sandboxed(servers, cwd, None)
91 }
92
93 pub fn new_sandboxed(
94 servers: HashMap<String, ServerConfig>,
95 cwd: PathBuf,
96 sandbox: Option<Arc<dyn vak_tools::sandbox::Sandbox>>,
97 ) -> Self {
98 McpManager {
99 servers,
100 clients: Mutex::new(HashMap::new()),
101 state: std::sync::Mutex::new(HashMap::new()),
102 cwd,
103 sandbox,
104 notifications: None,
105 observer: None,
106 }
107 }
108
109 /// Receive a server's name whenever what the pool knows about it
110 /// changes — a catalog learned or changed, a failure recorded or cleared.
111 pub fn with_observer(mut self, sink: tokio::sync::mpsc::UnboundedSender<String>) -> Self {
112 self.observer = Some(sink);
113 self
114 }
115
116 /// What demand has observed about every server, for the capability
117 /// declarations. Synchronous and free of I/O.
118 pub fn observations(&self) -> HashMap<String, ServerObservation> {
119 self.state
120 .lock()
121 .unwrap_or_else(std::sync::PoisonError::into_inner)
122 .iter()
123 .map(|(name, state)| (name.clone(), state.observation.clone()))
124 .collect()
125 }
126
127 /// Forget a server's catalog, e.g. after it announced
128 /// `notifications/tools/list_changed`. The next use re-lists it.
129 pub fn forget_catalog(&self, server: &str) {
130 let changed = {
131 let mut state = self
132 .state
133 .lock()
134 .unwrap_or_else(std::sync::PoisonError::into_inner);
135 state
136 .get_mut(server)
137 .and_then(|entry| entry.observation.tools.take())
138 .is_some()
139 };
140 if changed {
141 self.announce(server);
142 }
143 }
144
145 fn announce(&self, server: &str) {
146 if let Some(observer) = &self.observer {
147 let _ = observer.send(server.to_string());
148 }
149 }
150
151 fn record_tools(&self, server: &str, tools: &[McpToolInfo]) {
152 let observed: Vec<McpToolInfo> = tools
153 .iter()
154 .map(|tool| McpToolInfo {
155 name: tool.name.clone(),
156 description: self.redact(tool.description.clone()),
157 input_schema: self.redact_json(&tool.input_schema),
158 })
159 .collect();
160 let changed = {
161 let mut state = self
162 .state
163 .lock()
164 .unwrap_or_else(std::sync::PoisonError::into_inner);
165 let entry = state.entry(server.to_string()).or_default();
166 let changed = entry.observation.tools.as_ref() != Some(&observed)
167 || entry.observation.failure.is_some();
168 entry.observation.tools = Some(observed);
169 entry.observation.failure = None;
170 entry.attempts = 0;
171 entry.failed_at = None;
172 changed
173 };
174 if changed {
175 self.announce(server);
176 }
177 }
178
179 fn record_failure(&self, server: &str, reason: &str) {
180 let reason = self.redact(reason);
181 let changed = {
182 let mut state = self
183 .state
184 .lock()
185 .unwrap_or_else(std::sync::PoisonError::into_inner);
186 let entry = state.entry(server.to_string()).or_default();
187 let changed = entry.observation.failure.as_deref() != Some(reason.as_str());
188 entry.observation.failure = Some(reason);
189 entry.attempts = entry.attempts.saturating_add(1);
190 entry.failed_at = Some(Instant::now());
191 changed
192 };
193 if changed {
194 self.announce(server);
195 }
196 }
197
198 /// The failure still inside its backoff window, if any: repeated demand
199 /// gets the recorded reason back instead of a fresh spawn.
200 fn backing_off(&self, server: &str) -> Option<String> {
201 let state = self
202 .state
203 .lock()
204 .unwrap_or_else(std::sync::PoisonError::into_inner);
205 let entry = state.get(server)?;
206 let wait = entry.retry_after(Instant::now())?;
207 let reason = entry.observation.failure.as_deref().unwrap_or("failed");
208 Some(format!(
209 "mcp server '{server}' is unavailable: {reason} (next attempt in {}s)",
210 wait.as_secs().max(1)
211 ))
212 }
213
214 /// Route server-initiated notifications (catalog changes) to `sink`.
215 /// Applies to connections opened after this call.
216 pub fn with_notifications(mut self, sink: NotificationSink) -> Self {
217 self.notifications = Some(sink);
218 self
219 }
220
221 pub fn server_names(&self) -> Vec<String> {
222 let mut names: Vec<String> = self.servers.keys().cloned().collect();
223 names.sort();
224 names
225 }
226
227 async fn get(&self, server: &str) -> Result<Arc<McpClient>, McpError> {
228 if let Some(reason) = self.backing_off(server) {
229 return Err(McpError::Protocol(reason));
230 }
231 {
232 let mut pool = self.clients.lock().await;
233 if let Some(entry) = pool.get_mut(server) {
234 if entry.client.is_alive() {
235 entry.last_used = Instant::now();
236 return Ok(entry.client.clone());
237 }
238 // Dead: drop it and fall through to a fresh spawn rather
239 // than handing back a client whose child has exited.
240 pool.remove(server);
241 }
242 }
243 let config = self
244 .servers
245 .get(server)
246 .ok_or_else(|| McpError::Protocol(format!("unknown mcp server '{server}'")))?;
247 let config = Self::resolve(config, &self.cwd);
248 let client = Arc::new(
249 McpClient::connect_with_notifications(
250 server,
251 &config,
252 &self.cwd,
253 self.sandbox.as_ref(),
254 self.notifications.clone(),
255 )
256 .await?,
257 );
258 self.clients.lock().await.insert(
259 server.to_string(),
260 Pooled {
261 client: client.clone(),
262 last_used: Instant::now(),
263 },
264 );
265 Ok(client)
266 }
267
268 /// Shut down connections idle for longer than `ttl`, and drop any that
269 /// have died. Called from the reconcile loop; safe to call at any time.
270 /// Returns the names evicted, for the diagnostics surface.
271 pub async fn evict_idle(&self, ttl: Duration) -> Vec<String> {
272 let now = Instant::now();
273 let mut evicted = Vec::new();
274 let mut closing = Vec::new();
275 {
276 let mut pool = self.clients.lock().await;
277 pool.retain(|name, entry| {
278 let expired = now.duration_since(entry.last_used) > ttl;
279 let dead = !entry.client.is_alive();
280 if expired || dead {
281 evicted.push(name.clone());
282 if expired && !dead {
283 closing.push(entry.client.clone());
284 }
285 return false;
286 }
287 true
288 });
289 }
290 // Shut down outside the pool lock: `shutdown` waits on the child.
291 for client in closing {
292 client.shutdown().await;
293 }
294 evicted
295 }
296
297 /// Discover a server's current tool catalog immediately before dispatch.
298 /// MCP tool names are server-defined; validating them here keeps every
299 /// caller behind the same protocol boundary and turns model-invented
300 /// names into actionable errors before `tools/call` is sent.
301 pub async fn call_tool(
302 &self,
303 server: &str,
304 tool: &str,
305 arguments: Value,
306 ) -> Result<String, McpError> {
307 let tools = self.list_live(server).await?;
308 let Some(info) = tools.iter().find(|candidate| candidate.name == tool) else {
309 let available = tools
310 .iter()
311 .map(|candidate| candidate.name.as_str())
312 .collect::<Vec<_>>();
313 return Err(McpError::Protocol(format!(
314 "unknown tool '{tool}' on server '{server}'; discovered tools: {}. Use action \\\"list\\\" and call one of those exact names",
315 if available.is_empty() {
316 "(none)".to_string()
317 } else {
318 available.join(", ")
319 }
320 )));
321 };
322 crate::validate::validate_arguments(&arguments, &info.input_schema)
323 .map_err(McpError::Protocol)?;
324 let client = self.get(server).await?;
325 client.call_tool(tool, arguments).await
326 }
327
328 /// Connect if needed and list, recording what happened either way.
329 async fn list_live(&self, server: &str) -> Result<Vec<McpToolInfo>, McpError> {
330 let listed = async {
331 let client = self.get(server).await?;
332 client.list_tools().await
333 };
334 match listed.await {
335 Ok(tools) => {
336 self.record_tools(server, &tools);
337 Ok(tools)
338 }
339 Err(error) => {
340 if self.backing_off(server).is_none() {
341 self.record_failure(server, &error.to_string());
342 }
343 Err(error)
344 }
345 }
346 }
347
348 /// Remove all resolved MCP environment values from text that can cross
349 /// the MCP process boundary. MCP servers frequently include upstream
350 /// request details in errors, so this applies equally to results and
351 /// failures before either can enter a model transcript or UI timeline.
352 pub fn redact(&self, text: impl AsRef<str>) -> String {
353 let mut values = self
354 .servers
355 .values()
356 .flat_map(|config| config.env.iter().map(|(_, value)| value))
357 .filter(|value| value.len() >= 4)
358 .collect::<Vec<_>>();
359 values.sort_unstable_by_key(|value| std::cmp::Reverse(value.len()));
360 values.dedup();
361 let mut redacted = text.as_ref().to_string();
362 for value in values {
363 redacted = redacted.replace(value, "[REDACTED]");
364 }
365 redacted
366 }
367
368 /// The catalog is also model-visible, so schemas supplied by an MCP
369 /// server must cross the same secret boundary as its text output.
370 pub fn redact_json(&self, value: &Value) -> Value {
371 serde_json::from_str(&self.redact(value.to_string())).unwrap_or(Value::Null)
372 }
373
374 /// Relative commands resolve against the workspace cwd.
375 fn resolve(config: &ServerConfig, cwd: &std::path::Path) -> ServerConfig {
376 let mut c = config.clone();
377 let p = std::path::Path::new(&c.command);
378 if !p.is_absolute()
379 && p.components().count() > 1
380 && let joined = cwd.join(p)
381 && joined.is_file()
382 {
383 c.command = joined.display().to_string();
384 }
385 c
386 }
387
388 pub async fn shutdown_all(&self) {
389 let clients: Vec<_> = self
390 .clients
391 .lock()
392 .await
393 .values()
394 .map(|entry| entry.client.clone())
395 .collect();
396 for c in clients {
397 c.shutdown().await;
398 }
399 }
400}
401
402impl McpManager {
403 /// List one server's catalog under [`LIST_TIMEOUT`] — the demand path
404 /// behind the `mcp` tool's `list`. Connects on demand, records the
405 /// outcome, and returns the redacted catalog.
406 ///
407 /// The timeout wraps connect *and* list, because a server that hangs on
408 /// spawn is the same problem as one that hangs on `tools/list` and the
409 /// caller cannot act differently on the two.
410 pub async fn list_tools(&self, server: &str) -> ListOutcome {
411 match tokio::time::timeout(LIST_TIMEOUT, self.list_live(server)).await {
412 Ok(Ok(tools)) => Ok(tools
413 .iter()
414 .map(|t| McpToolInfo {
415 name: t.name.clone(),
416 description: self.redact(&t.description),
417 input_schema: self.redact_json(&t.input_schema),
418 })
419 .collect()),
420 Ok(Err(e)) => Err(self.redact(e.to_string())),
421 Err(_) => {
422 let reason = format!("did not respond within {}s", LIST_TIMEOUT.as_secs());
423 self.record_failure(server, &reason);
424 Err(reason)
425 }
426 }
427 }
428}
429

1use std::collections::HashMap;
2use std::path::PathBuf;
3use std::sync::Arc;
4use std::time::{Duration, Instant};
5
6use tokio::sync::Mutex;
7
8use serde_json::Value;
9
10use crate::client::{McpClient, McpError, McpToolInfo, NotificationSink, ServerConfig};
11
12/// How long listing a server (connect + `tools/list`) may take before the
13/// attempt counts as a failure. Far below the 60s per-request ceiling: a
14/// server that cannot even list its tools should cost the turn seconds, not
15/// a minute.
16pub const LIST_TIMEOUT: Duration = Duration::from_secs(10);
17
18/// First wait after a failed attempt; doubles per consecutive failure up to
19/// [`FAILURE_BACKOFF_CAP`]. The pool never retries on its own — this only
20/// stops repeated demand from respawning a broken server on every call.
21const FAILURE_BACKOFF_BASE: Duration = Duration::from_secs(5);
22const FAILURE_BACKOFF_CAP: Duration = Duration::from_secs(5 * 60);
23
24/// How long a pooled connection may sit unused before it is shut down. A
25/// process that stays up for weeks must not hold a subprocess for every
26/// server it has ever touched; the next call respawns on demand.
27pub const IDLE_TTL: Duration = Duration::from_secs(15 * 60);
28
29/// One pooled connection plus the bookkeeping the pool needs to decide
30/// whether to keep it.
31struct Pooled {
32 client: Arc<McpClient>,
33 last_used: Instant,
34}
35
36/// The outcome of listing one server's catalog.
37///
38/// A failure is a `Err(reason)`, never a synthesised tool. The old shape
39/// returned a tool literally named `error`, which made a dead server look
40/// like a working one to the prompt, to the alias table, and to the retry
41/// guard that then refused to try again.
42pub type ListOutcome = Result<Vec<McpToolInfo>, String>;
43
44/// What demand has taught the pool about one server. Nothing here is ever
45/// learned by spawning a server for its own sake: a server nobody has asked
46/// for has an empty observation.
47#[derive(Debug, Clone, Default, PartialEq)]
48pub struct ServerObservation {
49 /// The catalog the last live connection returned, redacted and with
50 /// descriptions trimmed. `None` until the server is first used.
51 pub tools: Option<Vec<McpToolInfo>>,
52 /// Why the last attempt failed, until an attempt succeeds.
53 pub failure: Option<String>,
54}
55
56#[derive(Default)]
57struct ServerState {
58 observation: ServerObservation,
59 attempts: u32,
60 failed_at: Option<Instant>,
61}
62
63impl ServerState {
64 fn retry_after(&self, now: Instant) -> Option<Duration> {
65 let failed_at = self.failed_at?;
66 let shift = self.attempts.saturating_sub(1).min(16);
67 let wait = FAILURE_BACKOFF_BASE
68 .saturating_mul(1u32 << shift)
69 .min(FAILURE_BACKOFF_CAP);
70 wait.checked_sub(now.duration_since(failed_at))
71 }
72}
73
74/// The on-demand connection pool: one client per configured server, spawned
75/// only when a call needs it, shared by every session of the owning `Core`,
76/// and shut down after [`IDLE_TTL`] unused. It never connects on its own
77/// initiative — no warm-up, no background probe, no respawn after eviction.
78pub struct McpManager {
79 servers: HashMap<String, ServerConfig>,
80 clients: Mutex<HashMap<String, Pooled>>,
81 state: std::sync::Mutex<HashMap<String, ServerState>>,
82 cwd: PathBuf,
83 sandbox: Option<Arc<dyn vak_tools::sandbox::Sandbox>>,
84 notifications: Option<NotificationSink>,
85 observer: Option<tokio::sync::mpsc::UnboundedSender<String>>,
86}
87
88impl McpManager {
89 pub fn new(servers: HashMap<String, ServerConfig>, cwd: PathBuf) -> Self {
90 Self::new_sandboxed(servers, cwd, None)
91 }
92
93 pub fn new_sandboxed(
94 servers: HashMap<String, ServerConfig>,
95 cwd: PathBuf,
96 sandbox: Option<Arc<dyn vak_tools::sandbox::Sandbox>>,
97 ) -> Self {
98 McpManager {
99 servers,
100 clients: Mutex::new(HashMap::new()),
101 state: std::sync::Mutex::new(HashMap::new()),
102 cwd,
103 sandbox,
104 notifications: None,
105 observer: None,
106 }
107 }
108
109 /// Receive a server's name whenever what the pool knows about it
110 /// changes — a catalog learned or changed, a failure recorded or cleared.
111 pub fn with_observer(mut self, sink: tokio::sync::mpsc::UnboundedSender<String>) -> Self {
112 self.observer = Some(sink);
113 self
114 }
115
116 /// What demand has observed about every server, for the capability
117 /// declarations. Synchronous and free of I/O.
118 pub fn observations(&self) -> HashMap<String, ServerObservation> {
119 self.state
120 .lock()
121 .unwrap_or_else(std::sync::PoisonError::into_inner)
122 .iter()
123 .map(|(name, state)| (name.clone(), state.observation.clone()))
124 .collect()
125 }
126
127 /// Forget a server's catalog, e.g. after it announced
128 /// `notifications/tools/list_changed`. The next use re-lists it.
129 pub fn forget_catalog(&self, server: &str) {
130 let changed = {
131 let mut state = self
132 .state
133 .lock()
134 .unwrap_or_else(std::sync::PoisonError::into_inner);
135 state
136 .get_mut(server)
137 .and_then(|entry| entry.observation.tools.take())
138 .is_some()
139 };
140 if changed {
141 self.announce(server);
142 }
143 }
144
145 fn announce(&self, server: &str) {
146 if let Some(observer) = &self.observer {
147 let _ = observer.send(server.to_string());
148 }
149 }
150
151 fn record_tools(&self, server: &str, tools: &[McpToolInfo]) {
152 let observed: Vec<McpToolInfo> = tools
153 .iter()
154 .map(|tool| McpToolInfo {
155 name: tool.name.clone(),
156 description: self.redact(tool.description.clone()),
157 input_schema: self.redact_json(&tool.input_schema),
158 })
159 .collect();
160 let changed = {
161 let mut state = self
162 .state
163 .lock()
164 .unwrap_or_else(std::sync::PoisonError::into_inner);
165 let entry = state.entry(server.to_string()).or_default();
166 let changed = entry.observation.tools.as_ref() != Some(&observed)
167 || entry.observation.failure.is_some();
168 entry.observation.tools = Some(observed);
169 entry.observation.failure = None;
170 entry.attempts = 0;
171 entry.failed_at = None;
172 changed
173 };
174 if changed {
175 self.announce(server);
176 }
177 }
178
179 fn record_failure(&self, server: &str, reason: &str) {
180 let reason = self.redact(reason);
181 let changed = {
182 let mut state = self
183 .state
184 .lock()
185 .unwrap_or_else(std::sync::PoisonError::into_inner);
186 let entry = state.entry(server.to_string()).or_default();
187 let changed = entry.observation.failure.as_deref() != Some(reason.as_str());
188 entry.observation.failure = Some(reason);
189 entry.attempts = entry.attempts.saturating_add(1);
190 entry.failed_at = Some(Instant::now());
191 changed
192 };
193 if changed {
194 self.announce(server);
195 }
196 }
197
198 /// The failure still inside its backoff window, if any: repeated demand
199 /// gets the recorded reason back instead of a fresh spawn.
200 fn backing_off(&self, server: &str) -> Option<String> {
201 let state = self
202 .state
203 .lock()
204 .unwrap_or_else(std::sync::PoisonError::into_inner);
205 let entry = state.get(server)?;
206 let wait = entry.retry_after(Instant::now())?;
207 let reason = entry.observation.failure.as_deref().unwrap_or("failed");
208 Some(format!(
209 "mcp server '{server}' is unavailable: {reason} (next attempt in {}s)",
210 wait.as_secs().max(1)
211 ))
212 }
213
214 /// Route server-initiated notifications (catalog changes) to `sink`.
215 /// Applies to connections opened after this call.
216 pub fn with_notifications(mut self, sink: NotificationSink) -> Self {
217 self.notifications = Some(sink);
218 self
219 }
220
221 pub fn server_names(&self) -> Vec<String> {
222 let mut names: Vec<String> = self.servers.keys().cloned().collect();
223 names.sort();
224 names
225 }
226
227 async fn get(&self, server: &str) -> Result<Arc<McpClient>, McpError> {
228 if let Some(reason) = self.backing_off(server) {
229 return Err(McpError::Protocol(reason));
230 }
231 {
232 let mut pool = self.clients.lock().await;
233 if let Some(entry) = pool.get_mut(server) {
234 if entry.client.is_alive() {
235 entry.last_used = Instant::now();
236 return Ok(entry.client.clone());
237 }
238 // Dead: drop it and fall through to a fresh spawn rather
239 // than handing back a client whose child has exited.
240 pool.remove(server);
241 }
242 }
243 let config = self
244 .servers
245 .get(server)
246 .ok_or_else(|| McpError::Protocol(format!("unknown mcp server '{server}'")))?;
247 let config = Self::resolve(config, &self.cwd);
248 let client = Arc::new(
249 McpClient::connect_with_notifications(
250 server,
251 &config,
252 &self.cwd,
253 self.sandbox.as_ref(),
254 self.notifications.clone(),
255 )
256 .await?,
257 );
258 self.clients.lock().await.insert(
259 server.to_string(),
260 Pooled {
261 client: client.clone(),
262 last_used: Instant::now(),
263 },
264 );
265 Ok(client)
266 }
267
268 /// Shut down connections idle for longer than `ttl`, and drop any that
269 /// have died. Called from the reconcile loop; safe to call at any time.
270 /// Returns the names evicted, for the diagnostics surface.
271 pub async fn evict_idle(&self, ttl: Duration) -> Vec<String> {
272 let now = Instant::now();
273 let mut evicted = Vec::new();
274 let mut closing = Vec::new();
275 {
276 let mut pool = self.clients.lock().await;
277 pool.retain(|name, entry| {
278 let expired = now.duration_since(entry.last_used) > ttl;
279 let dead = !entry.client.is_alive();
280 if expired || dead {
281 evicted.push(name.clone());
282 if expired && !dead {
283 closing.push(entry.client.clone());
284 }
285 return false;
286 }
287 true
288 });
289 }
290 // Shut down outside the pool lock: `shutdown` waits on the child.
291 for client in closing {
292 client.shutdown().await;
293 }
294 evicted
295 }
296
297 /// Discover a server's current tool catalog immediately before dispatch.
298 /// MCP tool names are server-defined; validating them here keeps every
299 /// caller behind the same protocol boundary and turns model-invented
300 /// names into actionable errors before `tools/call` is sent.
301 pub async fn call_tool(
302 &self,
303 server: &str,
304 tool: &str,
305 arguments: Value,
306 ) -> Result<String, McpError> {
307 let tools = self.list_live(server).await?;
308 let Some(info) = tools.iter().find(|candidate| candidate.name == tool) else {
309 let available = tools
310 .iter()
311 .map(|candidate| candidate.name.as_str())
312 .collect::<Vec<_>>();
313 return Err(McpError::Protocol(format!(
314 "unknown tool '{tool}' on server '{server}'; discovered tools: {}. Use action \\\"list\\\" and call one of those exact names",
315 if available.is_empty() {
316 "(none)".to_string()
317 } else {
318 available.join(", ")
319 }
320 )));
321 };
322 crate::validate::validate_arguments(&arguments, &info.input_schema)
323 .map_err(McpError::Protocol)?;
324 let client = self.get(server).await?;
325 client.call_tool(tool, arguments).await
326 }
327
328 /// Connect if needed and list, recording what happened either way.
329 async fn list_live(&self, server: &str) -> Result<Vec<McpToolInfo>, McpError> {
330 let listed = async {
331 let client = self.get(server).await?;
332 client.list_tools().await
333 };
334 match listed.await {
335 Ok(tools) => {
336 self.record_tools(server, &tools);
337 Ok(tools)
338 }
339 Err(error) => {
340 if self.backing_off(server).is_none() {
341 self.record_failure(server, &error.to_string());
342 }
343 Err(error)
344 }
345 }
346 }
347
348 /// Remove all resolved MCP environment values from text that can cross
349 /// the MCP process boundary. MCP servers frequently include upstream
350 /// request details in errors, so this applies equally to results and
351 /// failures before either can enter a model transcript or UI timeline.
352 pub fn redact(&self, text: impl AsRef<str>) -> String {
353 let mut values = self
354 .servers
355 .values()
356 .flat_map(|config| config.env.iter().map(|(_, value)| value))
357 .filter(|value| value.len() >= 4)
358 .collect::<Vec<_>>();
359 values.sort_unstable_by_key(|value| std::cmp::Reverse(value.len()));
360 values.dedup();
361 let mut redacted = text.as_ref().to_string();
362 for value in values {
363 redacted = redacted.replace(value, "[REDACTED]");
364 }
365 redacted
366 }
367
368 /// The catalog is also model-visible, so schemas supplied by an MCP
369 /// server must cross the same secret boundary as its text output.
370 pub fn redact_json(&self, value: &Value) -> Value {
371 serde_json::from_str(&self.redact(value.to_string())).unwrap_or(Value::Null)
372 }
373
374 /// Relative commands resolve against the workspace cwd.
375 fn resolve(config: &ServerConfig, cwd: &std::path::Path) -> ServerConfig {
376 let mut c = config.clone();
377 let p = std::path::Path::new(&c.command);
378 if !p.is_absolute()
379 && p.components().count() > 1
380 && let joined = cwd.join(p)
381 && joined.is_file()
382 {
383 c.command = joined.display().to_string();
384 }
385 c
386 }
387
388 pub async fn shutdown_all(&self) {
389 let clients: Vec<_> = self
390 .clients
391 .lock()
392 .await
393 .values()
394 .map(|entry| entry.client.clone())
395 .collect();
396 for c in clients {
397 c.shutdown().await;
398 }
399 }
400}
401
402impl McpManager {
403 /// List one server's catalog under [`LIST_TIMEOUT`] — the demand path
404 /// behind the `mcp` tool's `list`. Connects on demand, records the
405 /// outcome, and returns the redacted catalog.
406 ///
407 /// The timeout wraps connect *and* list, because a server that hangs on
408 /// spawn is the same problem as one that hangs on `tools/list` and the
409 /// caller cannot act differently on the two.
410 pub async fn list_tools(&self, server: &str) -> ListOutcome {
411 match tokio::time::timeout(LIST_TIMEOUT, self.list_live(server)).await {
412 Ok(Ok(tools)) => Ok(tools
413 .iter()
414 .map(|t| McpToolInfo {
415 name: t.name.clone(),
416 description: self.redact(&t.description),
417 input_schema: self.redact_json(&t.input_schema),
418 })
419 .collect()),
420 Ok(Err(e)) => Err(self.redact(e.to_string())),
421 Err(_) => {
422 let reason = format!("did not respond within {}s", LIST_TIMEOUT.as_secs());
423 self.record_failure(server, &reason);
424 Err(reason)
425 }
426 }
427 }
428}
429

Indexing the workspace…