Vakyartha documentation

1use std::collections::HashMap;
2use std::process::Stdio;
3use std::sync::Arc;
4use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
5use std::time::Duration;
6
7use serde_json::Value;
8use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
9use tokio::process::{Child, ChildStdin, Command};
10use tokio::sync::{Mutex, mpsc, oneshot};
11use vak_tools::sandbox::Sandbox;
12
13/// A server-initiated message that changes what the server offers.
14///
15/// MCP servers announce catalog changes rather than expecting the client to
16/// poll (`notifications/tools/list_changed`). Dropping these at the transport
17/// is what forced the old design to warm a catalog once and never notice a
18/// change; routing them out gives the capability registry a hint channel and
19/// costs one `match`.
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub enum McpNotification {
22 ToolsListChanged,
23 PromptsListChanged,
24 ResourcesListChanged,
25 /// Anything else the server sent. Kept rather than discarded so an
26 /// unknown-but-present signal is visible in diagnostics.
27 Other(String),
28}
29
30impl McpNotification {
31 fn from_method(method: &str) -> Self {
32 match method {
33 "notifications/tools/list_changed" => McpNotification::ToolsListChanged,
34 "notifications/prompts/list_changed" => McpNotification::PromptsListChanged,
35 "notifications/resources/list_changed" => McpNotification::ResourcesListChanged,
36 other => McpNotification::Other(other.to_string()),
37 }
38 }
39
40 /// Whether this invalidates the tool catalog the registry holds.
41 pub fn invalidates_tools(&self) -> bool {
42 matches!(self, McpNotification::ToolsListChanged)
43 }
44}
45
46/// Where a client posts server-initiated notifications, tagged with the
47/// server they came from. Unbounded because the producer is a server we do
48/// not control and blocking its reader task would stall responses too.
49pub type NotificationSink = mpsc::UnboundedSender<(String, McpNotification)>;
50
51#[derive(Debug, thiserror::Error)]
52pub enum McpError {
53 #[error("mcp spawn failed: {0}")]
54 Spawn(String),
55 #[error("mcp protocol error: {0}")]
56 Protocol(String),
57 #[error("mcp server error: {0}")]
58 Server(String),
59 #[error("mcp timed out")]
60 Timeout,
61 #[error("mcp connection closed")]
62 Closed,
63}
64
65/// Deserialized from TOML, so defaults live on the config side; this
66/// struct carries the resolved values.
67#[derive(Debug, Clone)]
68pub struct ServerConfig {
69 pub command: String,
70 pub args: Vec<String>,
71 pub env: Vec<(String, String)>,
72 /// Allow outbound network for this server (e.g. remote-API MCP tools
73 /// like web search). Opt-in per server via privileged config; when
74 /// false the platform sandbox applies as usual.
75 pub network: bool,
76}
77
78use std::path::PathBuf;
79
80type Pending = Arc<Mutex<HashMap<u64, oneshot::Sender<Result<Value, McpError>>>>>;
81
82pub struct McpClient {
83 server_name: String,
84 stdin: Mutex<ChildStdin>,
85 child: Mutex<Child>,
86 pending: Pending,
87 next_id: AtomicU64,
88 /// Set by the reader task when stdout ends, which is the earliest and
89 /// most reliable signal that this connection is dead. Checked before a
90 /// pooled client is handed out, so a caller never dispatches into a
91 /// corpse and gets a timeout instead of an actionable error.
92 closed: Arc<AtomicBool>,
93 /// What the server said it supports in its `initialize` result. Used to
94 /// report whether a stale catalog is the server's fault (no
95 /// `listChanged`, so we must poll) or ours.
96 server_capabilities: Value,
97}
98
99impl McpClient {
100 pub async fn connect(
101 server_name: &str,
102 config: &ServerConfig,
103 cwd: &std::path::Path,
104 sandbox: Option<&Arc<dyn Sandbox>>,
105 ) -> Result<Self, McpError> {
106 Self::connect_with_notifications(server_name, config, cwd, sandbox, None).await
107 }
108
109 pub async fn connect_with_notifications(
110 server_name: &str,
111 config: &ServerConfig,
112 cwd: &std::path::Path,
113 sandbox: Option<&Arc<dyn Sandbox>>,
114 notifications: Option<NotificationSink>,
115 ) -> Result<Self, McpError> {
116 // A network-egress server is a deliberate trust decision from
117 // privileged config; the OS command wrapper would deny its sockets,
118 // so it spawns directly (env still scrubbed to the explicit set).
119 let command_path = resolve_command(&config.command);
120 let mut cmd = if let (Some(sandbox), false) = (sandbox, config.network) {
121 let command = std::iter::once(command_path.to_string_lossy().to_string())
122 .chain(config.args.iter().cloned())
123 .map(|part| shell_quote(&part))
124 .collect::<Vec<_>>()
125 .join(" ");
126 let mut cmd = Command::new(vak_tools::bash::POSIX_SHELL);
127 cmd.arg("-c").arg(sandbox.wrap(&command));
128 cmd
129 } else {
130 let mut cmd = Command::new(&command_path);
131 cmd.args(&config.args);
132 cmd
133 };
134 cmd.current_dir(cwd)
135 .env_clear()
136 .kill_on_drop(true)
137 .stdin(Stdio::piped())
138 .stdout(Stdio::piped())
139 .stderr(Stdio::null());
140 for (k, v) in &config.env {
141 cmd.env(k, v);
142 }
143 // Operational basics every runtime needs (PATH for interpreters,
144 // HOME/TMPDIR for package caches). Secrets never ride along: the
145 // environment was cleared above and recipients are explicit.
146 // Prepend the server executable's own directory to PATH: under
147 // service managers PATH is minimal, and interpreters resolved via
148 // `#!/usr/bin/env` (npx→node) need the same prefix that worked for
149 // the command itself.
150 let mut path_parts: Vec<PathBuf> = Vec::new();
151 if let Some(dir) = command_path.parent() {
152 path_parts.push(dir.to_path_buf());
153 }
154 for toolchain_path in vak_config::paths::canonical_toolchain_paths() {
155 if !path_parts.contains(&toolchain_path) {
156 path_parts.push(toolchain_path);
157 }
158 }
159 if let Some(path) = std::env::var_os("PATH") {
160 for p in std::env::split_paths(&path).filter(|p| !p.as_os_str().is_empty()) {
161 if !path_parts.contains(&p) {
162 path_parts.push(p);
163 }
164 }
165 }
166 cmd.env(
167 "PATH",
168 std::env::join_paths(&path_parts).unwrap_or_default(),
169 );
170 for var in ["HOME", "TMPDIR"] {
171 if let Some(v) = std::env::var_os(var) {
172 cmd.env(var, v);
173 }
174 }
175 if std::env::var_os("VAK_MCP_DEBUG").is_some() {
176 cmd.stderr(Stdio::inherit());
177 }
178 vak_tools::bash::isolate_process_group(&mut cmd);
179
180 let mut child = cmd.spawn().map_err(|e| McpError::Spawn(e.to_string()))?;
181 let stdin = child
182 .stdin
183 .take()
184 .ok_or_else(|| McpError::Spawn("no stdin".into()))?;
185 let stdout = child
186 .stdout
187 .take()
188 .ok_or_else(|| McpError::Spawn("no stdout".into()))?;
189
190 let pending: Pending = Arc::new(Mutex::new(HashMap::new()));
191 let closed = Arc::new(AtomicBool::new(false));
192
193 // Reader: route responses by id and notifications to the sink.
194 // Server-initiated *requests* (those carry both an id and a method)
195 // are still unanswered — v1 does not serve sampling/roots — but a
196 // notification is a catalog-change signal the registry needs, so it
197 // is forwarded rather than dropped.
198 {
199 let pending = pending.clone();
200 let closed = closed.clone();
201 let server = server_name.to_string();
202 tokio::spawn(async move {
203 let mut lines = BufReader::new(stdout).lines();
204 loop {
205 match lines.next_line().await {
206 Ok(Some(line)) if line.trim().is_empty() => continue,
207 Ok(Some(line)) => {
208 let Ok(v) = serde_json::from_str::<Value>(&line) else {
209 continue;
210 };
211 let method = v.get("method").and_then(|m| m.as_str());
212 let Some(id) = v.get("id").and_then(|i| i.as_u64()) else {
213 // No id: a notification. Forward it.
214 if let (Some(method), Some(sink)) = (method, notifications.as_ref())
215 {
216 let _ = sink.send((
217 server.clone(),
218 McpNotification::from_method(method),
219 ));
220 }
221 continue;
222 };
223 if method.is_some() {
224 // id + method: a server-initiated request.
225 // Not served in v1; never matched to pending.
226 continue;
227 }
228 let responder = pending.lock().await.remove(&id);
229 if let Some(tx) = responder {
230 let result = match v.get("error") {
231 Some(err) => Err(McpError::Server(
232 err.get("message")
233 .and_then(|m| m.as_str())
234 .unwrap_or("unknown")
235 .to_string(),
236 )),
237 None => Ok(v.get("result").cloned().unwrap_or(Value::Null)),
238 };
239 let _ = tx.send(result);
240 }
241 }
242 _ => {
243 // stdout ended: this connection is dead. Mark it
244 // before waking waiters so a racing `get()` sees
245 // the flag rather than handing out this client.
246 closed.store(true, Ordering::Release);
247 let mut map = pending.lock().await;
248 for (_, tx) in map.drain() {
249 let _ = tx.send(Err(McpError::Closed));
250 }
251 return;
252 }
253 }
254 }
255 });
256 }
257
258 let mut client = McpClient {
259 server_name: server_name.to_string(),
260 stdin: Mutex::new(stdin),
261 child: Mutex::new(child),
262 pending,
263 next_id: AtomicU64::new(1),
264 closed,
265 server_capabilities: Value::Null,
266 };
267
268 let initialized = client
269 .request(
270 "initialize",
271 serde_json::json!({
272 "protocolVersion": "2025-06-18",
273 "capabilities": {},
274 "clientInfo": {"name": "vak", "version": env!("CARGO_PKG_VERSION")},
275 }),
276 )
277 .await?;
278 client.server_capabilities = initialized
279 .get("capabilities")
280 .cloned()
281 .unwrap_or(Value::Null);
282 client
283 .notify("notifications/initialized", serde_json::json!({}))
284 .await;
285 Ok(client)
286 }
287
288 pub fn server_name(&self) -> &str {
289 &self.server_name
290 }
291
292 /// False once stdout has ended or the child has exited. Cheap and
293 /// non-blocking: the pool checks this before reusing a client so a dead
294 /// server is replaced rather than dispatched into.
295 pub fn is_alive(&self) -> bool {
296 if self.closed.load(Ordering::Acquire) {
297 return false;
298 }
299 // `try_lock` because this runs on the hot path of every dispatch and
300 // a contended child lock means someone is mid-shutdown anyway.
301 match self.child.try_lock() {
302 Ok(mut child) => !matches!(child.try_wait(), Ok(Some(_))),
303 Err(_) => true,
304 }
305 }
306
307 /// Whether the server promised to announce tool-catalog changes. When
308 /// false the registry must fall back to periodic re-probing for this
309 /// server rather than trusting that silence means unchanged.
310 pub fn announces_tool_changes(&self) -> bool {
311 self.server_capabilities
312 .get("tools")
313 .and_then(|t| t.get("listChanged"))
314 .and_then(|l| l.as_bool())
315 .unwrap_or(false)
316 }
317
318 pub async fn shutdown(&self) {
319 self.cancel_pending().await;
320 let mut child = self.child.lock().await;
321 vak_tools::bash::kill_process_group(&child.id());
322 let _ = child.wait().await;
323 }
324
325 async fn cancel_pending(&self) {
326 let mut map = self.pending.lock().await;
327 for (_, tx) in map.drain() {
328 let _ = tx.send(Err(McpError::Closed));
329 }
330 }
331
332 async fn send_raw(&self, msg: &Value) -> Result<(), McpError> {
333 let mut line = serde_json::to_string(msg).map_err(|e| McpError::Protocol(e.to_string()))?;
334 line.push('\n');
335 let mut stdin = self.stdin.lock().await;
336 stdin
337 .write_all(line.as_bytes())
338 .await
339 .map_err(|_| McpError::Closed)?;
340 stdin.flush().await.map_err(|_| McpError::Closed)
341 }
342
343 async fn notify(&self, method: &str, params: Value) {
344 let msg = serde_json::json!({"jsonrpc": "2.0", "method": method, "params": params});
345 let _ = self.send_raw(&msg).await;
346 }
347
348 async fn request(&self, method: &str, params: Value) -> Result<Value, McpError> {
349 let id = self.next_id.fetch_add(1, Ordering::Relaxed);
350 let msg = serde_json::json!({
351 "jsonrpc": "2.0",
352 "id": id,
353 "method": method,
354 "params": params,
355 });
356
357 let (tx, rx) = oneshot::channel();
358 self.pending.lock().await.insert(id, tx);
359
360 if let Err(e) = self.send_raw(&msg).await {
361 self.pending.lock().await.remove(&id);
362 return Err(e);
363 }
364
365 match tokio::time::timeout(Duration::from_secs(60), rx).await {
366 Ok(Ok(result)) => result,
367 Ok(Err(_)) => Err(McpError::Closed),
368 Err(_) => {
369 self.pending.lock().await.remove(&id);
370 Err(McpError::Timeout)
371 }
372 }
373 }
374
375 pub async fn list_tools(&self) -> Result<Vec<McpToolInfo>, McpError> {
376 let result = self.request("tools/list", serde_json::json!({})).await?;
377 let Some(tools) = result.get("tools").and_then(|t| t.as_array()) else {
378 return Ok(Vec::new());
379 };
380 Ok(tools
381 .iter()
382 .map(|t| McpToolInfo {
383 name: t
384 .get("name")
385 .and_then(|n| n.as_str())
386 .unwrap_or_default()
387 .to_string(),
388 description: t
389 .get("description")
390 .and_then(|d| d.as_str())
391 .unwrap_or_default()
392 .to_string(),
393 input_schema: t
394 .get("inputSchema")
395 .cloned()
396 .unwrap_or_else(|| serde_json::json!({"type": "object"})),
397 })
398 .collect())
399 }
400
401 pub async fn call_tool(&self, tool: &str, arguments: Value) -> Result<String, McpError> {
402 let result = self
403 .request(
404 "tools/call",
405 serde_json::json!({"name": tool, "arguments": arguments}),
406 )
407 .await?;
408
409 let is_error = result
410 .get("isError")
411 .and_then(|e| e.as_bool())
412 .unwrap_or(false);
413 let text = text_tool_result(&result)?;
414
415 if is_error {
416 Err(McpError::Server(if text.is_empty() {
417 "tool reported an error".into()
418 } else {
419 text
420 }))
421 } else {
422 Ok(text)
423 }
424 }
425}
426
427/// The current model tool-result contract is text. Never report success after
428/// silently discarding an MCP image, audio, or resource block: that would
429/// make the model reason over an incomplete result while the ledger says the
430/// call succeeded. A richer result type can replace this boundary later.
431fn text_tool_result(result: &Value) -> Result<String, McpError> {
432 let blocks = result
433 .get("content")
434 .and_then(Value::as_array)
435 .ok_or_else(|| McpError::Protocol("tools/call returned no content array".into()))?;
436 let mut parts = Vec::with_capacity(blocks.len());
437 for block in blocks {
438 if block.get("type").and_then(Value::as_str) != Some("text") {
439 let kind = block
440 .get("type")
441 .and_then(Value::as_str)
442 .unwrap_or("unknown");
443 return Err(McpError::Protocol(format!(
444 "tools/call returned unsupported {kind} content; this tool boundary supports text only"
445 )));
446 }
447 let text = block
448 .get("text")
449 .and_then(Value::as_str)
450 .ok_or_else(|| McpError::Protocol("tools/call text block has no text".into()))?;
451 parts.push(text);
452 }
453 Ok(parts.join("\n"))
454}
455
456fn resolve_command(command: &str) -> PathBuf {
457 let path = PathBuf::from(command);
458 if path.is_absolute() || command.contains(std::path::MAIN_SEPARATOR) {
459 return path;
460 }
461 if let Some(paths) = std::env::var_os("PATH") {
462 for dir in std::env::split_paths(&paths) {
463 let candidate = dir.join(command);
464 if candidate.is_file() {
465 return candidate;
466 }
467 }
468 }
469 for dir in vak_config::paths::canonical_toolchain_paths() {
470 let candidate = dir.join(command);
471 if candidate.is_file() {
472 return candidate;
473 }
474 }
475 path
476}
477
478fn shell_quote(value: &str) -> String {
479 let mut out = String::with_capacity(value.len() + 2);
480 out.push('\'');
481 for character in value.chars() {
482 if character == '\'' {
483 out.push_str("'\\''");
484 } else {
485 out.push(character);
486 }
487 }
488 out.push('\'');
489 out
490}
491
492#[derive(Debug, Clone, PartialEq)]
493pub struct McpToolInfo {
494 pub name: String,
495 pub description: String,
496 pub input_schema: Value,
497}
498
499#[cfg(test)]
500mod tests {
501 use super::*;
502
503 #[test]
504 fn text_result_keeps_every_text_block() {
505 let result = serde_json::json!({"content": [
506 {"type": "text", "text": "first"},
507 {"type": "text", "text": "second"}
508 ]});
509 assert_eq!(
510 text_tool_result(&result).ok().as_deref(),
511 Some("first\nsecond")
512 );
513 }
514
515 #[test]
516 fn non_text_result_fails_instead_of_disappearing() {
517 for result in [
518 serde_json::json!({"content": [{"type": "image", "data": "..."}]}),
519 serde_json::json!({"content": [
520 {"type": "text", "text": "caption"},
521 {"type": "resource", "resource": {"uri": "file:///report"}}
522 ]}),
523 ] {
524 assert!(matches!(
525 text_tool_result(&result),
526 Err(McpError::Protocol(_))
527 ));
528 }
529 }
530}
531

1use std::collections::HashMap;
2use std::process::Stdio;
3use std::sync::Arc;
4use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
5use std::time::Duration;
6
7use serde_json::Value;
8use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
9use tokio::process::{Child, ChildStdin, Command};
10use tokio::sync::{Mutex, mpsc, oneshot};
11use vak_tools::sandbox::Sandbox;
12
13/// A server-initiated message that changes what the server offers.
14///
15/// MCP servers announce catalog changes rather than expecting the client to
16/// poll (`notifications/tools/list_changed`). Dropping these at the transport
17/// is what forced the old design to warm a catalog once and never notice a
18/// change; routing them out gives the capability registry a hint channel and
19/// costs one `match`.
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub enum McpNotification {
22 ToolsListChanged,
23 PromptsListChanged,
24 ResourcesListChanged,
25 /// Anything else the server sent. Kept rather than discarded so an
26 /// unknown-but-present signal is visible in diagnostics.
27 Other(String),
28}
29
30impl McpNotification {
31 fn from_method(method: &str) -> Self {
32 match method {
33 "notifications/tools/list_changed" => McpNotification::ToolsListChanged,
34 "notifications/prompts/list_changed" => McpNotification::PromptsListChanged,
35 "notifications/resources/list_changed" => McpNotification::ResourcesListChanged,
36 other => McpNotification::Other(other.to_string()),
37 }
38 }
39
40 /// Whether this invalidates the tool catalog the registry holds.
41 pub fn invalidates_tools(&self) -> bool {
42 matches!(self, McpNotification::ToolsListChanged)
43 }
44}
45
46/// Where a client posts server-initiated notifications, tagged with the
47/// server they came from. Unbounded because the producer is a server we do
48/// not control and blocking its reader task would stall responses too.
49pub type NotificationSink = mpsc::UnboundedSender<(String, McpNotification)>;
50
51#[derive(Debug, thiserror::Error)]
52pub enum McpError {
53 #[error("mcp spawn failed: {0}")]
54 Spawn(String),
55 #[error("mcp protocol error: {0}")]
56 Protocol(String),
57 #[error("mcp server error: {0}")]
58 Server(String),
59 #[error("mcp timed out")]
60 Timeout,
61 #[error("mcp connection closed")]
62 Closed,
63}
64
65/// Deserialized from TOML, so defaults live on the config side; this
66/// struct carries the resolved values.
67#[derive(Debug, Clone)]
68pub struct ServerConfig {
69 pub command: String,
70 pub args: Vec<String>,
71 pub env: Vec<(String, String)>,
72 /// Allow outbound network for this server (e.g. remote-API MCP tools
73 /// like web search). Opt-in per server via privileged config; when
74 /// false the platform sandbox applies as usual.
75 pub network: bool,
76}
77
78use std::path::PathBuf;
79
80type Pending = Arc<Mutex<HashMap<u64, oneshot::Sender<Result<Value, McpError>>>>>;
81
82pub struct McpClient {
83 server_name: String,
84 stdin: Mutex<ChildStdin>,
85 child: Mutex<Child>,
86 pending: Pending,
87 next_id: AtomicU64,
88 /// Set by the reader task when stdout ends, which is the earliest and
89 /// most reliable signal that this connection is dead. Checked before a
90 /// pooled client is handed out, so a caller never dispatches into a
91 /// corpse and gets a timeout instead of an actionable error.
92 closed: Arc<AtomicBool>,
93 /// What the server said it supports in its `initialize` result. Used to
94 /// report whether a stale catalog is the server's fault (no
95 /// `listChanged`, so we must poll) or ours.
96 server_capabilities: Value,
97}
98
99impl McpClient {
100 pub async fn connect(
101 server_name: &str,
102 config: &ServerConfig,
103 cwd: &std::path::Path,
104 sandbox: Option<&Arc<dyn Sandbox>>,
105 ) -> Result<Self, McpError> {
106 Self::connect_with_notifications(server_name, config, cwd, sandbox, None).await
107 }
108
109 pub async fn connect_with_notifications(
110 server_name: &str,
111 config: &ServerConfig,
112 cwd: &std::path::Path,
113 sandbox: Option<&Arc<dyn Sandbox>>,
114 notifications: Option<NotificationSink>,
115 ) -> Result<Self, McpError> {
116 // A network-egress server is a deliberate trust decision from
117 // privileged config; the OS command wrapper would deny its sockets,
118 // so it spawns directly (env still scrubbed to the explicit set).
119 let command_path = resolve_command(&config.command);
120 let mut cmd = if let (Some(sandbox), false) = (sandbox, config.network) {
121 let command = std::iter::once(command_path.to_string_lossy().to_string())
122 .chain(config.args.iter().cloned())
123 .map(|part| shell_quote(&part))
124 .collect::<Vec<_>>()
125 .join(" ");
126 let mut cmd = Command::new(vak_tools::bash::POSIX_SHELL);
127 cmd.arg("-c").arg(sandbox.wrap(&command));
128 cmd
129 } else {
130 let mut cmd = Command::new(&command_path);
131 cmd.args(&config.args);
132 cmd
133 };
134 cmd.current_dir(cwd)
135 .env_clear()
136 .kill_on_drop(true)
137 .stdin(Stdio::piped())
138 .stdout(Stdio::piped())
139 .stderr(Stdio::null());
140 for (k, v) in &config.env {
141 cmd.env(k, v);
142 }
143 // Operational basics every runtime needs (PATH for interpreters,
144 // HOME/TMPDIR for package caches). Secrets never ride along: the
145 // environment was cleared above and recipients are explicit.
146 // Prepend the server executable's own directory to PATH: under
147 // service managers PATH is minimal, and interpreters resolved via
148 // `#!/usr/bin/env` (npx→node) need the same prefix that worked for
149 // the command itself.
150 let mut path_parts: Vec<PathBuf> = Vec::new();
151 if let Some(dir) = command_path.parent() {
152 path_parts.push(dir.to_path_buf());
153 }
154 for toolchain_path in vak_config::paths::canonical_toolchain_paths() {
155 if !path_parts.contains(&toolchain_path) {
156 path_parts.push(toolchain_path);
157 }
158 }
159 if let Some(path) = std::env::var_os("PATH") {
160 for p in std::env::split_paths(&path).filter(|p| !p.as_os_str().is_empty()) {
161 if !path_parts.contains(&p) {
162 path_parts.push(p);
163 }
164 }
165 }
166 cmd.env(
167 "PATH",
168 std::env::join_paths(&path_parts).unwrap_or_default(),
169 );
170 for var in ["HOME", "TMPDIR"] {
171 if let Some(v) = std::env::var_os(var) {
172 cmd.env(var, v);
173 }
174 }
175 if std::env::var_os("VAK_MCP_DEBUG").is_some() {
176 cmd.stderr(Stdio::inherit());
177 }
178 vak_tools::bash::isolate_process_group(&mut cmd);
179
180 let mut child = cmd.spawn().map_err(|e| McpError::Spawn(e.to_string()))?;
181 let stdin = child
182 .stdin
183 .take()
184 .ok_or_else(|| McpError::Spawn("no stdin".into()))?;
185 let stdout = child
186 .stdout
187 .take()
188 .ok_or_else(|| McpError::Spawn("no stdout".into()))?;
189
190 let pending: Pending = Arc::new(Mutex::new(HashMap::new()));
191 let closed = Arc::new(AtomicBool::new(false));
192
193 // Reader: route responses by id and notifications to the sink.
194 // Server-initiated *requests* (those carry both an id and a method)
195 // are still unanswered — v1 does not serve sampling/roots — but a
196 // notification is a catalog-change signal the registry needs, so it
197 // is forwarded rather than dropped.
198 {
199 let pending = pending.clone();
200 let closed = closed.clone();
201 let server = server_name.to_string();
202 tokio::spawn(async move {
203 let mut lines = BufReader::new(stdout).lines();
204 loop {
205 match lines.next_line().await {
206 Ok(Some(line)) if line.trim().is_empty() => continue,
207 Ok(Some(line)) => {
208 let Ok(v) = serde_json::from_str::<Value>(&line) else {
209 continue;
210 };
211 let method = v.get("method").and_then(|m| m.as_str());
212 let Some(id) = v.get("id").and_then(|i| i.as_u64()) else {
213 // No id: a notification. Forward it.
214 if let (Some(method), Some(sink)) = (method, notifications.as_ref())
215 {
216 let _ = sink.send((
217 server.clone(),
218 McpNotification::from_method(method),
219 ));
220 }
221 continue;
222 };
223 if method.is_some() {
224 // id + method: a server-initiated request.
225 // Not served in v1; never matched to pending.
226 continue;
227 }
228 let responder = pending.lock().await.remove(&id);
229 if let Some(tx) = responder {
230 let result = match v.get("error") {
231 Some(err) => Err(McpError::Server(
232 err.get("message")
233 .and_then(|m| m.as_str())
234 .unwrap_or("unknown")
235 .to_string(),
236 )),
237 None => Ok(v.get("result").cloned().unwrap_or(Value::Null)),
238 };
239 let _ = tx.send(result);
240 }
241 }
242 _ => {
243 // stdout ended: this connection is dead. Mark it
244 // before waking waiters so a racing `get()` sees
245 // the flag rather than handing out this client.
246 closed.store(true, Ordering::Release);
247 let mut map = pending.lock().await;
248 for (_, tx) in map.drain() {
249 let _ = tx.send(Err(McpError::Closed));
250 }
251 return;
252 }
253 }
254 }
255 });
256 }
257
258 let mut client = McpClient {
259 server_name: server_name.to_string(),
260 stdin: Mutex::new(stdin),
261 child: Mutex::new(child),
262 pending,
263 next_id: AtomicU64::new(1),
264 closed,
265 server_capabilities: Value::Null,
266 };
267
268 let initialized = client
269 .request(
270 "initialize",
271 serde_json::json!({
272 "protocolVersion": "2025-06-18",
273 "capabilities": {},
274 "clientInfo": {"name": "vak", "version": env!("CARGO_PKG_VERSION")},
275 }),
276 )
277 .await?;
278 client.server_capabilities = initialized
279 .get("capabilities")
280 .cloned()
281 .unwrap_or(Value::Null);
282 client
283 .notify("notifications/initialized", serde_json::json!({}))
284 .await;
285 Ok(client)
286 }
287
288 pub fn server_name(&self) -> &str {
289 &self.server_name
290 }
291
292 /// False once stdout has ended or the child has exited. Cheap and
293 /// non-blocking: the pool checks this before reusing a client so a dead
294 /// server is replaced rather than dispatched into.
295 pub fn is_alive(&self) -> bool {
296 if self.closed.load(Ordering::Acquire) {
297 return false;
298 }
299 // `try_lock` because this runs on the hot path of every dispatch and
300 // a contended child lock means someone is mid-shutdown anyway.
301 match self.child.try_lock() {
302 Ok(mut child) => !matches!(child.try_wait(), Ok(Some(_))),
303 Err(_) => true,
304 }
305 }
306
307 /// Whether the server promised to announce tool-catalog changes. When
308 /// false the registry must fall back to periodic re-probing for this
309 /// server rather than trusting that silence means unchanged.
310 pub fn announces_tool_changes(&self) -> bool {
311 self.server_capabilities
312 .get("tools")
313 .and_then(|t| t.get("listChanged"))
314 .and_then(|l| l.as_bool())
315 .unwrap_or(false)
316 }
317
318 pub async fn shutdown(&self) {
319 self.cancel_pending().await;
320 let mut child = self.child.lock().await;
321 vak_tools::bash::kill_process_group(&child.id());
322 let _ = child.wait().await;
323 }
324
325 async fn cancel_pending(&self) {
326 let mut map = self.pending.lock().await;
327 for (_, tx) in map.drain() {
328 let _ = tx.send(Err(McpError::Closed));
329 }
330 }
331
332 async fn send_raw(&self, msg: &Value) -> Result<(), McpError> {
333 let mut line = serde_json::to_string(msg).map_err(|e| McpError::Protocol(e.to_string()))?;
334 line.push('\n');
335 let mut stdin = self.stdin.lock().await;
336 stdin
337 .write_all(line.as_bytes())
338 .await
339 .map_err(|_| McpError::Closed)?;
340 stdin.flush().await.map_err(|_| McpError::Closed)
341 }
342
343 async fn notify(&self, method: &str, params: Value) {
344 let msg = serde_json::json!({"jsonrpc": "2.0", "method": method, "params": params});
345 let _ = self.send_raw(&msg).await;
346 }
347
348 async fn request(&self, method: &str, params: Value) -> Result<Value, McpError> {
349 let id = self.next_id.fetch_add(1, Ordering::Relaxed);
350 let msg = serde_json::json!({
351 "jsonrpc": "2.0",
352 "id": id,
353 "method": method,
354 "params": params,
355 });
356
357 let (tx, rx) = oneshot::channel();
358 self.pending.lock().await.insert(id, tx);
359
360 if let Err(e) = self.send_raw(&msg).await {
361 self.pending.lock().await.remove(&id);
362 return Err(e);
363 }
364
365 match tokio::time::timeout(Duration::from_secs(60), rx).await {
366 Ok(Ok(result)) => result,
367 Ok(Err(_)) => Err(McpError::Closed),
368 Err(_) => {
369 self.pending.lock().await.remove(&id);
370 Err(McpError::Timeout)
371 }
372 }
373 }
374
375 pub async fn list_tools(&self) -> Result<Vec<McpToolInfo>, McpError> {
376 let result = self.request("tools/list", serde_json::json!({})).await?;
377 let Some(tools) = result.get("tools").and_then(|t| t.as_array()) else {
378 return Ok(Vec::new());
379 };
380 Ok(tools
381 .iter()
382 .map(|t| McpToolInfo {
383 name: t
384 .get("name")
385 .and_then(|n| n.as_str())
386 .unwrap_or_default()
387 .to_string(),
388 description: t
389 .get("description")
390 .and_then(|d| d.as_str())
391 .unwrap_or_default()
392 .to_string(),
393 input_schema: t
394 .get("inputSchema")
395 .cloned()
396 .unwrap_or_else(|| serde_json::json!({"type": "object"})),
397 })
398 .collect())
399 }
400
401 pub async fn call_tool(&self, tool: &str, arguments: Value) -> Result<String, McpError> {
402 let result = self
403 .request(
404 "tools/call",
405 serde_json::json!({"name": tool, "arguments": arguments}),
406 )
407 .await?;
408
409 let is_error = result
410 .get("isError")
411 .and_then(|e| e.as_bool())
412 .unwrap_or(false);
413 let text = text_tool_result(&result)?;
414
415 if is_error {
416 Err(McpError::Server(if text.is_empty() {
417 "tool reported an error".into()
418 } else {
419 text
420 }))
421 } else {
422 Ok(text)
423 }
424 }
425}
426
427/// The current model tool-result contract is text. Never report success after
428/// silently discarding an MCP image, audio, or resource block: that would
429/// make the model reason over an incomplete result while the ledger says the
430/// call succeeded. A richer result type can replace this boundary later.
431fn text_tool_result(result: &Value) -> Result<String, McpError> {
432 let blocks = result
433 .get("content")
434 .and_then(Value::as_array)
435 .ok_or_else(|| McpError::Protocol("tools/call returned no content array".into()))?;
436 let mut parts = Vec::with_capacity(blocks.len());
437 for block in blocks {
438 if block.get("type").and_then(Value::as_str) != Some("text") {
439 let kind = block
440 .get("type")
441 .and_then(Value::as_str)
442 .unwrap_or("unknown");
443 return Err(McpError::Protocol(format!(
444 "tools/call returned unsupported {kind} content; this tool boundary supports text only"
445 )));
446 }
447 let text = block
448 .get("text")
449 .and_then(Value::as_str)
450 .ok_or_else(|| McpError::Protocol("tools/call text block has no text".into()))?;
451 parts.push(text);
452 }
453 Ok(parts.join("\n"))
454}
455
456fn resolve_command(command: &str) -> PathBuf {
457 let path = PathBuf::from(command);
458 if path.is_absolute() || command.contains(std::path::MAIN_SEPARATOR) {
459 return path;
460 }
461 if let Some(paths) = std::env::var_os("PATH") {
462 for dir in std::env::split_paths(&paths) {
463 let candidate = dir.join(command);
464 if candidate.is_file() {
465 return candidate;
466 }
467 }
468 }
469 for dir in vak_config::paths::canonical_toolchain_paths() {
470 let candidate = dir.join(command);
471 if candidate.is_file() {
472 return candidate;
473 }
474 }
475 path
476}
477
478fn shell_quote(value: &str) -> String {
479 let mut out = String::with_capacity(value.len() + 2);
480 out.push('\'');
481 for character in value.chars() {
482 if character == '\'' {
483 out.push_str("'\\''");
484 } else {
485 out.push(character);
486 }
487 }
488 out.push('\'');
489 out
490}
491
492#[derive(Debug, Clone, PartialEq)]
493pub struct McpToolInfo {
494 pub name: String,
495 pub description: String,
496 pub input_schema: Value,
497}
498
499#[cfg(test)]
500mod tests {
501 use super::*;
502
503 #[test]
504 fn text_result_keeps_every_text_block() {
505 let result = serde_json::json!({"content": [
506 {"type": "text", "text": "first"},
507 {"type": "text", "text": "second"}
508 ]});
509 assert_eq!(
510 text_tool_result(&result).ok().as_deref(),
511 Some("first\nsecond")
512 );
513 }
514
515 #[test]
516 fn non_text_result_fails_instead_of_disappearing() {
517 for result in [
518 serde_json::json!({"content": [{"type": "image", "data": "..."}]}),
519 serde_json::json!({"content": [
520 {"type": "text", "text": "caption"},
521 {"type": "resource", "resource": {"uri": "file:///report"}}
522 ]}),
523 ] {
524 assert!(matches!(
525 text_tool_result(&result),
526 Err(McpError::Protocol(_))
527 ));
528 }
529 }
530}
531

Indexing the workspace…