deepstrike_sdk/runtime/
mcp_proxy_plane.rs

1use std::collections::{HashMap, HashSet};
2use std::pin::Pin;
3use std::sync::Arc;
4
5use async_stream::try_stream;
6use deepstrike_core::types::message::{ToolCall, ToolSchema};
7use futures::StreamExt;
8use futures::stream::Stream;
9use serde_json::Value;
10use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
11use tokio::sync::Mutex;
12
13use crate::run_event::RunEvent;
14use crate::runtime::credential_vault::CredentialVault;
15use crate::runtime::execution_plane::{ExecutionPlane, LocalExecutionPlane, RunContext};
16use crate::tools::RegisteredTool;
17use crate::{Error, Result};
18
19// ── Server config ─────────────────────────────────────────────────────────────
20
21pub struct McpServerConfig {
22    /// Executable to run (e.g. "npx", "python3", "/usr/local/bin/my-mcp-server").
23    pub command: String,
24    pub args: Vec<String>,
25    /// Vault keys injected as env vars into the subprocess. Never exposed to the model.
26    pub credential_keys: Vec<String>,
27    /// Additional static env vars forwarded to the subprocess.
28    pub env: HashMap<String, String>,
29}
30
31// ── Internal MCP connection ───────────────────────────────────────────────────
32
33struct Inner {
34    stdin: tokio::process::ChildStdin,
35    reader: BufReader<tokio::process::ChildStdout>,
36    next_id: u64,
37}
38
39struct McpConnection {
40    server_name: String,
41    inner: Mutex<Inner>,
42    child: Mutex<Option<tokio::process::Child>>,
43    schemas: Vec<ToolSchema>,
44}
45
46async fn do_request(
47    inner: &mut Inner,
48    server_name: &str,
49    method: &str,
50    params: Option<Value>,
51) -> Result<Value> {
52    let id = inner.next_id;
53    inner.next_id += 1;
54
55    let mut msg = serde_json::json!({ "jsonrpc": "2.0", "method": method, "id": id });
56    if let Some(p) = params {
57        msg["params"] = p;
58    }
59    let line = serde_json::to_string(&msg).map_err(|e| Error::Other(e.to_string()))? + "\n";
60
61    inner
62        .stdin
63        .write_all(line.as_bytes())
64        .await
65        .map_err(|e| Error::Other(e.to_string()))?;
66    inner
67        .stdin
68        .flush()
69        .await
70        .map_err(|e| Error::Other(e.to_string()))?;
71
72    loop {
73        let mut buf = String::new();
74        let n = inner
75            .reader
76            .read_line(&mut buf)
77            .await
78            .map_err(|e| Error::Other(e.to_string()))?;
79        if n == 0 {
80            return Err(Error::Other(format!(
81                "MCP server '{server_name}' disconnected"
82            )));
83        }
84        let trimmed = buf.trim();
85        if trimmed.is_empty() {
86            continue;
87        }
88        let Ok(resp) = serde_json::from_str::<Value>(trimmed) else {
89            continue;
90        };
91        // Skip notifications (no numeric id)
92        let Some(resp_id) = resp.get("id").and_then(Value::as_u64) else {
93            continue;
94        };
95        if resp_id != id {
96            continue;
97        }
98        if let Some(err) = resp.get("error") {
99            return Err(Error::Other(format!("MCP({server_name}) error: {err}")));
100        }
101        return Ok(resp.get("result").cloned().unwrap_or(Value::Null));
102    }
103}
104
105async fn do_notify(inner: &mut Inner, method: &str) {
106    let msg = serde_json::json!({ "jsonrpc": "2.0", "method": method });
107    if let Ok(s) = serde_json::to_string(&msg) {
108        inner.stdin.write_all((s + "\n").as_bytes()).await.ok();
109        inner.stdin.flush().await.ok();
110    }
111}
112
113impl McpConnection {
114    async fn start(
115        server_name: &str,
116        config: &McpServerConfig,
117        vault: &dyn CredentialVault,
118    ) -> Result<Arc<Self>> {
119        let mut env: HashMap<String, String> = std::env::vars().collect();
120        for (k, v) in &config.env {
121            env.insert(k.clone(), v.clone());
122        }
123        for key in &config.credential_keys {
124            if let Some(val) = vault.get(key).await {
125                env.insert(key.clone(), val);
126            }
127        }
128
129        let mut child = tokio::process::Command::new(&config.command)
130            .args(&config.args)
131            .envs(&env)
132            .stdin(std::process::Stdio::piped())
133            .stdout(std::process::Stdio::piped())
134            .stderr(std::process::Stdio::inherit())
135            .spawn()
136            .map_err(|e| {
137                Error::Other(format!("failed to spawn MCP server '{server_name}': {e}"))
138            })?;
139
140        let stdin = child
141            .stdin
142            .take()
143            .ok_or_else(|| Error::Other(format!("MCP server '{server_name}': no stdin handle")))?;
144        let stdout = child
145            .stdout
146            .take()
147            .ok_or_else(|| Error::Other(format!("MCP server '{server_name}': no stdout handle")))?;
148
149        let mut inner = Inner {
150            stdin,
151            reader: BufReader::new(stdout),
152            next_id: 1,
153        };
154
155        // MCP handshake
156        do_request(
157            &mut inner,
158            server_name,
159            "initialize",
160            Some(serde_json::json!({
161                "protocolVersion": "2024-11-05",
162                "capabilities": { "tools": {} },
163                "clientInfo": { "name": "deepstrike", "version": "0.1.0" },
164            })),
165        )
166        .await?;
167        do_notify(&mut inner, "notifications/initialized").await;
168
169        // Discover tool schemas
170        let list = do_request(&mut inner, server_name, "tools/list", None).await?;
171        let raw_tools = list
172            .get("tools")
173            .and_then(Value::as_array)
174            .cloned()
175            .unwrap_or_default();
176
177        let mut schemas = Vec::new();
178        for t in &raw_tools {
179            let name = t.get("name").and_then(Value::as_str).unwrap_or("");
180            if name.is_empty() {
181                continue;
182            }
183            let description = t.get("description").and_then(Value::as_str).unwrap_or(name);
184            let parameters = t
185                .get("inputSchema")
186                .cloned()
187                .unwrap_or_else(|| serde_json::json!({ "type": "object", "properties": {} }));
188            schemas.push(ToolSchema {
189                name: name.into(),
190                description: description.to_string(),
191                parameters,
192            });
193        }
194
195        Ok(Arc::new(McpConnection {
196            server_name: server_name.to_string(),
197            inner: Mutex::new(inner),
198            child: Mutex::new(Some(child)),
199            schemas,
200        }))
201    }
202
203    async fn execute(&self, call: &ToolCall) -> (String, bool) {
204        let params = serde_json::json!({
205            "name": call.name.as_str(),
206            "arguments": call.arguments,
207        });
208        let mut inner = self.inner.lock().await;
209        match do_request(&mut inner, &self.server_name, "tools/call", Some(params)).await {
210            Ok(result) => {
211                let text = result
212                    .get("content")
213                    .and_then(Value::as_array)
214                    .map(|arr| {
215                        arr.iter()
216                            .filter(|c| c.get("type").and_then(Value::as_str) == Some("text"))
217                            .filter_map(|c| c.get("text").and_then(Value::as_str))
218                            .collect::<Vec<_>>()
219                            .join("\n")
220                    })
221                    .unwrap_or_default();
222                let is_error = result
223                    .get("isError")
224                    .and_then(Value::as_bool)
225                    .unwrap_or(false);
226                let output = if text.is_empty() {
227                    serde_json::to_string(&result).unwrap_or_default()
228                } else {
229                    text
230                };
231                (output, is_error)
232            }
233            Err(e) => (e.to_string(), true),
234        }
235    }
236
237    async fn stop(&self) {
238        if let Some(mut child) = self.child.lock().await.take() {
239            child.kill().await.ok();
240            child.wait().await.ok();
241        }
242    }
243}
244
245// ── Public plane ──────────────────────────────────────────────────────────────
246
247/// `ExecutionPlane` that proxies tool calls to MCP servers over JSON-RPC 2.0 (stdio).
248///
249/// Credentials live in a `CredentialVault` and are injected into each server's subprocess
250/// environment — the model never sees the credential values.
251///
252/// Usage:
253/// ```ignore
254/// let mut plane = McpProxyPlane::new(Arc::new(EnvCredentialVault));
255/// plane.add_server("brave", McpServerConfig {
256///     command: "npx".into(),
257///     args: vec!["-y".into(), "@modelcontextprotocol/server-brave-search".into()],
258///     credential_keys: vec!["BRAVE_API_KEY".into()],
259///     env: Default::default(),
260/// });
261/// plane.connect().await?;
262/// // ... use with RuntimeRunner ...
263/// plane.disconnect().await;
264/// ```
265pub struct McpProxyPlane {
266    server_configs: HashMap<String, McpServerConfig>,
267    vault: Arc<dyn CredentialVault>,
268    connections: Vec<Arc<McpConnection>>,
269    tool_to_conn: HashMap<String, Arc<McpConnection>>,
270    local: LocalExecutionPlane,
271    local_names: HashSet<String>,
272}
273
274impl McpProxyPlane {
275    pub fn new(vault: Arc<dyn CredentialVault>) -> Self {
276        Self {
277            server_configs: HashMap::new(),
278            vault,
279            connections: Vec::new(),
280            tool_to_conn: HashMap::new(),
281            local: LocalExecutionPlane::new(),
282            local_names: HashSet::new(),
283        }
284    }
285
286    pub fn add_server(&mut self, name: impl Into<String>, config: McpServerConfig) -> &mut Self {
287        self.server_configs.insert(name.into(), config);
288        self
289    }
290
291    /// Start all configured MCP server processes and discover their tool schemas.
292    pub async fn connect(&mut self) -> Result<()> {
293        for (name, config) in &self.server_configs {
294            let conn = McpConnection::start(name, config, self.vault.as_ref()).await?;
295            for schema in &conn.schemas {
296                self.tool_to_conn
297                    .insert(schema.name.to_string(), Arc::clone(&conn));
298            }
299            self.connections.push(conn);
300        }
301        Ok(())
302    }
303
304    /// Gracefully stop all MCP server processes.
305    pub async fn disconnect(&mut self) {
306        for conn in &self.connections {
307            conn.stop().await;
308        }
309        self.connections.clear();
310        self.tool_to_conn.clear();
311    }
312
313    /// Register an in-process tool. Local tools take priority over MCP tools of the same name.
314    pub fn register(&mut self, tool: RegisteredTool) -> &mut Self {
315        self.local_names.insert(tool.schema.name.to_string());
316        self.local.register(tool);
317        self
318    }
319
320    pub fn unregister(&mut self, name: &str) -> &mut Self {
321        self.local_names.remove(name);
322        self.local.unregister(name);
323        self
324    }
325}
326
327impl ExecutionPlane for McpProxyPlane {
328    fn schemas(&self) -> Vec<ToolSchema> {
329        let mut result = self.local.schemas();
330        for conn in &self.connections {
331            result.extend(conn.schemas.iter().cloned());
332        }
333        result
334    }
335
336    fn execute_all<'a>(
337        &'a self,
338        calls: &'a [ToolCall],
339        ctx: RunContext<'a>,
340    ) -> Pin<Box<dyn Stream<Item = Result<RunEvent>> + Send + 'a>> {
341        Box::pin(try_stream! {
342            let RunContext {
343                agent_id, skill_dir, dream_store, knowledge_source, governance, on_tool_suspend, on_permission_request,
344            } = ctx;
345
346            let local_calls: Vec<_> = calls
347                .iter()
348                .filter(|c| self.local_names.contains(c.name.as_str()))
349                .cloned()
350                .collect();
351            let mcp_calls: Vec<_> = calls
352                .iter()
353                .filter(|c| !self.local_names.contains(c.name.as_str()))
354                .cloned()
355                .collect();
356
357            if !local_calls.is_empty() {
358                let local_ctx = RunContext {
359                    agent_id,
360                    skill_dir,
361                    dream_store,
362                    knowledge_source,
363                    governance: governance.clone(),
364                    on_tool_suspend: on_tool_suspend.clone(),
365                    on_permission_request: on_permission_request.clone(),
366                };
367                let mut s = self.local.execute_all(&local_calls, local_ctx);
368                while let Some(evt) = s.next().await {
369                    yield evt?;
370                }
371            }
372
373            // Group MCP calls by connection; concurrent across servers, sequential within.
374            let mut groups: HashMap<usize, (Arc<McpConnection>, Vec<ToolCall>)> = HashMap::new();
375            let mut unknown: Vec<ToolCall> = Vec::new();
376            for call in &mcp_calls {
377                if let Some(conn) = self.tool_to_conn.get(call.name.as_str()) {
378                    let key = Arc::as_ptr(conn) as usize;
379                    let entry = groups.entry(key).or_insert_with(|| (Arc::clone(conn), Vec::new()));
380                    entry.1.push(call.clone());
381                } else {
382                    unknown.push(call.clone());
383                }
384            }
385
386            for call in &unknown {
387                yield RunEvent::ToolResult {
388                    call_id: call.id.to_string(),
389                    content: format!("unknown MCP tool: {}", call.name),
390                    is_error: true,
391                    is_fatal: false,
392                    error_kind: None,
393                };
394            }
395
396            let futs: Vec<_> = groups
397                .into_values()
398                .map(|(conn, group_calls)| async move {
399                    let mut results = Vec::new();
400                    for call in &group_calls {
401                        let (output, is_error) = conn.execute(call).await;
402                        results.push((call.id.to_string(), output, is_error));
403                    }
404                    results
405                })
406                .collect();
407
408            let all_results = futures::future::join_all(futs).await;
409            for group in all_results {
410                for (call_id, content, is_error) in group {
411                    yield RunEvent::ToolResult {
412                        call_id,
413                        content,
414                        is_error,
415                        is_fatal: false,
416                        error_kind: None,
417                    };
418                }
419            }
420        })
421    }
422}
deepstrike_sdk/runtime/mcp_proxy_plane.rs

deepstrike_sdk/runtime/
mcp_proxy_plane.rs
