Skip to content

[codex-mcp] Avoid blocking on optional MCP startup #28407

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
bromano-oai wants to merge 1 commit into from
Closed

[codex-mcp] Avoid blocking on optional MCP startup #28407

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 39 additions & 29 deletions codex-rs/codex-mcp/src/connection_manager.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}

Expand Down
144 changes: 143 additions & 1 deletion codex-rs/codex-mcp/src/connection_manager_tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>>()
Expand Down Expand Up @@ -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();
Expand All @@ -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());
Expand All @@ -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 {
Expand Down
Loading