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

Move MCP connection helpers into the test module (#34962)

Browse Source 
Keep the `McpConnectionSet` test constructors and accessors alongside the
connection manager tests instead of defining test-only helpers in the
production module.

GitOrigin-RevId: 925392f8b7d903b131ef1b49afd7e50f41226b1f
This commit is contained in:
jif
2026-07-23 15:40:03 +00:00
committed by copyberry
parent 7d748d3bbc
commit 9e1f43dc2d
2 changed files with 72 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/codex-mcp/src/connection_manager.rs
-94
View File
@@ -63,15 +63,6 @@ use tokio::sync::RwLock;
use tokio::task::JoinSet;
use tracing::warn;
#[cfg(test)]
use codex_config::Constrained;
#[cfg(test)]
use codex_protocol::protocol::McpStartupFailureReason;
#[cfg(test)]
use codex_rmcp_client::McpAuthState;
#[cfg(test)]
use codex_rmcp_client::McpLoginRequirement;
pub(crate) struct McpServerConnection {
identity: Option<McpServerConnectionIdentity>,
client: AsyncManagedClient,
@@ -517,30 +508,6 @@ impl McpConnectionSet {
manager
}
#[cfg(test)]
fn new_uninitialized_with_permission_profile(
approval_policy: &Constrained<AskForApproval>,
permission_profile: &PermissionProfile,
prefix_mcp_tool_names: bool,
) -> Self {
Self {
servers: HashMap::new(),
required_servers: Vec::new(),
tool_catalog_revision: Arc::new(RwLock::new(0)),
codex_apps_tools_override: RwLock::new(None),
codex_apps_refresh_lock: Mutex::new(()),
tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()),
prefix_mcp_tool_names,
elicitation_requests: ElicitationRequestManager::new(
approval_policy.value(),
permission_profile.clone(),
/*reviewer*/ None,
/*lifecycle*/ None,
ElicitationRequestRouter::default(),
),
}
}
pub fn empty(prefix_mcp_tool_names: bool) -> Self {
Self {
servers: HashMap::new(),
@@ -685,67 +652,6 @@ impl McpConnectionSet {
}
server_infos
}
#[cfg(test)]
fn new_uninitialized(
approval_policy: &Constrained<AskForApproval>,
permission_profile: &Constrained<PermissionProfile>,
prefix_mcp_tool_names: bool,
) -> Self {
Self::new_uninitialized_with_permission_profile(
approval_policy,
permission_profile.get(),
prefix_mcp_tool_names,
)
}
#[cfg(test)]
fn insert_test_client(&mut self, name: impl Into<String>, client: AsyncManagedClient) {
let name = name.into();
self.servers.insert(
name,
McpServerView {
tool_filter: ToolFilter::default(),
connection: Arc::new(McpServerConnection {
identity: None,
client,
}),
metadata: McpServerMetadata {
environment_id: String::new(),
pollutes_memory: true,
origin: None,
supports_parallel_tool_calls: false,
default_tools_approval_mode: None,
tool_approval_modes: HashMap::new(),
},
tool_timeout: None,
},
);
}
#[cfg(test)]
fn test_client(&self, name: &str) -> &AsyncManagedClient {
&self.servers[name].connection.client
}
#[cfg(test)]
fn set_test_server_metadata(&mut self, name: &str, metadata: McpServerMetadata) {
self.servers
.get_mut(name)
.expect("test server exists")
.metadata = metadata;
}
#[cfg(test)]
fn shares_test_connection_with(&self, other: &Self, name: &str) -> bool {
let Some(left) = self.servers.get(name) else {
return false;
};
let Some(right) = other.servers.get(name) else {
return false;
};
Arc::ptr_eq(&left.connection, &right.connection)
}
}
#[cfg(test)]
codex-rs/codex-mcp/src/connection_manager_tests.rs
+72
View File
@@ -37,8 +37,11 @@ use codex_protocol::ToolName;
use codex_protocol::mcp::McpServerInfo;
use codex_protocol::models::PermissionProfile;
use codex_protocol::protocol::GranularApprovalConfig;
use codex_protocol::protocol::McpStartupFailureReason;
use codex_rmcp_client::ElicitationResponse;
use codex_rmcp_client::InProcessTransportFactory;
use codex_rmcp_client::McpAuthState;
use codex_rmcp_client::McpLoginRequirement;
use codex_rmcp_client::RmcpClient;
use futures::FutureExt;
use futures::future::BoxFuture;
@@ -74,6 +77,75 @@ use tokio::io::DuplexStream;
use tokio::sync::Notify;
use tokio_util::sync::CancellationToken;
impl McpConnectionSet {
fn new_uninitialized(
approval_policy: &Constrained<AskForApproval>,
permission_profile: &Constrained<PermissionProfile>,
prefix_mcp_tool_names: bool,
) -> Self {
Self {
servers: HashMap::new(),
required_servers: Vec::new(),
tool_catalog_revision: Arc::new(RwLock::new(0)),
codex_apps_tools_override: RwLock::new(None),
codex_apps_refresh_lock: Mutex::new(()),
tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()),
prefix_mcp_tool_names,
elicitation_requests: ElicitationRequestManager::new(
approval_policy.value(),
permission_profile.get().clone(),
/*reviewer*/ None,
/*lifecycle*/ None,
ElicitationRequestRouter::default(),
),
}
}
fn insert_test_client(&mut self, name: impl Into<String>, client: AsyncManagedClient) {
let name = name.into();
self.servers.insert(
name,
McpServerView {
tool_filter: ToolFilter::default(),
connection: Arc::new(McpServerConnection {
identity: None,
client,
}),
metadata: McpServerMetadata {
environment_id: String::new(),
pollutes_memory: true,
origin: None,
supports_parallel_tool_calls: false,
default_tools_approval_mode: None,
tool_approval_modes: HashMap::new(),
},
tool_timeout: None,
},
);
}
fn test_client(&self, name: &str) -> &AsyncManagedClient {
&self.servers[name].connection.client
}
fn set_test_server_metadata(&mut self, name: &str, metadata: McpServerMetadata) {
self.servers
.get_mut(name)
.expect("test server exists")
.metadata = metadata;
}
fn shares_test_connection_with(&self, other: &Self, name: &str) -> bool {
let Some(left) = self.servers.get(name) else {
return false;
};
let Some(right) = other.servers.get(name) else {
return false;
};
Arc::ptr_eq(&left.connection, &right.connection)
}
}
fn create_test_tool(server_name: &str, tool_name: &str) -> ToolInfo {
ToolInfo {
server_name: server_name.to_string(),