feat(core): pool stdio MCP backends in thread manager

This commit is contained in:
Joe Gershenson
2026-03-20 11:50:47 -07:00
parent 7591ac9184
commit 05c88b6639
10 changed files with 1498 additions and 731 deletions
+7
View File
@@ -8,6 +8,7 @@ use crate::error::CodexErr;
use crate::error::Result as CodexResult;
use crate::find_archived_thread_path_by_id_str;
use crate::find_thread_path_by_id_str;
use crate::mcp_connection_manager::SharedMcpBackendPool;
use crate::rollout::RolloutRecorder;
use crate::session_prefix::format_subagent_context_line;
use crate::session_prefix::format_subagent_notification_message;
@@ -90,6 +91,12 @@ impl AgentControl {
}
}
pub(crate) fn shared_mcp_backend_pool(&self) -> Option<Arc<SharedMcpBackendPool>> {
self.manager
.upgrade()
.map(|state| state.shared_mcp_backend_pool())
}
/// Spawn a new agent thread and submit the initial prompt.
pub(crate) async fn spawn_agent(
&self,
+154 -46
View File
@@ -220,6 +220,7 @@ use crate::mcp::auth::compute_auth_statuses;
use crate::mcp::maybe_prompt_and_install_mcp_dependencies;
use crate::mcp::with_codex_apps_mcp;
use crate::mcp_connection_manager::McpConnectionManager;
use crate::mcp_connection_manager::SharedMcpBackendAcquireMode;
use crate::mcp_connection_manager::codex_apps_tools_cache_key;
use crate::mcp_connection_manager::filter_non_codex_apps_mcp_tools_only;
use crate::memories;
@@ -1249,6 +1250,23 @@ impl Session {
per_turn_config
}
fn sandbox_state_for_session_configuration(
session_configuration: &SessionConfiguration,
) -> SandboxState {
SandboxState {
sandbox_policy: session_configuration.sandbox_policy.get().clone(),
codex_linux_sandbox_exe: session_configuration
.original_config_do_not_use
.codex_linux_sandbox_exe
.clone(),
sandbox_cwd: session_configuration.cwd.clone(),
use_legacy_landlock: session_configuration
.original_config_do_not_use
.features
.use_legacy_landlock(),
}
}
pub(crate) async fn codex_home(&self) -> PathBuf {
let state = self.state.lock().await;
state.session_configuration.codex_home().clone()
@@ -1800,6 +1818,7 @@ impl Session {
&config.permissions.approval_policy,
))),
mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()),
shared_mcp_backend_pool: agent_control.shared_mcp_backend_pool(),
unified_exec_manager: UnifiedExecProcessManager::new(
config.background_terminal_max_timeout,
),
@@ -1924,17 +1943,40 @@ impl Session {
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(),
&session_configuration.approval_policy,
tx_event.clone(),
sandbox_state,
config.codex_home.clone(),
codex_apps_tools_cache_key(auth),
tool_plugin_provenance,
)
let (mcp_connection_manager, cancel_token) = async {
match sess.services.shared_mcp_backend_pool.as_ref() {
Some(shared_mcp_backend_pool) => {
McpConnectionManager::new_with_pool(
shared_mcp_backend_pool.as_ref(),
SharedMcpBackendAcquireMode::ReuseExisting,
&mcp_servers,
config.mcp_oauth_credentials_store_mode,
auth_statuses.clone(),
&session_configuration.approval_policy,
tx_event.clone(),
sandbox_state,
config.codex_home.clone(),
codex_apps_tools_cache_key(auth),
tool_plugin_provenance,
)
.await
}
None => {
McpConnectionManager::new(
&mcp_servers,
config.mcp_oauth_credentials_store_mode,
auth_statuses.clone(),
&session_configuration.approval_policy,
tx_event.clone(),
sandbox_state,
config.codex_home.clone(),
codex_apps_tools_cache_key(auth),
tool_plugin_provenance,
)
.await
}
}
}
.instrument(info_span!(
"session_init.mcp_manager_init",
otel.name = "session_init.mcp_manager_init",
@@ -2308,23 +2350,29 @@ impl Session {
) -> ConstraintResult<Arc<TurnContext>> {
let (
session_configuration,
sandbox_policy_changed,
sandbox_state,
sandbox_state_changed,
previous_cwd,
codex_home,
session_source,
) = {
let mut state = self.state.lock().await;
match state.session_configuration.clone().apply(&updates) {
let previous_session_configuration = state.session_configuration.clone();
match previous_session_configuration.apply(&updates) {
Ok(next) => {
let previous_cwd = state.session_configuration.cwd.clone();
let sandbox_policy_changed =
state.session_configuration.sandbox_policy != next.sandbox_policy;
let previous_cwd = previous_session_configuration.cwd.clone();
let previous_sandbox_state = Self::sandbox_state_for_session_configuration(
&previous_session_configuration,
);
let sandbox_state = Self::sandbox_state_for_session_configuration(&next);
let sandbox_state_changed = previous_sandbox_state != sandbox_state;
let codex_home = next.codex_home.clone();
let session_source = next.session_source.clone();
state.session_configuration = next.clone();
(
next,
sandbox_policy_changed,
sandbox_state,
sandbox_state_changed,
previous_cwd,
codex_home,
session_source,
@@ -2357,7 +2405,8 @@ impl Session {
sub_id,
session_configuration,
updates.final_output_json_schema,
sandbox_policy_changed,
sandbox_state,
sandbox_state_changed,
)
.await)
}
@@ -2367,23 +2416,55 @@ impl Session {
sub_id: String,
session_configuration: SessionConfiguration,
final_output_json_schema: Option<Option<Value>>,
sandbox_policy_changed: bool,
sandbox_state: SandboxState,
sandbox_state_changed: bool,
) -> Arc<TurnContext> {
let per_turn_config = Self::build_per_turn_config(&session_configuration);
self.services
.mcp_connection_manager
.read()
.await
.set_approval_policy(&session_configuration.approval_policy);
if sandbox_policy_changed {
let sandbox_state = SandboxState {
sandbox_policy: per_turn_config.permissions.sandbox_policy.get().clone(),
codex_linux_sandbox_exe: per_turn_config.codex_linux_sandbox_exe.clone(),
sandbox_cwd: per_turn_config.cwd.clone(),
use_legacy_landlock: per_turn_config.features.use_legacy_landlock(),
};
if let Err(e) = self
if sandbox_state_changed {
if let Some(shared_mcp_backend_pool) = self.services.shared_mcp_backend_pool.as_ref() {
let auth = self.services.auth_manager.auth().await;
let config = session_configuration.original_config_do_not_use.clone();
let mcp_servers = self
.services
.mcp_manager
.effective_servers(config.as_ref(), auth.as_ref());
let auth_statuses = compute_auth_statuses(
mcp_servers.iter(),
config.mcp_oauth_credentials_store_mode,
)
.await;
let tool_plugin_provenance = self
.services
.mcp_manager
.tool_plugin_provenance(config.as_ref());
{
let manager = self.services.mcp_connection_manager.read().await;
if let Err(e) = manager
.notify_local_sandbox_state_change(&sandbox_state)
.await
{
warn!("failed to notify local MCP sandbox state change: {e:#}");
}
let rebuilt_manager = manager
.rebuild_pooled_backend(
shared_mcp_backend_pool.as_ref(),
SharedMcpBackendAcquireMode::ReuseExisting,
&mcp_servers,
config.mcp_oauth_credentials_store_mode,
auth_statuses,
self.get_tx_event(),
sandbox_state.clone(),
config.codex_home.clone(),
codex_apps_tools_cache_key(auth.as_ref()),
tool_plugin_provenance,
)
.await;
drop(manager);
let mut manager = self.services.mcp_connection_manager.write().await;
*manager = rebuilt_manager;
}
} else if let Err(e) = self
.services
.mcp_connection_manager
.read()
@@ -2394,6 +2475,11 @@ impl Session {
warn!("Failed to notify sandbox state change to MCP servers: {e:#}");
}
}
self.services
.mcp_connection_manager
.read()
.await
.set_approval_policy(&session_configuration.approval_policy);
let model_info = self
.services
@@ -2536,9 +2622,10 @@ impl Session {
};
self.new_turn_from_configuration(
sub_id,
session_configuration,
session_configuration.clone(),
/*final_output_json_schema*/ None,
/*sandbox_policy_changed*/ false,
Self::sandbox_state_for_session_configuration(&session_configuration),
/*sandbox_state_changed*/ false,
)
.await
}
@@ -4059,18 +4146,39 @@ impl Session {
guard.cancel();
*guard = CancellationToken::new();
}
let (refreshed_manager, cancel_token) = McpConnectionManager::new(
&mcp_servers,
store_mode,
auth_statuses,
&turn_context.config.permissions.approval_policy,
self.get_tx_event(),
sandbox_state,
config.codex_home.clone(),
codex_apps_tools_cache_key(auth.as_ref()),
tool_plugin_provenance,
)
.await;
let (refreshed_manager, cancel_token) = match self.services.shared_mcp_backend_pool.as_ref()
{
Some(shared_mcp_backend_pool) => {
McpConnectionManager::new_with_pool(
shared_mcp_backend_pool.as_ref(),
SharedMcpBackendAcquireMode::ForceCreate,
&mcp_servers,
store_mode,
auth_statuses,
&turn_context.config.permissions.approval_policy,
self.get_tx_event(),
sandbox_state,
config.codex_home.clone(),
codex_apps_tools_cache_key(auth.as_ref()),
tool_plugin_provenance,
)
.await
}
None => {
McpConnectionManager::new(
&mcp_servers,
store_mode,
auth_statuses,
&turn_context.config.permissions.approval_policy,
self.get_tx_event(),
sandbox_state,
config.codex_home.clone(),
codex_apps_tools_cache_key(auth.as_ref()),
tool_plugin_provenance,
)
.await
}
};
{
let mut guard = self.services.mcp_startup_cancellation_token.lock().await;
if guard.is_cancelled() {
+81 -6
View File
@@ -10,6 +10,9 @@ use crate::config_loader::Sourced;
use crate::exec::ExecCapturePolicy;
use crate::exec::ExecToolCallOutput;
use crate::function_tool::FunctionCallError;
use crate::mcp::ToolPluginProvenance;
use crate::mcp_connection_manager::SharedMcpBackendAcquireMode;
use crate::mcp_connection_manager::SharedMcpBackendPool;
use crate::mcp_connection_manager::ToolInfo;
use crate::models_manager::model_info;
use crate::shell::default_user_shell;
@@ -2466,6 +2469,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
),
)),
mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()),
shared_mcp_backend_pool: None,
unified_exec_manager: UnifiedExecProcessManager::new(
config.background_terminal_max_timeout,
),
@@ -3172,6 +3176,17 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx(
Arc<Session>,
Arc<TurnContext>,
async_channel::Receiver<Event>,
) {
make_session_and_context_with_dynamic_tools_and_rx_and_pool(dynamic_tools, false).await
}
async fn make_session_and_context_with_dynamic_tools_and_rx_and_pool(
dynamic_tools: Vec<DynamicToolSpec>,
pooled: bool,
) -> (
Arc<Session>,
Arc<TurnContext>,
async_channel::Receiver<Event>,
) {
let (tx_event, rx_event) = async_channel::unbounded();
let codex_home = tempfile::tempdir().expect("create temp dir");
@@ -3256,15 +3271,44 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx(
.await
.expect("create environment"),
);
let shared_mcp_backend_pool = pooled.then(|| Arc::new(SharedMcpBackendPool::new()));
let sandbox_state = SandboxState {
sandbox_policy: session_configuration.sandbox_policy.get().clone(),
codex_linux_sandbox_exe: config.codex_linux_sandbox_exe.clone(),
sandbox_cwd: session_configuration.cwd.clone(),
use_legacy_landlock: config.features.use_legacy_landlock(),
};
let (mcp_connection_manager, mcp_startup_cancellation_token) =
match shared_mcp_backend_pool.as_ref() {
Some(shared_mcp_backend_pool) => {
McpConnectionManager::new_with_pool(
shared_mcp_backend_pool.as_ref(),
SharedMcpBackendAcquireMode::ReuseExisting,
&HashMap::new(),
config.mcp_oauth_credentials_store_mode,
HashMap::new(),
&config.permissions.approval_policy,
tx_event.clone(),
sandbox_state,
config.codex_home.clone(),
codex_apps_tools_cache_key(Some(&auth_manager.auth().await.expect("auth"))),
ToolPluginProvenance::default(),
)
.await
}
None => (
McpConnectionManager::new_mcp_connection_manager_for_tests(
&config.permissions.approval_policy,
),
CancellationToken::new(),
),
};
let file_watcher = Arc::new(FileWatcher::noop());
let services = SessionServices {
mcp_connection_manager: Arc::new(RwLock::new(
McpConnectionManager::new_mcp_connection_manager_for_tests(
&config.permissions.approval_policy,
),
)),
mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()),
mcp_connection_manager: Arc::new(RwLock::new(mcp_connection_manager)),
mcp_startup_cancellation_token: Mutex::new(mcp_startup_cancellation_token),
shared_mcp_backend_pool,
unified_exec_manager: UnifiedExecProcessManager::new(
config.background_terminal_max_timeout,
),
@@ -3365,6 +3409,14 @@ pub(crate) async fn make_session_and_context_with_rx() -> (
make_session_and_context_with_dynamic_tools_and_rx(Vec::new()).await
}
async fn make_pooled_session_and_context_with_rx() -> (
Arc<Session>,
Arc<TurnContext>,
async_channel::Receiver<Event>,
) {
make_session_and_context_with_dynamic_tools_and_rx_and_pool(Vec::new(), true).await
}
#[tokio::test]
async fn refresh_mcp_servers_is_deferred_until_next_turn() {
let (session, turn_context) = make_session_and_context().await;
@@ -3407,6 +3459,29 @@ async fn refresh_mcp_servers_is_deferred_until_next_turn() {
assert!(!new_token.is_cancelled());
}
#[tokio::test]
async fn pooled_new_turn_keeps_local_mcp_startup_token_when_sandbox_state_changes() {
let (session, _turn_context, _rx) = make_pooled_session_and_context_with_rx().await;
let old_token = session.mcp_startup_cancellation_token().await;
assert!(!old_token.is_cancelled());
let new_cwd = session.get_config().await.cwd.join("other-worktree");
session
.new_turn_with_sub_id(
"sandbox-change".to_string(),
SessionSettingsUpdate {
cwd: Some(new_cwd),
..SessionSettingsUpdate::default()
},
)
.await
.expect("new turn");
assert!(!old_token.is_cancelled());
let new_token = session.mcp_startup_cancellation_token().await;
assert!(!new_token.is_cancelled());
}
#[tokio::test]
async fn record_model_warning_appends_user_message() {
let (mut session, turn_context) = make_session_and_context().await;
File diff suppressed because it is too large Load Diff
@@ -1,7 +1,9 @@
use super::*;
use codex_protocol::protocol::GranularApprovalConfig;
use codex_protocol::protocol::McpAuthStatus;
use pretty_assertions::assert_eq;
use rmcp::model::JsonObject;
use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::Arc;
use tempfile::tempdir;
@@ -60,6 +62,98 @@ fn create_codex_apps_tools_cache_context(
}
}
fn test_codex_apps_tools_cache_key() -> CodexAppsToolsCacheKey {
CodexAppsToolsCacheKey {
account_id: None,
chatgpt_user_id: None,
is_workspace_account: false,
}
}
fn stdio_server_config(command: &str) -> McpServerConfig {
McpServerConfig {
transport: McpServerTransportConfig::Stdio {
command: command.to_string(),
args: Vec::new(),
env: None,
env_vars: Vec::new(),
cwd: None,
},
enabled: true,
required: false,
disabled_reason: None,
startup_timeout_sec: Some(Duration::from_secs(1)),
tool_timeout_sec: Some(Duration::from_secs(1)),
enabled_tools: None,
disabled_tools: None,
scopes: None,
oauth_resource: None,
}
}
fn http_server_config(url: &str) -> McpServerConfig {
McpServerConfig {
transport: McpServerTransportConfig::StreamableHttp {
url: url.to_string(),
bearer_token_env_var: None,
http_headers: None,
env_http_headers: None,
},
enabled: true,
required: false,
disabled_reason: None,
startup_timeout_sec: Some(Duration::from_secs(1)),
tool_timeout_sec: Some(Duration::from_secs(1)),
enabled_tools: None,
disabled_tools: None,
scopes: None,
oauth_resource: None,
}
}
fn test_sandbox_state(cwd: PathBuf) -> SandboxState {
SandboxState {
sandbox_policy: SandboxPolicy::DangerFullAccess,
codex_linux_sandbox_exe: None,
sandbox_cwd: cwd,
use_legacy_landlock: false,
}
}
fn pooled_cache_key_for_tests(
mcp_servers: &HashMap<String, McpServerConfig>,
store_mode: OAuthCredentialsStoreMode,
sandbox_state: &SandboxState,
) -> SharedMcpBackendCacheKey {
let (pooled_servers, _) = split_poolable_mcp_servers(mcp_servers);
SharedMcpBackendCacheKey::new(&pooled_servers, store_mode, sandbox_state)
}
async fn new_pooled_manager_for_tests(
pool: &SharedMcpBackendPool,
acquire_mode: SharedMcpBackendAcquireMode,
mcp_servers: &HashMap<String, McpServerConfig>,
sandbox_state: SandboxState,
) -> McpConnectionManager {
let approval_policy = Constrained::allow_any(AskForApproval::OnRequest);
let (tx_event, _rx_event) = async_channel::unbounded();
let (manager, _cancel_token) = McpConnectionManager::new_with_pool(
pool,
acquire_mode,
mcp_servers,
OAuthCredentialsStoreMode::Auto,
HashMap::new(),
&approval_policy,
tx_event,
sandbox_state,
PathBuf::from("/tmp"),
test_codex_apps_tools_cache_key(),
ToolPluginProvenance::default(),
)
.await;
manager
}
#[test]
fn elicitation_granular_policy_defaults_to_prompting() {
assert!(!elicitation_is_rejected_by_policy(
@@ -518,6 +612,246 @@ async fn list_all_tools_uses_startup_snapshot_when_client_startup_fails() {
assert_eq!(tool.tool_name, "calendar_create_event");
}
#[tokio::test]
async fn parse_tool_name_searches_shared_backend() {
let approval_policy = Constrained::allow_any(AskForApproval::OnFailure);
let shared_tool = create_test_tool("shared_stdio", "tool_a");
let pending_client = futures::future::pending::<Result<ManagedClient, StartupOutcomeError>>()
.boxed()
.shared();
let mut shared_backend = SharedMcpBackend::new_uninitialized();
shared_backend.clients.insert(
"shared_stdio".to_string(),
AsyncManagedClient {
client: pending_client,
startup_snapshot: Some(vec![shared_tool]),
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()),
},
);
let manager = McpConnectionManager::from_parts(
Arc::new(SharedMcpBackend::new_uninitialized()),
Some(Arc::new(shared_backend)),
SessionMcpHandle::new(approval_policy.value()),
None,
);
assert_eq!(
manager.parse_tool_name("mcp__shared_stdio__tool_a").await,
Some(("shared_stdio".to_string(), "tool_a".to_string()))
);
}
#[test]
fn shared_mcp_backend_cache_key_is_stable_for_equivalent_stdio_configs() {
let mut left_servers = HashMap::new();
left_servers.insert("beta".to_string(), stdio_server_config("missing-beta"));
left_servers.insert("alpha".to_string(), stdio_server_config("missing-alpha"));
let mut right_servers = HashMap::new();
right_servers.insert("alpha".to_string(), stdio_server_config("missing-alpha"));
right_servers.insert("beta".to_string(), stdio_server_config("missing-beta"));
let sandbox_state = test_sandbox_state(PathBuf::from("/tmp/shared"));
assert_eq!(
pooled_cache_key_for_tests(
&left_servers,
OAuthCredentialsStoreMode::Auto,
&sandbox_state,
),
pooled_cache_key_for_tests(
&right_servers,
OAuthCredentialsStoreMode::Auto,
&sandbox_state,
),
);
}
#[test]
fn shared_mcp_backend_cache_key_ignores_http_servers() {
let sandbox_state = test_sandbox_state(PathBuf::from("/tmp/shared"));
let mut left_servers = HashMap::new();
left_servers.insert("stdio".to_string(), stdio_server_config("missing-stdio"));
left_servers.insert(
"http".to_string(),
http_server_config("http://127.0.0.1:9/left"),
);
let mut right_servers = HashMap::new();
right_servers.insert("stdio".to_string(), stdio_server_config("missing-stdio"));
right_servers.insert(
"http".to_string(),
http_server_config("http://127.0.0.1:9/right"),
);
assert_eq!(
pooled_cache_key_for_tests(
&left_servers,
OAuthCredentialsStoreMode::Auto,
&sandbox_state,
),
pooled_cache_key_for_tests(
&right_servers,
OAuthCredentialsStoreMode::Auto,
&sandbox_state,
),
);
}
#[test]
fn split_poolable_mcp_servers_keeps_http_servers_local() {
let mut servers = HashMap::new();
servers.insert("stdio".to_string(), stdio_server_config("missing-stdio"));
servers.insert(
"http".to_string(),
http_server_config("http://127.0.0.1:9/http"),
);
let (pooled_servers, local_servers) = split_poolable_mcp_servers(&servers);
assert_eq!(pooled_servers.len(), 1);
assert!(pooled_servers.contains_key("stdio"));
assert_eq!(local_servers.len(), 1);
assert!(local_servers.contains_key("http"));
}
#[test]
fn shared_mcp_backend_cache_key_separates_sandbox_state() {
let mut servers = HashMap::new();
servers.insert("stdio".to_string(), stdio_server_config("missing-stdio"));
assert_ne!(
pooled_cache_key_for_tests(
&servers,
OAuthCredentialsStoreMode::Auto,
&test_sandbox_state(PathBuf::from("/tmp/left")),
),
pooled_cache_key_for_tests(
&servers,
OAuthCredentialsStoreMode::Auto,
&test_sandbox_state(PathBuf::from("/tmp/right")),
),
);
}
#[tokio::test]
async fn shared_mcp_backend_pool_reuses_backend_for_same_stdio_config() {
let pool = SharedMcpBackendPool::new();
let sandbox_state = test_sandbox_state(PathBuf::from("/tmp/shared"));
let mut servers = HashMap::new();
servers.insert("stdio".to_string(), stdio_server_config("missing-stdio"));
let manager_1 = new_pooled_manager_for_tests(
&pool,
SharedMcpBackendAcquireMode::ReuseExisting,
&servers,
sandbox_state.clone(),
)
.await;
let manager_2 = new_pooled_manager_for_tests(
&pool,
SharedMcpBackendAcquireMode::ReuseExisting,
&servers,
sandbox_state,
)
.await;
assert!(Arc::ptr_eq(
manager_1
.shared_backend
.as_ref()
.expect("stdio backend should be pooled"),
manager_2
.shared_backend
.as_ref()
.expect("stdio backend should be pooled"),
));
}
#[tokio::test]
async fn shared_mcp_backend_pool_separates_backends_for_different_sandbox_states() {
let pool = SharedMcpBackendPool::new();
let mut servers = HashMap::new();
servers.insert("stdio".to_string(), stdio_server_config("missing-stdio"));
let manager_1 = new_pooled_manager_for_tests(
&pool,
SharedMcpBackendAcquireMode::ReuseExisting,
&servers,
test_sandbox_state(PathBuf::from("/tmp/left")),
)
.await;
let manager_2 = new_pooled_manager_for_tests(
&pool,
SharedMcpBackendAcquireMode::ReuseExisting,
&servers,
test_sandbox_state(PathBuf::from("/tmp/right")),
)
.await;
assert!(!Arc::ptr_eq(
manager_1
.shared_backend
.as_ref()
.expect("stdio backend should be pooled"),
manager_2
.shared_backend
.as_ref()
.expect("stdio backend should be pooled"),
));
}
#[tokio::test]
async fn shared_mcp_backend_pool_force_create_replaces_pool_entry_for_same_key() {
let pool = SharedMcpBackendPool::new();
let mut servers = HashMap::new();
servers.insert("stdio".to_string(), stdio_server_config("missing-stdio"));
let sandbox_state = test_sandbox_state(PathBuf::from("/tmp/shared"));
let manager_1 = new_pooled_manager_for_tests(
&pool,
SharedMcpBackendAcquireMode::ReuseExisting,
&servers,
sandbox_state.clone(),
)
.await;
let manager_2 = new_pooled_manager_for_tests(
&pool,
SharedMcpBackendAcquireMode::ForceCreate,
&servers,
sandbox_state.clone(),
)
.await;
let manager_3 = new_pooled_manager_for_tests(
&pool,
SharedMcpBackendAcquireMode::ReuseExisting,
&servers,
sandbox_state,
)
.await;
let shared_1 = manager_1
.shared_backend
.as_ref()
.expect("stdio backend should be pooled");
let shared_2 = manager_2
.shared_backend
.as_ref()
.expect("stdio backend should be pooled");
let shared_3 = manager_3
.shared_backend
.as_ref()
.expect("stdio backend should be pooled");
assert!(!Arc::ptr_eq(shared_1, shared_2));
assert!(Arc::ptr_eq(shared_2, shared_3));
assert!(!Arc::ptr_eq(shared_1, shared_3));
}
#[test]
fn elicitation_capability_enabled_only_for_codex_apps() {
let codex_apps_capability = elicitation_capability_for_server(CODEX_APPS_MCP_SERVER_NAME);
+2
View File
@@ -11,6 +11,7 @@ use crate::exec_policy::ExecPolicyManager;
use crate::file_watcher::FileWatcher;
use crate::mcp::McpManager;
use crate::mcp_connection_manager::McpConnectionManager;
use crate::mcp_connection_manager::SharedMcpBackendPool;
use crate::models_manager::manager::ModelsManager;
use crate::plugins::PluginsManager;
use crate::skills::SkillsManager;
@@ -33,6 +34,7 @@ use tokio_util::sync::CancellationToken;
pub(crate) struct SessionServices {
pub(crate) mcp_connection_manager: Arc<RwLock<McpConnectionManager>>,
pub(crate) mcp_startup_cancellation_token: Mutex<CancellationToken>,
pub(crate) shared_mcp_backend_pool: Option<Arc<SharedMcpBackendPool>>,
pub(crate) unified_exec_manager: UnifiedExecProcessManager,
#[cfg_attr(not(unix), allow(dead_code))]
pub(crate) shell_zsh_path: Option<PathBuf>,
+8
View File
@@ -14,6 +14,7 @@ use crate::error::Result as CodexResult;
use crate::file_watcher::FileWatcher;
use crate::file_watcher::FileWatcherEvent;
use crate::mcp::McpManager;
use crate::mcp_connection_manager::SharedMcpBackendPool;
use crate::models_manager::collaboration_mode_presets::CollaborationModesConfig;
use crate::models_manager::manager::ModelsManager;
use crate::plugins::PluginsManager;
@@ -155,6 +156,7 @@ pub(crate) struct ThreadManagerState {
skills_manager: Arc<SkillsManager>,
plugins_manager: Arc<PluginsManager>,
mcp_manager: Arc<McpManager>,
shared_mcp_backend_pool: Arc<SharedMcpBackendPool>,
file_watcher: Arc<FileWatcher>,
session_source: SessionSource,
// Captures submitted ops for testing purpose when test mode is enabled.
@@ -202,6 +204,7 @@ impl ThreadManager {
skills_manager,
plugins_manager,
mcp_manager,
shared_mcp_backend_pool: Arc::new(SharedMcpBackendPool::new()),
file_watcher,
auth_manager,
session_source,
@@ -266,6 +269,7 @@ impl ThreadManager {
skills_manager,
plugins_manager,
mcp_manager,
shared_mcp_backend_pool: Arc::new(SharedMcpBackendPool::new()),
file_watcher,
auth_manager,
session_source: SessionSource::Exec,
@@ -582,6 +586,10 @@ impl ThreadManager {
}
impl ThreadManagerState {
pub(crate) fn shared_mcp_backend_pool(&self) -> Arc<SharedMcpBackendPool> {
Arc::clone(&self.shared_mcp_backend_pool)
}
pub(crate) async fn list_thread_ids(&self) -> Vec<ThreadId> {
self.threads.read().await.keys().copied().collect()
}