Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 28 additions & 0 deletions codex-rs/codex-mcp/src/connection_manager.rs
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,7 @@ pub struct McpConnectionManager {
clients: HashMap<String, AsyncManagedClient>,
server_origins: HashMap<String, String>,
elicitation_requests: ElicitationRequestManager,
startup_cancellation_token: CancellationToken,
}

impl McpConnectionManager {
Expand All @@ -85,13 +86,32 @@ impl McpConnectionManager {
approval_policy.value(),
permission_profile.get().clone(),
),
startup_cancellation_token: CancellationToken::new(),
}
}

pub fn has_servers(&self) -> bool {
!self.clients.is_empty()
}

/// Drain all MCP clients from this manager and return a future that stops
/// them and terminates their stdio server processes.
pub fn begin_shutdown(&mut self) -> impl std::future::Future<Output = ()> + Send + 'static {
self.startup_cancellation_token.cancel();
let clients = std::mem::take(&mut self.clients);
self.server_origins.clear();
async move {
for client in clients.into_values() {
client.shutdown().await;
}
}
}

/// Stop all MCP clients owned by this manager and terminate stdio server processes.
pub async fn shutdown(&mut self) {
self.begin_shutdown().await;
}

pub fn server_origin(&self, server_name: &str) -> Option<&str> {
self.server_origins.get(server_name).map(String::as_str)
}
Expand Down Expand Up @@ -221,6 +241,7 @@ impl McpConnectionManager {
clients,
server_origins,
elicitation_requests: elicitation_requests.clone(),
startup_cancellation_token: cancel_token.clone(),
};
tokio::spawn(async move {
let outcomes = join_set.join_all().await;
Expand Down Expand Up @@ -609,6 +630,13 @@ impl McpConnectionManager {
}
}

impl Drop for McpConnectionManager {
fn drop(&mut self) {
self.startup_cancellation_token.cancel();
self.clients.clear();
}
}

async fn emit_update(
submit_id: &str,
tx_event: &Sender<Event>,
Expand Down
5 changes: 5 additions & 0 deletions codex-rs/codex-mcp/src/connection_manager_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -662,6 +662,7 @@ async fn list_all_tools_uses_startup_snapshot_while_client_is_pending() {
startup_snapshot: Some(startup_tools),
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()),
cancel_token: CancellationToken::new(),
},
);

Expand Down Expand Up @@ -690,6 +691,7 @@ async fn resolve_tool_info_accepts_canonical_namespaced_tool_names() {
startup_snapshot: Some(startup_tools),
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()),
cancel_token: CancellationToken::new(),
},
);

Expand Down Expand Up @@ -726,6 +728,7 @@ async fn list_all_tools_blocks_while_client_is_pending_without_startup_snapshot(
startup_snapshot: None,
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()),
cancel_token: CancellationToken::new(),
},
);

Expand All @@ -750,6 +753,7 @@ async fn list_all_tools_does_not_block_when_startup_snapshot_cache_hit_is_empty(
startup_snapshot: Some(Vec::new()),
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()),
cancel_token: CancellationToken::new(),
},
);

Expand Down Expand Up @@ -784,6 +788,7 @@ async fn list_all_tools_uses_startup_snapshot_when_client_startup_fails() {
startup_snapshot: Some(startup_tools),
startup_complete,
tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()),
cancel_token: CancellationToken::new(),
},
);

Expand Down
31 changes: 23 additions & 8 deletions codex-rs/codex-mcp/src/rmcp_client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@ use rmcp::model::Implementation;
use rmcp::model::InitializeRequestParams;
use rmcp::model::ProtocolVersion;
use tokio_util::sync::CancellationToken;
use tracing::warn;

/// MCP server capability indicating that Codex should include [`SandboxState`]
/// in tool-call request `_meta` under this key.
Expand Down Expand Up @@ -115,6 +116,7 @@ pub(crate) struct AsyncManagedClient {
pub(crate) startup_snapshot: Option<Vec<ToolInfo>>,
pub(crate) startup_complete: Arc<AtomicBool>,
pub(crate) tool_plugin_provenance: Arc<ToolPluginProvenance>,
pub(crate) cancel_token: CancellationToken,
}

impl AsyncManagedClient {
Expand Down Expand Up @@ -142,8 +144,9 @@ impl AsyncManagedClient {
let startup_tool_filter = tool_filter;
let startup_complete = Arc::new(AtomicBool::new(false));
let startup_complete_for_fut = Arc::clone(&startup_complete);
let cancel_token_for_fut = cancel_token.clone();
let fut = async move {
let outcome = async {
let outcome = match async {
if let Err(error) = validate_mcp_server_name(&server_name) {
return Err(error.into());
}
Expand All @@ -158,7 +161,7 @@ impl AsyncManagedClient {
)
.await?,
);
match start_server_task(
start_server_task(
server_name,
client,
StartServerTaskParams {
Expand All @@ -172,14 +175,14 @@ impl AsyncManagedClient {
codex_apps_tools_cache_context,
},
)
.or_cancel(&cancel_token)
.await
{
Ok(result) => result,
Err(CancelErr::Cancelled) => Err(StartupOutcomeError::Cancelled),
}
}
.await;
.or_cancel(&cancel_token_for_fut)
.await
{
Ok(result) => result,
Err(CancelErr::Cancelled) => Err(StartupOutcomeError::Cancelled),
};

startup_complete_for_fut.store(true, Ordering::Release);
outcome
Expand All @@ -197,13 +200,25 @@ impl AsyncManagedClient {
startup_snapshot,
startup_complete,
tool_plugin_provenance,
cancel_token,
}
}

pub(crate) async fn client(&self) -> Result<ManagedClient, StartupOutcomeError> {
self.client.clone().await
}

pub(crate) async fn shutdown(&self) {
self.cancel_token.cancel();
match self.client().await {
Ok(client) => client.client.shutdown().await,
Err(StartupOutcomeError::Cancelled) => {}
Err(error) => {
warn!("failed to initialize MCP client during shutdown: {error:#}");
}
}
}

fn startup_snapshot_while_initializing(&self) -> Option<Vec<ToolInfo>> {
if !self.startup_complete.load(Ordering::Acquire) {
return self.startup_snapshot.clone();
Expand Down
3 changes: 2 additions & 1 deletion codex-rs/core/src/connectors.rs
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_environment_manager(
.default_environment()
.unwrap_or_else(|| environment_manager.local_environment());

let (mcp_connection_manager, cancel_token) = McpConnectionManager::new(
let (mut mcp_connection_manager, cancel_token) = McpConnectionManager::new(
&mcp_servers,
config.mcp_oauth_credentials_store_mode,
auth_status_entries,
Expand Down Expand Up @@ -346,6 +346,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_environment_manager(
}
let accessible_connectors =
with_app_plugin_sources(accessible_connectors, &tool_plugin_provenance);
mcp_connection_manager.shutdown().await;
Ok(AccessibleConnectorsStatus {
connectors: accessible_connectors,
codex_apps_ready,
Expand Down
20 changes: 18 additions & 2 deletions codex-rs/core/src/session/handlers.rs
Original file line number Diff line number Diff line change
Expand Up @@ -943,6 +943,11 @@ pub async fn shutdown(sess: &Arc<Session>, sub_id: String) -> bool {
.unified_exec_manager
.terminate_all_processes()
.await;
let mcp_shutdown = {
let mut manager = sess.services.mcp_connection_manager.write().await;
manager.begin_shutdown()
};
mcp_shutdown.await;
sess.guardian_review_session.shutdown().await;
info!("Shutting down Codex instance");
let history = sess.clone_history().await;
Expand Down Expand Up @@ -1231,8 +1236,19 @@ pub(super) async fn submission_loop(
break;
}
}
// Also drain cached guardian state if the submission loop exits because
// the channel closed without receiving an explicit shutdown op.
// If the submission loop exits because the channel closed without an
// explicit shutdown op, still run process teardown for child processes
// owned by this session.
sess.services
.unified_exec_manager
.terminate_all_processes()
.await;
let mcp_shutdown = {
let mut manager = sess.services.mcp_connection_manager.write().await;
manager.begin_shutdown()
};
mcp_shutdown.await;
// Also drain cached guardian state on this implicit shutdown path.
sess.guardian_review_session.shutdown().await;
debug!("Agent loop exited");
}
Expand Down
7 changes: 5 additions & 2 deletions codex-rs/core/src/session/mcp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -255,8 +255,11 @@ impl Session {
*guard = cancel_token;
}

let mut manager = self.services.mcp_connection_manager.write().await;
*manager = refreshed_manager;
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;
}

pub(crate) async fn refresh_mcp_servers_if_requested(&self, turn_context: &TurnContext) {
Expand Down
3 changes: 3 additions & 0 deletions codex-rs/rmcp-client/src/bin/test_stdio_server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -755,6 +755,9 @@ fn parse_data_url(url: &str) -> Option<(String, String)> {
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
eprintln!("starting rmcp test server");
if let Ok(pid_file) = std::env::var("MCP_TEST_PID_FILE") {
std::fs::write(pid_file, std::process::id().to_string())?;
}
// Run the server with STDIO transport. If the client disconnects we simply
// bubble up the error so the process exits.
let service = TestToolServer::new();
Expand Down
37 changes: 37 additions & 0 deletions codex-rs/rmcp-client/src/rmcp_client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,7 @@ use crate::oauth::OAuthPersistor;
use crate::oauth::StoredOAuthTokens;
use crate::stdio_server_launcher::StdioServerCommand;
use crate::stdio_server_launcher::StdioServerLauncher;
use crate::stdio_server_launcher::StdioServerProcessHandle;
use crate::stdio_server_launcher::StdioServerTransport;
use crate::utils::apply_default_headers;
use crate::utils::build_default_headers;
Expand All @@ -93,6 +94,7 @@ enum ClientState {
service: Arc<RunningService<RoleClient, ElicitationClientService>>,
oauth: Option<OAuthPersistor>,
},
Closed,
}

#[derive(Clone)]
Expand Down Expand Up @@ -265,6 +267,7 @@ pub struct ListToolsWithConnectorIdResult {
/// https://github.com/modelcontextprotocol/rust-sdk
pub struct RmcpClient {
state: Mutex<ClientState>,
stdio_process: Option<StdioServerProcessHandle>,
transport_recipe: TransportRecipe,
initialize_context: Mutex<Option<InitializeContext>>,
session_recovery_lock: Semaphore,
Expand All @@ -287,11 +290,17 @@ impl RmcpClient {
let transport = Self::create_pending_transport(&transport_recipe)
.await
.map_err(io::Error::other)?;
let stdio_process = match &transport {
PendingTransport::Stdio { transport } => Some(transport.process_handle()),
PendingTransport::StreamableHttp { .. }
| PendingTransport::StreamableHttpWithOAuth { .. } => None,
};

Ok(Self {
state: Mutex::new(ClientState::Connecting {
transport: Some(transport),
}),
stdio_process,
transport_recipe,
initialize_context: Mutex::new(None),
session_recovery_lock: Semaphore::new(/*permits*/ 1),
Expand Down Expand Up @@ -325,6 +334,7 @@ impl RmcpClient {
state: Mutex::new(ClientState::Connecting {
transport: Some(transport),
}),
stdio_process: None,
transport_recipe,
initialize_context: Mutex::new(None),
session_recovery_lock: Semaphore::new(/*permits*/ 1),
Expand Down Expand Up @@ -353,6 +363,7 @@ impl RmcpClient {
None => return Err(anyhow!("client already initializing")),
},
ClientState::Ready { .. } => return Err(anyhow!("client already initialized")),
ClientState::Closed => return Err(anyhow!("MCP client is shut down")),
}
};

Expand All @@ -376,6 +387,9 @@ impl RmcpClient {

{
let mut guard = self.state.lock().await;
if matches!(*guard, ClientState::Closed) {
return Err(anyhow!("MCP client is shut down"));
}
*guard = ClientState::Ready {
service,
oauth: oauth_persistor.clone(),
Expand Down Expand Up @@ -623,6 +637,7 @@ impl RmcpClient {
match &*guard {
ClientState::Ready { service, .. } => Ok(Arc::clone(service)),
ClientState::Connecting { .. } => Err(anyhow!("MCP client not initialized")),
ClientState::Closed => Err(anyhow!("MCP client is shut down")),
}
}

Expand All @@ -637,6 +652,22 @@ impl RmcpClient {
}
}

/// Stop the MCP transport and any stdio server process owned by this client.
pub async fn shutdown(&self) {
let previous_state = {
let mut guard = self.state.lock().await;
std::mem::replace(&mut *guard, ClientState::Closed)
};

if let Some(process) = &self.stdio_process
&& let Err(error) = process.terminate().await
{
warn!("failed to terminate MCP stdio server process: {error}");
}

drop(previous_state);
}

/// This should be called after every tool call so that if a given tool call triggered
/// a refresh of the OAuth tokens, they are persisted.
async fn persist_oauth_tokens(&self) {
Expand Down Expand Up @@ -900,6 +931,9 @@ impl RmcpClient {
ClientState::Connecting { .. } => {
return Err(anyhow!("MCP client not initialized"));
}
ClientState::Closed => {
return Err(anyhow!("MCP client is shut down"));
}
}
}

Expand All @@ -919,6 +953,9 @@ impl RmcpClient {

{
let mut guard = self.state.lock().await;
if matches!(*guard, ClientState::Closed) {
return Err(anyhow!("MCP client is shut down"));
}
*guard = ClientState::Ready {
service,
oauth: oauth_persistor.clone(),
Expand Down
Loading
Loading