openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-15 12:08:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

wip: adding the ability to refresh rmcp client per thread after changes

Browse Source
This commit is contained in:
shijie-openai
2025-12-08 15:16:25 -08:00
parent bf53b8c3c7
commit f4e3b2f945
7 changed files with 153 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
codex-rs/app-server/src/codex_message_processor.rs
View File

@@ -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,

64
codex-rs/core/src/codex.rs
View File

@@ -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<ModelsManager>,
conversation_history: InitialHistory,
session_source: SessionSource,
mcp_oauth_refresh_clock: Arc<AtomicI64>,
) -> CodexResult<CodexSpawnOk> {
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<Config>,
@@ -474,6 +479,7 @@ impl Session {
tx_event: Sender<Event>,
initial_history: InitialHistory,
session_source: SessionSource,
mcp_oauth_refresh_clock: Arc<AtomicI64>,
) -> anyhow::Result<Arc<Self>> {
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<ListResourcesRequestParams>,
) -> anyhow::Result<ListResourcesResult> {
self.refresh_mcp_clients_if_needed().await?;
self.services
.mcp_connection_manager
.read()
@@ -1399,6 +1409,7 @@ impl Session {
server: &str,
params: Option<ListResourceTemplatesRequestParams>,
) -> anyhow::Result<ListResourceTemplatesResult> {
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<ReadResourceResult> {
self.refresh_mcp_clients_if_needed().await?;
self.services
.mcp_connection_manager
.read()
@@ -1426,6 +1438,7 @@ impl Session {
tool: &str,
arguments: Option<serde_json::Value>,
) -> anyhow::Result<CallToolResult> {
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<Self>) {
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),

1
codex-rs/core/src/codex_delegate.rs
View File

@@ -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);

16
codex-rs/core/src/conversation_manager.rs
View File

@@ -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<AuthManager>,
models_manager: Arc<ModelsManager>,
session_source: SessionSource,
mcp_oauth_refresh_clock: Arc<AtomicI64>,
}
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<AtomicI64> {
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<NewConversation> {
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?;

5
codex-rs/core/src/mcp/mod.rs
View File

@@ -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();

66
codex-rs/core/src/mcp_connection_manager.rs
View File

@@ -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<String, AsyncManagedClient>,
elicitation_requests: ElicitationRequestManager,
mcp_oauth_refresh_clock: Arc<AtomicI64>,
last_refresh_seen: AtomicI64,
config_snapshot: HashMap<String, McpServerConfig>,
store_mode_snapshot: Option<OAuthCredentialsStoreMode>,
auth_entries_snapshot: HashMap<String, McpAuthStatusEntry>,
}
impl McpConnectionManager {
pub(crate) fn new(mcp_oauth_refresh_clock: Arc<AtomicI64>) -> 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<String, McpServerConfig>,
store_mode: OAuthCredentialsStoreMode,
auth_entries: &HashMap<String, McpAuthStatusEntry>,
) {
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<String, McpServerConfig>,
store_mode: OAuthCredentialsStoreMode,
auth_entries: HashMap<String, McpAuthStatusEntry>,
tx_event: Sender<Event>,
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<String, McpServerConfig>,
@@ -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();

2
codex-rs/core/src/state/service.rs
View File

@@ -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<RwLock<McpConnectionManager>>,
pub(crate) mcp_startup_cancellation_token: CancellationToken,
pub(crate) mcp_oauth_refresh_clock: Arc<AtomicI64>,
pub(crate) unified_exec_manager: UnifiedExecSessionManager,
pub(crate) notifier: UserNotifier,
pub(crate) rollout: Mutex<Option<RolloutRecorder>>,