mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
Use latest-wins MCP manager replacement (#27259)
## Summary We originally addressed startup prewarming holding the read side of `RwLock<McpConnectionManager>` by snapshotting tool-list state. Review feedback identified the broader ownership problem: the outer synchronization should only publish or retrieve the current manager, while MCP operations rely on the manager's internal synchronization. A follow-up preserved operation retirement with a separate gate, but further review questioned whether that synchronization was actually required and whether we could support latest-wins replacement instead. This PR now stores the current MCP manager in `ArcSwap`. Each operation uses `load_full()` to obtain an owned `Arc<McpConnectionManager>`, then performs MCP I/O without retaining the publication mechanism. Refresh cancels obsolete startup work, constructs a replacement, and atomically publishes it. New operations see the latest manager, while operations that already loaded the previous manager retain a valid handle. Refresh happens at a turn boundary, so there should be no active user tool calls to drain. Git history supports dropping the outer `RwLock`. It was introduced in `03ffe4d595` on November 17, 2025 for non-blocking MCP startup: the session published an empty manager, startup initialized that same object while holding the write lock, and readers waited for initialization. `7cd2e84026` on February 19, 2026 removed that two-phase initialization in favor of constructing a fresh manager and swapping it in, explicitly noting that `Option` or `OnceCell` could replace the placeholder design. Hot reload later reused the existing lock to publish a replacement, but I found no indication that the lock was introduced to guarantee in-flight tool calls finish before refresh or shutdown. Terminal shutdown remains separate from refresh: it aborts startup prewarming and active tasks before shutting down the current manager, so tool calls may be interrupted and no model WebSocket work continues after shutdown. Focused regression coverage exercises pending tool-list cancellation, deferred refresh, and startup-prewarm shutdown.
This commit is contained in:
committed by
GitHub
Unverified
parent
d2f6d23c6c
commit
41b4fabbb4
@@ -580,6 +580,9 @@ pub async fn set_thread_memory_mode(sess: &Arc<Session>, sub_id: String, mode: T
|
||||
}
|
||||
|
||||
async fn shutdown_session_runtime(sess: &Arc<Session>) {
|
||||
if let Some(startup_prewarm) = sess.take_session_startup_prewarm().await {
|
||||
startup_prewarm.abort().await;
|
||||
}
|
||||
sess.abort_all_tasks(TurnAbortReason::Interrupted).await;
|
||||
let _ = sess.conversation.shutdown().await;
|
||||
sess.services
|
||||
@@ -589,11 +592,11 @@ async fn shutdown_session_runtime(sess: &Arc<Session>) {
|
||||
if let Err(err) = sess.services.code_mode_service.shutdown().await {
|
||||
warn!("failed to shutdown code mode session: {err}");
|
||||
}
|
||||
let mcp_shutdown = {
|
||||
let mut manager = sess.services.mcp_connection_manager.write().await;
|
||||
manager.begin_shutdown()
|
||||
};
|
||||
mcp_shutdown.await;
|
||||
sess.services
|
||||
.mcp_connection_manager
|
||||
.load_full()
|
||||
.shutdown()
|
||||
.await;
|
||||
sess.guardian_review_session.shutdown().await;
|
||||
}
|
||||
|
||||
|
||||
@@ -91,8 +91,7 @@ impl Session {
|
||||
if self
|
||||
.services
|
||||
.mcp_connection_manager
|
||||
.read()
|
||||
.await
|
||||
.load_full()
|
||||
.elicitations_auto_deny()
|
||||
{
|
||||
return McpServerElicitationOutcome {
|
||||
@@ -226,16 +225,11 @@ impl Session {
|
||||
|
||||
self.services
|
||||
.mcp_connection_manager
|
||||
.read()
|
||||
.await
|
||||
.load_full()
|
||||
.resolve_elicitation(server_name, id, response)
|
||||
.await
|
||||
}
|
||||
|
||||
#[expect(
|
||||
clippy::await_holding_invalid_type,
|
||||
reason = "MCP resource calls are serialized through the session-owned manager guard"
|
||||
)]
|
||||
pub async fn list_resources(
|
||||
&self,
|
||||
server: &str,
|
||||
@@ -243,16 +237,11 @@ impl Session {
|
||||
) -> anyhow::Result<ListResourcesResult> {
|
||||
self.services
|
||||
.mcp_connection_manager
|
||||
.read()
|
||||
.await
|
||||
.load_full()
|
||||
.list_resources(server, params)
|
||||
.await
|
||||
}
|
||||
|
||||
#[expect(
|
||||
clippy::await_holding_invalid_type,
|
||||
reason = "MCP resource calls are serialized through the session-owned manager guard"
|
||||
)]
|
||||
pub async fn list_resource_templates(
|
||||
&self,
|
||||
server: &str,
|
||||
@@ -260,16 +249,11 @@ impl Session {
|
||||
) -> anyhow::Result<ListResourceTemplatesResult> {
|
||||
self.services
|
||||
.mcp_connection_manager
|
||||
.read()
|
||||
.await
|
||||
.load_full()
|
||||
.list_resource_templates(server, params)
|
||||
.await
|
||||
}
|
||||
|
||||
#[expect(
|
||||
clippy::await_holding_invalid_type,
|
||||
reason = "MCP resource calls are serialized through the session-owned manager guard"
|
||||
)]
|
||||
pub async fn read_resource(
|
||||
&self,
|
||||
server: &str,
|
||||
@@ -277,16 +261,11 @@ impl Session {
|
||||
) -> anyhow::Result<ReadResourceResult> {
|
||||
self.services
|
||||
.mcp_connection_manager
|
||||
.read()
|
||||
.await
|
||||
.load_full()
|
||||
.read_resource(server, params)
|
||||
.await
|
||||
}
|
||||
|
||||
#[expect(
|
||||
clippy::await_holding_invalid_type,
|
||||
reason = "MCP tool calls are serialized through the session-owned manager guard"
|
||||
)]
|
||||
pub async fn call_tool(
|
||||
&self,
|
||||
server: &str,
|
||||
@@ -296,8 +275,7 @@ impl Session {
|
||||
) -> anyhow::Result<CallToolResult> {
|
||||
self.services
|
||||
.mcp_connection_manager
|
||||
.read()
|
||||
.await
|
||||
.load_full()
|
||||
.call_tool(server, tool, arguments, meta)
|
||||
.await
|
||||
}
|
||||
@@ -366,14 +344,12 @@ impl Session {
|
||||
)
|
||||
.await;
|
||||
{
|
||||
let current_manager = self.services.mcp_connection_manager.read().await;
|
||||
let current_manager = self.services.mcp_connection_manager.load_full();
|
||||
refreshed_manager.set_elicitations_auto_deny(current_manager.elicitations_auto_deny());
|
||||
}
|
||||
let mut old_manager = {
|
||||
let mut manager = self.services.mcp_connection_manager.write().await;
|
||||
std::mem::replace(&mut *manager, refreshed_manager)
|
||||
};
|
||||
old_manager.shutdown().await;
|
||||
self.services
|
||||
.mcp_connection_manager
|
||||
.store(Arc::new(refreshed_manager));
|
||||
}
|
||||
|
||||
pub(crate) async fn refresh_mcp_servers_if_requested(
|
||||
|
||||
@@ -793,7 +793,7 @@ impl Codex {
|
||||
..Default::default()
|
||||
})
|
||||
.await?;
|
||||
let mcp_connection_manager = self.session.services.mcp_connection_manager.read().await;
|
||||
let mcp_connection_manager = self.session.services.mcp_connection_manager.load_full();
|
||||
mcp_connection_manager.set_elicitations_auto_deny(mcp_elicitations_auto_deny);
|
||||
Ok(())
|
||||
}
|
||||
@@ -2748,10 +2748,6 @@ impl Session {
|
||||
}
|
||||
}
|
||||
|
||||
#[expect(
|
||||
clippy::await_holding_invalid_type,
|
||||
reason = "MCP app context rendering reads through the session-owned manager guard"
|
||||
)]
|
||||
pub(crate) async fn build_initial_context(
|
||||
&self,
|
||||
turn_context: &TurnContext,
|
||||
@@ -2844,7 +2840,7 @@ impl Session {
|
||||
}
|
||||
}
|
||||
if turn_context.config.include_apps_instructions && turn_context.apps_enabled() {
|
||||
let mcp_connection_manager = self.services.mcp_connection_manager.read().await;
|
||||
let mcp_connection_manager = self.services.mcp_connection_manager.load_full();
|
||||
let accessible_and_enabled_connectors =
|
||||
connectors::list_accessible_and_enabled_connectors_from_manager(
|
||||
&mcp_connection_manager,
|
||||
|
||||
@@ -981,13 +981,13 @@ impl Session {
|
||||
// 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(
|
||||
mcp_connection_manager: arc_swap::ArcSwap::from_pointee(
|
||||
McpConnectionManager::new_uninitialized_with_permission_profile(
|
||||
&config.permissions.approval_policy,
|
||||
config.permissions.permission_profile(),
|
||||
config.prefix_mcp_tool_names(),
|
||||
),
|
||||
)),
|
||||
),
|
||||
mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()),
|
||||
unified_exec_manager: UnifiedExecProcessManager::new(
|
||||
config.background_terminal_max_timeout,
|
||||
|
||||
@@ -319,8 +319,7 @@ async fn request_mcp_server_elicitation_auto_accepts_when_auto_deny_is_enabled()
|
||||
session
|
||||
.services
|
||||
.mcp_connection_manager
|
||||
.read()
|
||||
.await
|
||||
.load_full()
|
||||
.set_elicitations_auto_deny(/*auto_deny*/ true);
|
||||
|
||||
let requested_schema: McpElicitationSchema = serde_json::from_value(json!({
|
||||
@@ -4816,13 +4815,13 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
|
||||
);
|
||||
|
||||
let services = SessionServices {
|
||||
mcp_connection_manager: Arc::new(RwLock::new(
|
||||
mcp_connection_manager: arc_swap::ArcSwap::from_pointee(
|
||||
McpConnectionManager::new_uninitialized_with_permission_profile(
|
||||
&config.permissions.approval_policy,
|
||||
config.permissions.permission_profile(),
|
||||
config.prefix_mcp_tool_names(),
|
||||
),
|
||||
)),
|
||||
),
|
||||
mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()),
|
||||
unified_exec_manager: UnifiedExecProcessManager::new(
|
||||
config.background_terminal_max_timeout,
|
||||
@@ -6894,13 +6893,13 @@ where
|
||||
);
|
||||
|
||||
let services = SessionServices {
|
||||
mcp_connection_manager: Arc::new(RwLock::new(
|
||||
mcp_connection_manager: arc_swap::ArcSwap::from_pointee(
|
||||
McpConnectionManager::new_uninitialized_with_permission_profile(
|
||||
&config.permissions.approval_policy,
|
||||
config.permissions.permission_profile(),
|
||||
config.prefix_mcp_tool_names(),
|
||||
),
|
||||
)),
|
||||
),
|
||||
mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()),
|
||||
unified_exec_manager: UnifiedExecProcessManager::new(
|
||||
config.background_terminal_max_timeout,
|
||||
@@ -9440,18 +9439,13 @@ async fn abort_review_task_emits_exited_then_aborted_and_records_history() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[expect(
|
||||
clippy::await_holding_invalid_type,
|
||||
reason = "test builds a router from session-owned MCP manager state"
|
||||
)]
|
||||
async fn fatal_tool_error_stops_turn_and_reports_error() {
|
||||
let (session, turn_context, _rx) = make_session_and_context_with_rx().await;
|
||||
let tools = {
|
||||
session
|
||||
.services
|
||||
.mcp_connection_manager
|
||||
.read()
|
||||
.await
|
||||
.load_full()
|
||||
.list_all_tools()
|
||||
.await
|
||||
};
|
||||
|
||||
@@ -432,10 +432,6 @@ async fn run_hooks_and_record_inputs(
|
||||
blocked_input && !accepted_user_input
|
||||
}
|
||||
|
||||
#[expect(
|
||||
clippy::await_holding_invalid_type,
|
||||
reason = "MCP tool listing borrows the read guard across cancellation-aware await"
|
||||
)]
|
||||
#[instrument(level = "trace", skip_all)]
|
||||
async fn build_skills_and_plugins(
|
||||
sess: &Arc<Session>,
|
||||
@@ -473,8 +469,7 @@ async fn build_skills_and_plugins(
|
||||
match sess
|
||||
.services
|
||||
.mcp_connection_manager
|
||||
.read()
|
||||
.await
|
||||
.load_full()
|
||||
.list_all_tools()
|
||||
.or_cancel(cancellation_token)
|
||||
.await
|
||||
@@ -1078,10 +1073,6 @@ async fn run_sampling_request(
|
||||
}
|
||||
}
|
||||
|
||||
#[expect(
|
||||
clippy::await_holding_invalid_type,
|
||||
reason = "tool router construction reads through the session-owned manager guard"
|
||||
)]
|
||||
#[instrument(level = "trace",
|
||||
skip_all,
|
||||
fields(
|
||||
@@ -1095,18 +1086,12 @@ pub(crate) async fn built_tools(
|
||||
turn_context: &TurnContext,
|
||||
cancellation_token: &CancellationToken,
|
||||
) -> CodexResult<Arc<ToolRouter>> {
|
||||
let mcp_connection_manager = sess
|
||||
.services
|
||||
.mcp_connection_manager
|
||||
.read()
|
||||
.instrument(trace_span!("read_mcp_connection_manager"))
|
||||
.await;
|
||||
let mcp_connection_manager = sess.services.mcp_connection_manager.load_full();
|
||||
let has_mcp_servers = mcp_connection_manager.has_servers();
|
||||
let all_mcp_tools = mcp_connection_manager
|
||||
.list_all_tools()
|
||||
.or_cancel(cancellation_token)
|
||||
.await?;
|
||||
drop(mcp_connection_manager);
|
||||
let loaded_plugins = sess
|
||||
.services
|
||||
.plugins_manager
|
||||
|
||||
@@ -724,7 +724,7 @@ impl Session {
|
||||
.unwrap_or_else(|| session_configuration.cwd().clone());
|
||||
let per_turn_config = Self::build_per_turn_config(&session_configuration, cwd.clone());
|
||||
{
|
||||
let mcp_connection_manager = self.services.mcp_connection_manager.read().await;
|
||||
let mcp_connection_manager = self.services.mcp_connection_manager.load_full();
|
||||
mcp_connection_manager.set_approval_policy(&session_configuration.approval_policy);
|
||||
mcp_connection_manager
|
||||
.set_permission_profile(session_configuration.permission_profile());
|
||||
|
||||
Reference in New Issue
Block a user