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

[codex-mcp] Avoid blocking on optional MCP startup

Browse Source
This commit is contained in:
bromano
2026-06-15 23:24:50 +00:00
parent 6e50b22e55
commit ff17e5d9b3
2 changed files with 182 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

68
codex-rs/codex-mcp/src/connection_manager.rs
View File

@@ -446,41 +446,51 @@ impl McpConnectionManager {
/// Returns all tools with model-visible names normalized.
#[instrument(level = "trace", skip_all, fields(mcp_server_count = self.clients.len()))]
pub async fn list_all_tools(&self) -> Vec<ToolInfo> {
let mut tools = Vec::new();
for (server_name, managed_client) in &self.clients {
let has_cached_tool_info_snapshot = managed_client.cached_tool_info_snapshot.is_some();
let startup_complete = managed_client
.startup_complete
.load(std::sync::atomic::Ordering::Acquire);
trace!(
server_name = %server_name,
has_cached_tool_info_snapshot,
startup_complete,
"waiting for MCP server tools while building tool list"
);
let Some(server_tools) = managed_client
.listed_tools()
let tools: Vec<ToolInfo> =
futures::future::join_all(self.clients.iter().map(|(server_name, managed_client)| {
let required = self.required_servers.binary_search(server_name).is_ok();
let has_cached_tool_info_snapshot =
managed_client.cached_tool_info_snapshot.is_some();
let startup_complete = managed_client.startup_complete.load(Ordering::Acquire);
async move {
trace!(
server_name = %server_name,
required,
has_cached_tool_info_snapshot,
startup_complete,
"waiting for MCP server tools while building tool list"
);
if !required && !startup_complete && !has_cached_tool_info_snapshot {
return None;
}
if required && !startup_complete && has_cached_tool_info_snapshot {
let _ = managed_client.client().await;
}
let server_tools = managed_client.listed_tools().await;
if let Some(server_tools) = &server_tools {
trace!(
server_name = %server_name,
required,
tool_count = server_tools.len(),
"listed MCP server tools while building tool list"
);
}
server_tools
}
.instrument(trace_span!(
"list_tools_for_server",
server_name = %server_name,
required,
has_cached_tool_info_snapshot,
startup_complete
))
.await
else {
continue;
};
trace!(
server_name = %server_name,
tool_count = server_tools.len(),
"listed MCP server tools while building tool list"
);
tools.extend(
server_tools
.into_iter()
.map(|tool| self.with_server_metadata(tool)),
);
}
}))
.await
.into_iter()
.flatten()
.flatten()
.map(|tool| self.with_server_metadata(tool))
.collect();
normalize_tools_for_model_with_prefix(tools, self.prefix_mcp_tool_names)
}

144
codex-rs/codex-mcp/src/connection_manager_tests.rs
View File

@@ -833,6 +833,76 @@ async fn list_all_tools_uses_cached_tool_info_snapshot_while_client_is_pending()
assert_eq!(tool.callable_name, "calendar_create_event");
}
#[tokio::test]
async fn list_all_tools_polls_clients_concurrently() {
let barrier = Arc::new(tokio::sync::Barrier::new(3));
let (first_started_tx, first_started_rx) = tokio::sync::oneshot::channel();
let first_barrier = Arc::clone(&barrier);
let first_pending_client = async move {
let _ = first_started_tx.send(());
first_barrier.wait().await;
Err(StartupOutcomeError::Cancelled)
}
.boxed()
.shared();
let (second_started_tx, second_started_rx) = tokio::sync::oneshot::channel();
let second_barrier = Arc::clone(&barrier);
let second_pending_client = async move {
let _ = second_started_tx.send(());
second_barrier.wait().await;
Err(StartupOutcomeError::Cancelled)
}
.boxed()
.shared();
let approval_policy = Constrained::allow_any(AskForApproval::OnFailure);
let permission_profile = Constrained::allow_any(PermissionProfile::default());
let mut manager = McpConnectionManager::new_uninitialized(
&approval_policy,
&permission_profile,
/*prefix_mcp_tool_names*/ true,
);
manager.required_servers = vec!["first".to_string(), "second".to_string()];
manager.clients.insert(
"first".to_string(),
AsyncManagedClient {
client: first_pending_client,
cached_tool_info_snapshot: None,
cached_server_info: None,
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()),
cancel_token: CancellationToken::new(),
},
);
manager.clients.insert(
"second".to_string(),
AsyncManagedClient {
client: second_pending_client,
cached_tool_info_snapshot: None,
cached_server_info: None,
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()),
cancel_token: CancellationToken::new(),
},
);
let list_task = tokio::spawn(async move { manager.list_all_tools().await });
let (first_started, second_started) =
tokio::time::timeout(Duration::from_secs(1), async move {
tokio::join!(first_started_rx, second_started_rx)
})
.await
.expect("all client futures should be polled before either is released");
first_started.expect("first client future should start");
second_started.expect("second client future should start");
barrier.wait().await;
let tools = tokio::time::timeout(Duration::from_secs(1), list_task)
.await
.expect("tool listing should finish after clients are released")
.expect("tool listing task should not panic");
assert!(tools.is_empty());
}
#[tokio::test]
async fn list_available_server_infos_uses_cache_while_client_is_pending() {
let pending_client = futures::future::pending::<Result<ManagedClient, StartupOutcomeError>>()
@@ -957,7 +1027,7 @@ async fn list_all_tools_applies_legacy_mcp_prefix_by_default() {
}
#[tokio::test]
async fn list_all_tools_blocks_while_client_is_pending_without_cached_tool_info_snapshot() {
async fn list_all_tools_skips_optional_client_pending_without_cached_tool_info_snapshot() {
let pending_client = futures::future::pending::<Result<ManagedClient, StartupOutcomeError>>()
.boxed()
.shared();
@@ -980,6 +1050,75 @@ async fn list_all_tools_blocks_while_client_is_pending_without_cached_tool_info_
},
);
let timeout_result =
tokio::time::timeout(Duration::from_millis(10), manager.list_all_tools()).await;
let tools = timeout_result.expect("optional client tool listing should not block on startup");
assert!(tools.is_empty());
}
#[tokio::test]
async fn list_all_tools_blocks_while_required_client_is_pending_without_cached_tool_info_snapshot()
{
let pending_client = futures::future::pending::<Result<ManagedClient, StartupOutcomeError>>()
.boxed()
.shared();
let approval_policy = Constrained::allow_any(AskForApproval::OnFailure);
let permission_profile = Constrained::allow_any(PermissionProfile::default());
let mut manager = McpConnectionManager::new_uninitialized(
&approval_policy,
&permission_profile,
/*prefix_mcp_tool_names*/ true,
);
manager
.required_servers
.push(CODEX_APPS_MCP_SERVER_NAME.to_string());
manager.clients.insert(
CODEX_APPS_MCP_SERVER_NAME.to_string(),
AsyncManagedClient {
client: pending_client,
cached_tool_info_snapshot: None,
cached_server_info: None,
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()),
cancel_token: CancellationToken::new(),
},
);
let timeout_result =
tokio::time::timeout(Duration::from_millis(10), manager.list_all_tools()).await;
assert!(timeout_result.is_err());
}
#[tokio::test]
async fn list_all_tools_blocks_while_required_client_is_pending_with_cached_tool_info_snapshot() {
let pending_client = futures::future::pending::<Result<ManagedClient, StartupOutcomeError>>()
.boxed()
.shared();
let approval_policy = Constrained::allow_any(AskForApproval::OnFailure);
let permission_profile = Constrained::allow_any(PermissionProfile::default());
let mut manager = McpConnectionManager::new_uninitialized(
&approval_policy,
&permission_profile,
/*prefix_mcp_tool_names*/ true,
);
manager
.required_servers
.push(CODEX_APPS_MCP_SERVER_NAME.to_string());
manager.clients.insert(
CODEX_APPS_MCP_SERVER_NAME.to_string(),
AsyncManagedClient {
client: pending_client,
cached_tool_info_snapshot: Some(vec![create_test_tool(
CODEX_APPS_MCP_SERVER_NAME,
"calendar_create_event",
)]),
cached_server_info: None,
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()),
cancel_token: CancellationToken::new(),
},
);
let timeout_result =
tokio::time::timeout(Duration::from_millis(10), manager.list_all_tools()).await;
assert!(timeout_result.is_err());
@@ -1004,6 +1143,9 @@ async fn shutdown_cancels_pending_tool_listing() {
&permission_profile,
/*prefix_mcp_tool_names*/ true,
);
manager
.required_servers
.push(CODEX_APPS_MCP_SERVER_NAME.to_string());
manager.clients.insert(
CODEX_APPS_MCP_SERVER_NAME.to_string(),
AsyncManagedClient {