From f4e3b2f9455eddff02288592c48c1aad0ddfa075 Mon Sep 17 00:00:00 2001 From: shijie-openai Date: Mon, 8 Dec 2025 15:16:25 -0800 Subject: [PATCH] wip: adding the ability to refresh rmcp client per thread after changes --- .../app-server/src/codex_message_processor.rs | 5 ++ codex-rs/core/src/codex.rs | 64 +++++++++++++++++- codex-rs/core/src/codex_delegate.rs | 1 + codex-rs/core/src/conversation_manager.rs | 16 +++++ codex-rs/core/src/mcp/mod.rs | 5 +- codex-rs/core/src/mcp_connection_manager.rs | 66 ++++++++++++++++++- codex-rs/core/src/state/service.rs | 2 + 7 files changed, 153 insertions(+), 6 deletions(-) diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 0a8445055d..79d11e1654 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -2011,6 +2011,7 @@ impl CodexMessageProcessor { let authorization_url = handle.authorization_url().to_string(); let notification_name = name.clone(); let outgoing = Arc::clone(&self.outgoing); + let conversation_manager = Arc::clone(&self.conversation_manager); tokio::spawn(async move { let (success, error) = match handle.wait().await { @@ -2018,6 +2019,10 @@ impl CodexMessageProcessor { Err(err) => (false, Some(err.to_string())), }; + if success { + conversation_manager.mark_mcp_oauth_success(Utc::now().timestamp()); + } + let notification = ServerNotification::McpServerOauthLoginCompleted( McpServerOauthLoginCompletedNotification { name: notification_name, diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 042ae1a37a..279e3db0c7 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -55,6 +55,8 @@ use mcp_types::ReadResourceResult; use mcp_types::RequestId; use serde_json; use serde_json::Value; +use std::sync::atomic::AtomicI64; +use std::sync::atomic::Ordering; use tokio::sync::Mutex; use tokio::sync::RwLock; use tokio::sync::oneshot; @@ -170,6 +172,7 @@ impl Codex { models_manager: Arc, conversation_history: InitialHistory, session_source: SessionSource, + mcp_oauth_refresh_clock: Arc, ) -> CodexResult { let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); let (tx_event, rx_event) = async_channel::unbounded(); @@ -210,6 +213,7 @@ impl Codex { tx_event.clone(), conversation_history, session_source_clone, + mcp_oauth_refresh_clock.clone(), ) .await .map_err(|e| { @@ -466,6 +470,7 @@ impl Session { } } + #[allow(clippy::too_many_arguments)] async fn new( session_configuration: SessionConfiguration, config: Arc, @@ -474,6 +479,7 @@ impl Session { tx_event: Sender, initial_history: InitialHistory, session_source: SessionSource, + mcp_oauth_refresh_clock: Arc, ) -> anyhow::Result> { debug!( "Configuring session: model={}; provider={:?}", @@ -583,8 +589,11 @@ impl Session { let state = SessionState::new(session_configuration.clone()); let services = SessionServices { - mcp_connection_manager: Arc::new(RwLock::new(McpConnectionManager::default())), + mcp_connection_manager: Arc::new(RwLock::new(McpConnectionManager::new( + mcp_oauth_refresh_clock.clone(), + ))), mcp_startup_cancellation_token: CancellationToken::new(), + mcp_oauth_refresh_clock, unified_exec_manager: UnifiedExecSessionManager::default(), notifier: UserNotifier::new(config.notify.clone()), rollout: Mutex::new(Some(rollout_recorder)), @@ -1386,6 +1395,7 @@ impl Session { server: &str, params: Option, ) -> anyhow::Result { + self.refresh_mcp_clients_if_needed().await?; self.services .mcp_connection_manager .read() @@ -1399,6 +1409,7 @@ impl Session { server: &str, params: Option, ) -> anyhow::Result { + self.refresh_mcp_clients_if_needed().await?; self.services .mcp_connection_manager .read() @@ -1412,6 +1423,7 @@ impl Session { server: &str, params: ReadResourceRequestParams, ) -> anyhow::Result { + self.refresh_mcp_clients_if_needed().await?; self.services .mcp_connection_manager .read() @@ -1426,6 +1438,7 @@ impl Session { tool: &str, arguments: Option, ) -> anyhow::Result { + self.refresh_mcp_clients_if_needed().await?; self.services .mcp_connection_manager .read() @@ -1435,6 +1448,7 @@ impl Session { } pub(crate) async fn parse_mcp_tool_name(&self, tool_name: &str) -> Option<(String, String)> { + self.refresh_mcp_clients_if_needed().await.ok()?; self.services .mcp_connection_manager .read() @@ -1443,6 +1457,42 @@ impl Session { .await } + async fn refresh_mcp_clients_if_needed(&self) -> anyhow::Result<()> { + let current_clock = self.services.mcp_oauth_refresh_clock.load(Ordering::SeqCst); + let last_seen = { + let manager = self.services.mcp_connection_manager.read().await; + manager.last_refresh_seen() + }; + if current_clock <= last_seen { + return Ok(()); + } + + let config = { + let state = self.state.lock().await; + state + .session_configuration + .original_config_do_not_use + .clone() + }; + let store_mode = config.mcp_oauth_credentials_store_mode; + let auth_statuses = compute_auth_statuses(config.mcp_servers.iter(), store_mode).await; + + { + let mut manager = self.services.mcp_connection_manager.write().await; + manager + .refresh_if_needed( + &config.mcp_servers, + store_mode, + auth_statuses, + self.tx_event.clone(), + self.services.mcp_startup_cancellation_token.clone(), + ) + .await; + } + + Ok(()) + } + pub async fn interrupt_task(self: &Arc) { info!("interrupt received: abort current task, if any"); let has_active_turn = { self.active_turn.lock().await.is_some() }; @@ -2882,9 +2932,13 @@ mod tests { let state = SessionState::new(session_configuration.clone()); + let mcp_oauth_refresh_clock = Arc::new(AtomicI64::new(0)); let services = SessionServices { - mcp_connection_manager: Arc::new(RwLock::new(McpConnectionManager::default())), + mcp_connection_manager: Arc::new(RwLock::new(McpConnectionManager::new( + mcp_oauth_refresh_clock.clone(), + ))), mcp_startup_cancellation_token: CancellationToken::new(), + mcp_oauth_refresh_clock, unified_exec_manager: UnifiedExecSessionManager::default(), notifier: UserNotifier::new(None), rollout: Mutex::new(None), @@ -2964,9 +3018,13 @@ mod tests { let state = SessionState::new(session_configuration.clone()); + let mcp_oauth_refresh_clock = Arc::new(AtomicI64::new(0)); let services = SessionServices { - mcp_connection_manager: Arc::new(RwLock::new(McpConnectionManager::default())), + mcp_connection_manager: Arc::new(RwLock::new(McpConnectionManager::new( + mcp_oauth_refresh_clock.clone(), + ))), mcp_startup_cancellation_token: CancellationToken::new(), + mcp_oauth_refresh_clock, unified_exec_manager: UnifiedExecSessionManager::default(), notifier: UserNotifier::new(None), rollout: Mutex::new(None), diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 670225ead0..efadc90029 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -51,6 +51,7 @@ pub(crate) async fn run_codex_conversation_interactive( models_manager, initial_history.unwrap_or(InitialHistory::New), SessionSource::SubAgent(SubAgentSource::Review), + parent_session.services.mcp_oauth_refresh_clock.clone(), ) .await?; let codex = Arc::new(codex); diff --git a/codex-rs/core/src/conversation_manager.rs b/codex-rs/core/src/conversation_manager.rs index b1818849eb..570020e622 100644 --- a/codex-rs/core/src/conversation_manager.rs +++ b/codex-rs/core/src/conversation_manager.rs @@ -22,6 +22,8 @@ use codex_protocol::protocol::SessionSource; use std::collections::HashMap; use std::path::PathBuf; use std::sync::Arc; +use std::sync::atomic::AtomicI64; +use std::sync::atomic::Ordering; use tokio::sync::RwLock; /// Represents a newly created Codex conversation, including the first event @@ -39,6 +41,7 @@ pub struct ConversationManager { auth_manager: Arc, models_manager: Arc, session_source: SessionSource, + mcp_oauth_refresh_clock: Arc, } impl ConversationManager { @@ -48,6 +51,7 @@ impl ConversationManager { auth_manager: auth_manager.clone(), session_source, models_manager: Arc::new(ModelsManager::new(auth_manager)), + mcp_oauth_refresh_clock: Arc::new(AtomicI64::new(0)), } } @@ -65,6 +69,15 @@ impl ConversationManager { self.session_source.clone() } + pub fn mcp_oauth_refresh_clock(&self) -> Arc { + self.mcp_oauth_refresh_clock.clone() + } + + pub fn mark_mcp_oauth_success(&self, timestamp_secs: i64) { + self.mcp_oauth_refresh_clock + .store(timestamp_secs, Ordering::SeqCst); + } + pub async fn new_conversation(&self, config: Config) -> CodexResult { self.spawn_conversation( config, @@ -89,6 +102,7 @@ impl ConversationManager { models_manager, InitialHistory::New, self.session_source.clone(), + self.mcp_oauth_refresh_clock.clone(), ) .await?; self.finalize_spawn(codex, conversation_id).await @@ -166,6 +180,7 @@ impl ConversationManager { self.models_manager.clone(), initial_history, self.session_source.clone(), + self.mcp_oauth_refresh_clock.clone(), ) .await?; self.finalize_spawn(codex, conversation_id).await @@ -207,6 +222,7 @@ impl ConversationManager { self.models_manager.clone(), history, self.session_source.clone(), + self.mcp_oauth_refresh_clock.clone(), ) .await?; diff --git a/codex-rs/core/src/mcp/mod.rs b/codex-rs/core/src/mcp/mod.rs index ed5f2ea69f..8c538f7153 100644 --- a/codex-rs/core/src/mcp/mod.rs +++ b/codex-rs/core/src/mcp/mod.rs @@ -1,5 +1,7 @@ pub mod auth; use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::AtomicI64; use async_channel::unbounded; use codex_protocol::protocol::McpListToolsResponseEvent; @@ -29,7 +31,8 @@ pub async fn collect_mcp_snapshot(config: &Config) -> McpListToolsResponseEvent ) .await; - let mut mcp_connection_manager = McpConnectionManager::default(); + let mcp_oauth_refresh_clock = Arc::new(AtomicI64::new(0)); + let mut mcp_connection_manager = McpConnectionManager::new(mcp_oauth_refresh_clock); let (tx_event, rx_event) = unbounded(); drop(rx_event); let cancel_token = CancellationToken::new(); diff --git a/codex-rs/core/src/mcp_connection_manager.rs b/codex-rs/core/src/mcp_connection_manager.rs index 11a90f77a8..b081787517 100644 --- a/codex-rs/core/src/mcp_connection_manager.rs +++ b/codex-rs/core/src/mcp_connection_manager.rs @@ -12,6 +12,8 @@ use std::env; use std::ffi::OsString; use std::path::PathBuf; use std::sync::Arc; +use std::sync::atomic::AtomicI64; +use std::sync::atomic::Ordering; use std::time::Duration; use crate::mcp::auth::McpAuthStatusEntry; @@ -260,13 +262,70 @@ pub struct SandboxState { } /// A thin wrapper around a set of running [`RmcpClient`] instances. -#[derive(Default)] pub(crate) struct McpConnectionManager { clients: HashMap, elicitation_requests: ElicitationRequestManager, + mcp_oauth_refresh_clock: Arc, + last_refresh_seen: AtomicI64, + config_snapshot: HashMap, + store_mode_snapshot: Option, + auth_entries_snapshot: HashMap, } impl McpConnectionManager { + pub(crate) fn new(mcp_oauth_refresh_clock: Arc) -> Self { + Self { + clients: HashMap::new(), + elicitation_requests: ElicitationRequestManager::default(), + mcp_oauth_refresh_clock, + last_refresh_seen: AtomicI64::new(0), + config_snapshot: HashMap::new(), + store_mode_snapshot: None, + auth_entries_snapshot: HashMap::new(), + } + } + + fn update_snapshots( + &mut self, + mcp_servers: &HashMap, + store_mode: OAuthCredentialsStoreMode, + auth_entries: &HashMap, + ) { + self.config_snapshot = mcp_servers.clone(); + self.store_mode_snapshot = Some(store_mode); + self.auth_entries_snapshot = auth_entries.clone(); + let now = self.mcp_oauth_refresh_clock.load(Ordering::SeqCst); + self.last_refresh_seen.store(now, Ordering::SeqCst); + } + + pub(crate) fn last_refresh_seen(&self) -> i64 { + self.last_refresh_seen.load(Ordering::SeqCst) + } + + pub(crate) async fn refresh_if_needed( + &mut self, + config: &HashMap, + store_mode: OAuthCredentialsStoreMode, + auth_entries: HashMap, + tx_event: Sender, + cancel_token: CancellationToken, + ) { + let current = self.mcp_oauth_refresh_clock.load(Ordering::SeqCst); + if current <= self.last_refresh_seen() { + return; + } + + self.initialize( + config.clone(), + store_mode, + auth_entries, + tx_event, + cancel_token, + ) + .await; + self.last_refresh_seen.store(current, Ordering::SeqCst); + } + pub async fn initialize( &mut self, mcp_servers: HashMap, @@ -281,7 +340,9 @@ impl McpConnectionManager { let mut clients = HashMap::new(); let mut join_set = JoinSet::new(); let elicitation_requests = ElicitationRequestManager::default(); - for (server_name, cfg) in mcp_servers.into_iter().filter(|(_, cfg)| cfg.enabled) { + for (server_name, cfg) in mcp_servers.iter().filter(|(_, cfg)| cfg.enabled) { + let server_name = server_name.to_string(); + let cfg = cfg.clone(); let cancel_token = cancel_token.child_token(); let _ = emit_update( &tx_event, @@ -333,6 +394,7 @@ impl McpConnectionManager { } self.clients = clients; self.elicitation_requests = elicitation_requests.clone(); + self.update_snapshots(&mcp_servers, store_mode, &auth_entries); tokio::spawn(async move { let outcomes = join_set.join_all().await; let mut summary = McpStartupCompleteEvent::default(); diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 7387bcedae..4410304277 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -8,6 +8,7 @@ use crate::tools::sandboxing::ApprovalStore; use crate::unified_exec::UnifiedExecSessionManager; use crate::user_notification::UserNotifier; use codex_otel::otel_event_manager::OtelEventManager; +use std::sync::atomic::AtomicI64; use tokio::sync::Mutex; use tokio::sync::RwLock; use tokio_util::sync::CancellationToken; @@ -15,6 +16,7 @@ use tokio_util::sync::CancellationToken; pub(crate) struct SessionServices { pub(crate) mcp_connection_manager: Arc>, pub(crate) mcp_startup_cancellation_token: CancellationToken, + pub(crate) mcp_oauth_refresh_clock: Arc, pub(crate) unified_exec_manager: UnifiedExecSessionManager, pub(crate) notifier: UserNotifier, pub(crate) rollout: Mutex>,