openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-15 12:08:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Michael Bolin
2025-05-06 10:52:28 -07:00
parent 2cf7aeeeb6
commit b4edb124db
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));
}
}