valknar/llmx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: update McpClient::new_stdio_client() to accept an env (#831)

Browse Source 
Cleans up the signature for `new_stdio_client()` to more closely mirror
how MCP servers are declared in config files (`command`, `args`, `env`).
Also takes a cue from Claude Code where the MCP server is launched with
a restricted `env` so that it only includes "safe" things like `USER`
and `PATH` (see the `create_env_for_mcp_server()` function introduced in
this PR for details) by default, as it is common for developers to have
sensitive API keys present in their environment that should only be
forwarded to the MCP server when the user has explicitly configured it
to do so.

---
[//]: # (BEGIN SAPLING FOOTER)
Stack created with [Sapling](https://sapling-scm.com). Best reviewed
with [ReviewStack](https://reviewstack.dev/openai/codex/pull/831).
* #829
* __->__ #831
This commit is contained in:
Michael Bolin
2025-05-06 11:14:47 -07:00
committed by GitHub
parent 2cf7aeeeb6
commit 5f1b8f707c
4 changed files with 112 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
codex-rs/Cargo.lock generated
View File

@@ -561,7 +561,6 @@ name = "codex-mcp-client"
version = "0.1.0"
dependencies = [
"anyhow",
"codex-core",
"mcp-types",
"pretty_assertions",
"serde",

1
codex-rs/mcp-client/Cargo.toml
View File

@@ -5,7 +5,6 @@ edition = "2021"
[dependencies]
anyhow = "1"
codex-core = { path = "../core", features = ["cli"] }
mcp-types = { path = "../mcp-types" }
serde = { version = "1", features = ["derive"] }
serde_json = "1"

11
codex-rs/mcp-client/src/main.rs
View File

@@ -18,17 +18,20 @@ use mcp_types::ListToolsRequestParams;
#[tokio::main]
async fn main() -> Result<()> {
// Collect command-line arguments excluding the program name itself.
let cmd_args: Vec<String> = std::env::args().skip(1).collect();
let mut args: Vec<String> = std::env::args().skip(1).collect();
if cmd_args.is_empty() || cmd_args[0] == "--help" || cmd_args[0] == "-h" {
if args.is_empty() || args[0] == "--help" || args[0] == "-h" {
eprintln!("Usage: mcp-client <program> [args..]\n\nExample: mcp-client codex-mcp-server");
std::process::exit(1);
}
let original_args = args.clone();
// Spawn the subprocess and connect the client.
let client = McpClient::new_stdio_client(cmd_args.clone())
let program = args.remove(0);
let env = None;
let client = McpClient::new_stdio_client(program, args, env)
.await
.with_context(|| format!("failed to spawn subprocess: {:?}", cmd_args))?;
.with_context(|| format!("failed to spawn subprocess: {original_args:?}"))?;
// Issue `tools/list` request (no params).
let tools = client

140
codex-rs/mcp-client/src/mcp_client.rs
View File

@@ -18,6 +18,8 @@ use std::sync::Arc;
use anyhow::anyhow;
use anyhow::Result;
use mcp_types::CallToolRequest;
use mcp_types::CallToolRequestParams;
use mcp_types::JSONRPCMessage;
use mcp_types::JSONRPCNotification;
use mcp_types::JSONRPCRequest;
@@ -37,6 +39,7 @@ use tokio::process::Command;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio::sync::Mutex;
use tracing::debug;
use tracing::error;
use tracing::info;
use tracing::warn;
@@ -69,40 +72,22 @@ pub struct McpClient {
impl McpClient {
/// Spawn the given command and establish an MCP session over its STDIO.
///
/// `args` follows the Unix convention where the first element is the
/// executable path and the rest are arguments. For example:
///
/// ```no_run
/// # use codex_mcp_client::McpClient;
/// # async fn run() -> anyhow::Result<()> {
/// let client = McpClient::new_stdio_client(vec![
/// "codex-mcp-server".to_string(),
/// ]).await?;
/// # Ok(()) }
/// ```
pub async fn new_stdio_client(args: Vec<String>) -> std::io::Result<Self> {
if args.is_empty() {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"expected at least one element in `args` - the program to spawn",
));
}
let program = &args[0];
let mut command = Command::new(program);
if args.len() > 1 {
command.args(&args[1..]);
}
command.stdin(std::process::Stdio::piped());
command.stdout(std::process::Stdio::piped());
command.stderr(std::process::Stdio::null());
// As noted in the `kill_on_drop` documentation, the Tokio runtime makes
// a "best effort" to reap-after-exit to avoid zombie processes, but it
// is not a guarantee.
command.kill_on_drop(true);
let mut child = command.spawn()?;
pub async fn new_stdio_client(
program: String,
args: Vec<String>,
env: Option<HashMap<String, String>>,
) -> std::io::Result<Self> {
let mut child = Command::new(program)
.args(args)
.envs(create_env_for_mcp_server(env))
.stdin(std::process::Stdio::piped())
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::null())
// As noted in the `kill_on_drop` documentation, the Tokio runtime makes
// a "best effort" to reap-after-exit to avoid zombie processes, but it
// is not a guarantee.
.kill_on_drop(true)
.spawn()?;
let stdin = child.stdin.take().ok_or_else(|| {
std::io::Error::new(std::io::ErrorKind::Other, "failed to capture child stdin")
@@ -122,6 +107,7 @@ impl McpClient {
while let Some(msg) = outgoing_rx.recv().await {
match serde_json::to_string(&msg) {
Ok(json) => {
debug!("MCP message to server: {json}");
if stdin.write_all(json.as_bytes()).await.is_err() {
error!("failed to write message to child stdin");
break;
@@ -149,6 +135,7 @@ impl McpClient {
tokio::spawn(async move {
while let Ok(Some(line)) = lines.next_line().await {
debug!("MCP message from server: {line}");
match serde_json::from_str::<JSONRPCMessage>(&line) {
Ok(JSONRPCMessage::Response(resp)) => {
Self::dispatch_response(resp, &pending).await;
@@ -229,7 +216,7 @@ impl McpClient {
// Send to writer task.
if self.outgoing_tx.send(message).await.is_err() {
return Err(anyhow!(
"failed to send message to writer task – channel closed"
"failed to send message to writer task - channel closed"
));
}
@@ -262,6 +249,17 @@ impl McpClient {
self.send_request::<ListToolsRequest>(params).await
}
/// Convenience wrapper around `tools/call`.
pub async fn call_tool(
&self,
name: String,
arguments: Option<serde_json::Value>,
) -> Result<mcp_types::CallToolResult> {
let params = CallToolRequestParams { name, arguments };
debug!("MCP tool call: {params:?}");
self.send_request::<CallToolRequest>(params).await
}
/// Internal helper: route a JSON-RPC *response* object to the pending map.
async fn dispatch_response(
resp: JSONRPCResponse,
@@ -310,3 +308,75 @@ impl Drop for McpClient {
let _ = self.child.try_wait();
}
}
/// Environment variables that are always included when spawning a new MCP
/// server.
#[rustfmt::skip]
#[cfg(unix)]
const DEFAULT_ENV_VARS: &[&str] = &[
// https://modelcontextprotocol.io/docs/tools/debugging#environment-variables
// states:
//
// > MCP servers inherit only a subset of environment variables automatically,
// > like `USER`, `HOME`, and `PATH`.
//
// But it does not fully enumerate the list. Empirically, when spawning a
// an MCP server via Claude Desktop on macOS, it reports the following
// environment variables:
"HOME",
"LOGNAME",
"PATH",
"SHELL",
"USER",
"__CF_USER_TEXT_ENCODING",
// Additional environment variables Codex chooses to include by default:
"LANG",
"LC_ALL",
"TERM",
"TMPDIR",
"TZ",
];
#[cfg(windows)]
const DEFAULT_ENV_VARS: &[&str] = &[
// TODO: More research is necessary to curate this list.
"PATH",
"PATHEXT",
"USERNAME",
"USERDOMAIN",
"USERPROFILE",
"TEMP",
"TMP",
];
/// `extra_env` comes from the config for an entry in `mcp_servers` in
/// `config.toml`.
fn create_env_for_mcp_server(
extra_env: Option<HashMap<String, String>>,
) -> HashMap<String, String> {
DEFAULT_ENV_VARS
.iter()
.filter_map(|var| match std::env::var(var) {
Ok(value) => Some((var.to_string(), value)),
Err(_) => None,
})
.chain(extra_env.unwrap_or_default())
.collect::<HashMap<_, _>>()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_create_env_for_mcp_server() {
let env_var = "USER";
let env_var_existing_value = std::env::var(env_var).unwrap_or_default();
let env_var_new_value = format!("{env_var_existing_value}-extra");
let extra_env = HashMap::from([(env_var.to_owned(), env_var_new_value.clone())]);
let mcp_server_env = create_env_for_mcp_server(Some(extra_env));
assert!(mcp_server_env.contains_key("PATH"));
assert_eq!(Some(&env_var_new_value), mcp_server_env.get(env_var));
}
}