github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-30 00:09:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

Centralize thread MCP connections in McpRuntime (#33889)

Browse Source 
## What changed

- Add a thread-owned `McpRuntime` that atomically publishes connection-set
  replacements and owns shutdown.
- Route `McpResourceClient` snapshots and cache keys through that runtime so
  extension clients observe refreshed MCP servers.
- Keep model-step configuration snapshots separate from ownership of live MCP
  connections.

## Testing

- Add coverage that captures an extension's session resource client, refreshes
  the MCP server configuration, and verifies that the existing client sees the
  newly published server.

GitOrigin-RevId: 351b2e25594b8eb6b91e50ff8d06b7d49768ae1c
This commit is contained in:
jif
2026-07-17 18:43:14 +00:00
committed by copyberry
parent b99cc315fd
commit f24e695470
10 changed files with 203 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/session/handlers.rs
+1 -5
View File
@@ -607,11 +607,7 @@ 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}");
}
sess.services
.latest_mcp_runtime()
.manager_arc()
.shutdown()
.await;
sess.services.mcp_runtime.shutdown().await;
sess.guardian_review_session.shutdown().await;
}
codex-rs/core/src/session/mcp.rs
+3 -1
View File
@@ -185,7 +185,9 @@ impl Session {
current.runtime_context().clone(),
ready_selected_capability_roots,
));
self.services.mcp_runtime.store(Some(Arc::clone(&runtime)));
self.services
.mcp_runtime_snapshot
.store(Some(Arc::clone(&runtime)));
return runtime;
}
self.refresh_mcp_servers_inner(
codex-rs/core/src/session/mod.rs
+1
View File
@@ -73,6 +73,7 @@ use codex_login::CodexAuth;
use codex_login::auth_env_telemetry::collect_auth_env_telemetry;
use codex_mcp::McpConnectionManager;
use codex_mcp::McpResourceClient;
use codex_mcp::McpRuntime;
use codex_mcp::McpRuntimeContext;
use codex_models_manager::manager::RefreshStrategy;
use codex_models_manager::manager::SharedModelsManager;
codex-rs/core/src/session/session.rs
+8 -17
View File
@@ -1036,20 +1036,16 @@ impl Session {
config.analytics_enabled,
)
});
// Keep one stable manager handle for the session so extension resource clients
// automatically observe the manager installed at startup and on later refreshes.
let mcp_connection_manager = Arc::new(arc_swap::ArcSwap::from_pointee(
let mcp_runtime = Arc::new(McpRuntime::new(Arc::new(
McpConnectionManager::new_uninitialized_with_permission_profile(
&config.permissions.approval_policy,
config.permissions.permission_profile(),
config.prefix_mcp_tool_names(),
),
));
)));
let session_extension_data =
codex_extension_api::ExtensionData::new(session_id.to_string());
session_extension_data.insert(McpResourceClient::new(Arc::clone(
&mcp_connection_manager,
)));
session_extension_data.insert(McpResourceClient::new(Arc::clone(&mcp_runtime)));
for contributor in extensions.thread_lifecycle_contributors() {
contributor.on_thread_start(codex_extension_api::ThreadStartInput {
config: config.as_ref(),
@@ -1062,15 +1058,10 @@ impl Session {
}
let services = SessionServices {
// 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,
mcp_runtime: arc_swap::ArcSwapOption::empty(),
// Start with an empty connection set. The initialized set is
// published after SessionConfigured so MCP events follow it.
mcp_runtime,
mcp_runtime_snapshot: arc_swap::ArcSwapOption::empty(),
mcp_projection_lock: Mutex::new(()),
mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()),
unified_exec_manager: UnifiedExecProcessManager::new(
@@ -1257,7 +1248,7 @@ impl Session {
))
.await;
sess.services
.install_mcp_connection_manager(
.install_mcp_runtime(
Arc::new(mcp_projection.config),
mcp_projection.plugins_available,
mcp_runtime_context,
codex-rs/core/src/session/tests.rs
+12 -6
View File
@@ -5305,11 +5305,14 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
/*bundled_skills_enabled*/ true,
));
let network_approval = Arc::new(NetworkApprovalService::default());
let mcp_runtime =
let mcp_runtime_snapshot =
crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(config.as_ref());
let mcp_runtime = Arc::new(codex_mcp::McpRuntime::new(
mcp_runtime_snapshot.manager_arc(),
));
let services = SessionServices {
mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from(mcp_runtime.manager_arc())),
mcp_runtime: arc_swap::ArcSwapOption::from(Some(mcp_runtime)),
mcp_runtime,
mcp_runtime_snapshot: arc_swap::ArcSwapOption::from(Some(mcp_runtime_snapshot)),
mcp_projection_lock: Mutex::new(()),
mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()),
unified_exec_manager: UnifiedExecProcessManager::new(
@@ -7460,11 +7463,14 @@ where
/*bundled_skills_enabled*/ true,
));
let network_approval = Arc::new(NetworkApprovalService::default());
let mcp_runtime =
let mcp_runtime_snapshot =
crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(config.as_ref());
let mcp_runtime = Arc::new(codex_mcp::McpRuntime::new(
mcp_runtime_snapshot.manager_arc(),
));
let services = SessionServices {
mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from(mcp_runtime.manager_arc())),
mcp_runtime: arc_swap::ArcSwapOption::from(Some(mcp_runtime)),
mcp_runtime,
mcp_runtime_snapshot: arc_swap::ArcSwapOption::from(Some(mcp_runtime_snapshot)),
mcp_projection_lock: Mutex::new(()),
mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()),
unified_exec_manager: UnifiedExecProcessManager::new(
codex-rs/core/src/state/service.rs
+15 -17
View File
@@ -34,6 +34,7 @@ use codex_hooks::Hooks;
use codex_login::AuthManager;
use codex_mcp::McpConfig;
use codex_mcp::McpConnectionManager;
use codex_mcp::McpRuntime;
use codex_mcp::McpRuntimeContext;
use codex_models_manager::manager::SharedModelsManager;
use codex_otel::SessionTelemetry;
@@ -48,10 +49,10 @@ use tokio::sync::Mutex;
use tokio_util::sync::CancellationToken;
pub(crate) struct SessionServices {
/// Mirror of the latest manager for extension resource clients that predate runtime snapshots.
pub(crate) mcp_connection_manager: Arc<ArcSwap<McpConnectionManager>>,
/// The latest atomically published MCP config and manager pair.
pub(crate) mcp_runtime: ArcSwapOption<McpRuntimeSnapshot>,
/// The single owner of live MCP connections for this thread.
pub(crate) mcp_runtime: Arc<McpRuntime>,
/// The latest atomically published MCP config and connection snapshot.
pub(crate) mcp_runtime_snapshot: ArcSwapOption<McpRuntimeSnapshot>,
/// Serializes environment-driven runtime rebuilds.
pub(crate) mcp_projection_lock: Mutex<()>,
pub(crate) mcp_startup_cancellation_token: Mutex<CancellationToken>,
@@ -104,22 +105,22 @@ pub(crate) struct SessionServices {
}
impl SessionServices {
/// Installs the manager before validating required servers so startup-time elicitation can
/// resolve through the session's manager while validation waits.
pub(crate) async fn install_mcp_connection_manager(
/// Publishes the initial connections before validating required servers so startup-time
/// elicitation can resolve through the thread runtime while validation waits.
pub(crate) async fn install_mcp_runtime(
&self,
config: Arc<McpConfig>,
plugins_available: bool,
runtime_context: McpRuntimeContext,
ready_selected_capability_roots: Vec<SelectedCapabilityRoot>,
manager: McpConnectionManager,
connections: McpConnectionManager,
) -> Result<()> {
let runtime = self.publish_mcp_runtime(
config,
plugins_available,
runtime_context,
ready_selected_capability_roots,
manager,
connections,
);
runtime.manager().validate_required_servers().await
}
@@ -130,25 +131,22 @@ impl SessionServices {
plugins_available: bool,
runtime_context: McpRuntimeContext,
ready_selected_capability_roots: Vec<SelectedCapabilityRoot>,
manager: McpConnectionManager,
connections: McpConnectionManager,
) -> Arc<McpRuntimeSnapshot> {
let manager = Arc::new(manager);
// Publish the manager for legacy resource clients first. Once the paired snapshot is
// visible, every model-scoped consumer observes this exact manager.
self.mcp_connection_manager.store(Arc::clone(&manager));
let connections = self.mcp_runtime.replace(connections);
let runtime = Arc::new(McpRuntimeSnapshot::new(
config,
plugins_available,
manager,
connections,
runtime_context,
ready_selected_capability_roots,
));
self.mcp_runtime.store(Some(Arc::clone(&runtime)));
self.mcp_runtime_snapshot.store(Some(Arc::clone(&runtime)));
runtime
}
pub(crate) fn latest_mcp_runtime(&self) -> Arc<McpRuntimeSnapshot> {
let Some(runtime) = self.mcp_runtime.load_full() else {
let Some(runtime) = self.mcp_runtime_snapshot.load_full() else {
unreachable!("MCP runtime must be installed before handling requests");
};
runtime