chore: consolidate new() and initialize() for McpConnectionManager (#12255)

## Why
`McpConnectionManager` used a two-phase setup (`new()` followed by
`initialize()`), which forced call sites to construct placeholder state
and then mutate it asynchronously. That made MCP startup/refresh flows
harder to follow and easier to misuse, especially around cancellation
token ownership.

## What changed
- Replaced the two-phase initialization flow with a single async
constructor: `McpConnectionManager::new(...) -> (Self,
CancellationToken)`.
- Added `McpConnectionManager::new_uninitialized()` for places that need
an empty manager before async startup begins.
- Added `McpConnectionManager::new_mcp_connection_manager_for_tests()`
for test-only construction.
- Updated MCP startup and refresh call sites in
`codex-rs/core/src/codex.rs` to build a fresh manager via `new(...)`,
swap it in, and update the startup cancellation token consistently.
- Updated MCP snapshot/connector call sites in
`codex-rs/core/src/mcp/mod.rs` and `codex-rs/core/src/connectors.rs` to
use the consolidated constructor.
- Removed the now-obsolete `reset_mcp_startup_cancellation_token()`
helper in favor of explicit token replacement at the call sites.

## Testing
- Not run (refactor-only change; no new behavior was intended).
This commit is contained in:
Michael Bolin
2026-02-19 10:59:51 -08:00
committed by GitHub
Unverified
parent 9719dc502c
commit 7cd2e84026
4 changed files with 99 additions and 75 deletions
+62 -40
View File
@@ -1281,7 +1281,16 @@ impl Session {
.await;
let services = SessionServices {
mcp_connection_manager: Arc::new(RwLock::new(McpConnectionManager::default())),
// Initialize the MCP connection manager with an uninitialized
// instance. It will be replaced with one created via
// McpConnectionManager::new() once all its constructor args are
// available. This also ensures `SessionConfigured` is emitted
// before any MCP-related events. It is reasonable to consider
// changing this to use Option or OnceCell, though the current
// setup is straightforward enough and performs well.
mcp_connection_manager: Arc::new(
RwLock::new(McpConnectionManager::new_uninitialized()),
),
mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()),
unified_exec_manager: UnifiedExecProcessManager::new(
config.background_terminal_max_timeout,
@@ -1385,7 +1394,7 @@ impl Session {
// Start the watcher after SessionConfigured so it cannot emit earlier events.
sess.start_file_watcher_listener();
// Construct sandbox_state before initialize() so it can be sent to each
// Construct sandbox_state before MCP startup so it can be sent to each
// MCP server immediately after it becomes ready (avoiding blocking).
let sandbox_state = SandboxState {
sandbox_policy: session_configuration.sandbox_policy.get().clone(),
@@ -1399,21 +1408,30 @@ impl Session {
.map(|(name, _)| name.clone())
.collect();
required_mcp_servers.sort();
let cancel_token = sess.mcp_startup_cancellation_token().await;
sess.services
.mcp_connection_manager
.write()
.await
.initialize(
&mcp_servers,
config.mcp_oauth_credentials_store_mode,
auth_statuses.clone(),
tx_event.clone(),
cancel_token,
sandbox_state,
)
.await;
{
let mut cancel_guard = sess.services.mcp_startup_cancellation_token.lock().await;
cancel_guard.cancel();
*cancel_guard = CancellationToken::new();
}
let (mcp_connection_manager, cancel_token) = McpConnectionManager::new(
&mcp_servers,
config.mcp_oauth_credentials_store_mode,
auth_statuses.clone(),
tx_event.clone(),
sandbox_state,
)
.await;
{
let mut manager_guard = sess.services.mcp_connection_manager.write().await;
*manager_guard = mcp_connection_manager;
}
{
let mut cancel_guard = sess.services.mcp_startup_cancellation_token.lock().await;
if cancel_guard.is_cancelled() {
cancel_token.cancel();
}
*cancel_guard = cancel_token;
}
if !required_mcp_servers.is_empty() {
let failures = sess
.services
@@ -3032,19 +3050,26 @@ impl Session {
sandbox_cwd: turn_context.cwd.clone(),
use_linux_sandbox_bwrap: turn_context.features.enabled(Feature::UseLinuxSandboxBwrap),
};
let cancel_token = self.reset_mcp_startup_cancellation_token().await;
let mut refreshed_manager = McpConnectionManager::default();
refreshed_manager
.initialize(
&mcp_servers,
store_mode,
auth_statuses,
self.get_tx_event(),
cancel_token,
sandbox_state,
)
.await;
{
let mut cancel_guard = self.services.mcp_startup_cancellation_token.lock().await;
cancel_guard.cancel();
*cancel_guard = CancellationToken::new();
}
let (refreshed_manager, cancel_token) = McpConnectionManager::new(
&mcp_servers,
store_mode,
auth_statuses,
self.get_tx_event(),
sandbox_state,
)
.await;
{
let mut cancel_guard = self.services.mcp_startup_cancellation_token.lock().await;
if cancel_guard.is_cancelled() {
cancel_token.cancel();
}
*cancel_guard = cancel_token;
}
let mut manager = self.services.mcp_connection_manager.write().await;
*manager = refreshed_manager;
@@ -3093,6 +3118,7 @@ impl Session {
.await;
}
#[cfg(test)]
async fn mcp_startup_cancellation_token(&self) -> CancellationToken {
self.services
.mcp_startup_cancellation_token
@@ -3101,14 +3127,6 @@ impl Session {
.clone()
}
async fn reset_mcp_startup_cancellation_token(&self) -> CancellationToken {
let mut guard = self.services.mcp_startup_cancellation_token.lock().await;
guard.cancel();
let cancel_token = CancellationToken::new();
*guard = cancel_token.clone();
cancel_token
}
fn show_raw_agent_reasoning(&self) -> bool {
self.services.show_raw_agent_reasoning
}
@@ -7364,7 +7382,9 @@ mod tests {
let file_watcher = Arc::new(FileWatcher::noop());
let services = SessionServices {
mcp_connection_manager: Arc::new(RwLock::new(McpConnectionManager::default())),
mcp_connection_manager: Arc::new(RwLock::new(
McpConnectionManager::new_mcp_connection_manager_for_tests(),
)),
mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()),
unified_exec_manager: UnifiedExecProcessManager::new(
config.background_terminal_max_timeout,
@@ -7516,7 +7536,9 @@ mod tests {
let file_watcher = Arc::new(FileWatcher::noop());
let services = SessionServices {
mcp_connection_manager: Arc::new(RwLock::new(McpConnectionManager::default())),
mcp_connection_manager: Arc::new(RwLock::new(
McpConnectionManager::new_mcp_connection_manager_for_tests(),
)),
mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()),
unified_exec_manager: UnifiedExecProcessManager::new(
config.background_terminal_max_timeout,