From b940adae8e6d7e291bcf33feddcf875c6c571420 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Thu, 8 May 2025 22:49:15 -0700 Subject: [PATCH 1/2] fix: get responses API working again in Rust (#872) I inadvertently regressed support for the Responses API when adding support for the chat completions API in https://github.com/openai/codex/pull/862. This should get both APIs working again, but the chat completions codepath seems more complex than necessary. I'll try to clean that up shortly, but I want to get things working again ASAP. --- codex-rs/core/src/client.rs | 27 ++++++++++++++++++++++++++- codex-rs/core/src/codex.rs | 3 +-- 2 files changed, 27 insertions(+), 3 deletions(-) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 1b21f6e0c5..5f4f2a1cb8 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -19,6 +19,7 @@ use tracing::debug; use tracing::trace; use tracing::warn; +use crate::chat_completions::AggregateStreamExt; use crate::chat_completions::stream_chat_completions; use crate::client_common::Payload; use crate::client_common::Prompt; @@ -111,7 +112,31 @@ impl ModelClient { match self.provider.wire_api { WireApi::Responses => self.stream_responses(prompt).await, WireApi::Chat => { - stream_chat_completions(prompt, &self.model, &self.client, &self.provider).await + // Create the raw streaming connection first. + let response_stream = + stream_chat_completions(prompt, &self.model, &self.client, &self.provider) + .await?; + + // Wrap it with the aggregation adapter so callers see *only* + // the final assistant message per turn (matching the + // behaviour of the Responses API). + let mut aggregated = response_stream.aggregate(); + + // Bridge the aggregated stream back into a standard + // `ResponseStream` by forwarding events through a channel. + let (tx, rx) = mpsc::channel::>(16); + + tokio::spawn(async move { + use futures::StreamExt; + while let Some(ev) = aggregated.next().await { + // Exit early if receiver hung up. + if tx.send(ev).await.is_err() { + break; + } + } + }); + + Ok(ResponseStream { rx_event: rx }) } } } diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index f68eb73f48..7d056adcd9 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -32,7 +32,6 @@ use tracing::trace; use tracing::warn; use crate::WireApi; -use crate::chat_completions::AggregateStreamExt; use crate::client::ModelClient; use crate::client_common::Prompt; use crate::client_common::ResponseEvent; @@ -864,7 +863,7 @@ async fn try_run_turn( sub_id: &str, prompt: &Prompt, ) -> CodexResult> { - let mut stream = sess.client.clone().stream(prompt).await?.aggregate(); + let mut stream = sess.client.clone().stream(prompt).await?; // Buffer all the incoming messages from the stream first, then execute them. // If we execute a function call in the middle of handling the stream, it can time out. From 009403b02bf8a1f460575b2c995324dc5d6d95d7 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Thu, 8 May 2025 23:16:16 -0700 Subject: [PATCH 2/2] fix: make McpConnectionManager tolerant of MCPs that fail to start --- codex-rs/core/src/codex.rs | 27 ++++++++++++-- codex-rs/core/src/mcp_connection_manager.rs | 39 ++++++++++++++------- 2 files changed, 51 insertions(+), 15 deletions(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 7d056adcd9..bc8b900b38 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -561,15 +561,36 @@ async fn submission_loop( let writable_roots = Mutex::new(get_writable_roots(&cwd)); - let mcp_connection_manager = + let (mcp_connection_manager, failed_clients) = match McpConnectionManager::new(config.mcp_servers.clone()).await { - Ok(mgr) => mgr, + Ok((mgr, failures)) => (mgr, failures), Err(e) => { error!("Failed to create MCP connection manager: {e:#}"); - McpConnectionManager::default() + (McpConnectionManager::default(), Default::default()) } }; + // Surface individual client start-up failures to the user. + if !failed_clients.is_empty() { + for (server_name, err) in failed_clients { + // Log the failure for debugging. + error!("MCP client for '{server_name}' failed to start: {err:#}"); + + // Emit an error event so the front-end can inform the user. + let event = Event { + id: sub.id.clone(), + msg: EventMsg::Error { + message: format!( + "Failed to start MCP server '{server_name}': {err}" + ), + }, + }; + + // Ignore send failures (agent might have died already). + let _ = tx_event.send(event).await; + } + } + // Attempt to create a RolloutRecorder *before* moving the // `instructions` value into the Session struct. let rollout_recorder = match RolloutRecorder::new(instructions.clone()).await { diff --git a/codex-rs/core/src/mcp_connection_manager.rs b/codex-rs/core/src/mcp_connection_manager.rs index 734c351478..e29a0c4ba5 100644 --- a/codex-rs/core/src/mcp_connection_manager.rs +++ b/codex-rs/core/src/mcp_connection_manager.rs @@ -29,6 +29,10 @@ const MCP_TOOL_NAME_DELIMITER: &str = "__OAI_CODEX_MCP__"; /// Timeout for the `tools/list` request. const LIST_TOOLS_TIMEOUT: Duration = Duration::from_secs(10); +/// Map that holds a startup error for every MCP server that could **not** be +/// spawned successfully. +pub type ClientStartErrors = HashMap; + fn fully_qualified_tool_name(server: &str, tool: &str) -> String { format!("{server}{MCP_TOOL_NAME_DELIMITER}{tool}") } @@ -60,40 +64,51 @@ impl McpConnectionManager { /// * `mcp_servers` – Map loaded from the user configuration where *keys* /// are human-readable server identifiers and *values* are the spawn /// instructions. - pub async fn new(mcp_servers: HashMap) -> Result { + /// + /// The function no longer errors out when *individual* MCP servers fail + /// to start. Instead, it returns a tuple `(Self, ClientStartErrors)` where + /// the map stores the error for every server that failed to spawn. + /// Call-sites are expected to inspect the map and surface the failures to + /// the user (e.g. via `EventMsg::Error`). + pub async fn new( + mcp_servers: HashMap, + ) -> Result<(Self, ClientStartErrors)> { // Early exit if no servers are configured. if mcp_servers.is_empty() { - return Ok(Self::default()); + return Ok((Self::default(), ClientStartErrors::default())); } - // Spin up all servers concurrently. + // Launch all configured servers concurrently. let mut join_set = JoinSet::new(); - // Spawn tasks to launch each server. for (server_name, cfg) in mcp_servers { - // TODO: Verify server name: require `^[a-zA-Z0-9_-]+$`? join_set.spawn(async move { let McpServerConfig { command, args, env } = cfg; let client_res = McpClient::new_stdio_client(command, args, env).await; - (server_name, client_res) }); } let mut clients: HashMap> = HashMap::with_capacity(join_set.len()); + let mut errors: ClientStartErrors = HashMap::new(); + while let Some(res) = join_set.join_next().await { - let (server_name, client_res) = res?; + let (server_name, client_res) = res?; // JoinError propagation - let client = client_res - .with_context(|| format!("failed to spawn MCP server `{server_name}`"))?; - - clients.insert(server_name, std::sync::Arc::new(client)); + match client_res { + Ok(client) => { + clients.insert(server_name, std::sync::Arc::new(client)); + } + Err(e) => { + errors.insert(server_name, e.into()); + } + } } let tools = list_all_tools(&clients).await?; - Ok(Self { clients, tools }) + Ok((Self { clients, tools }, errors)) } /// Returns a single map that contains **all** tools. Each key is the