diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index 1a97defe2d..f824f3b766 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -76,6 +76,8 @@ pub struct McpRuntime { struct PublishedMcpRuntime { connections: Arc, config: Option>, + auth: Option, + auth_token: Option, plugins_available: bool, ready_selected_capability_roots: Vec, } @@ -125,6 +127,8 @@ impl McpRuntime { current: ArcSwap::from_pointee(PublishedMcpRuntime { connections: Arc::new(McpConnectionSet::empty(prefix_mcp_tool_names)), config: None, + auth: None, + auth_token: None, plugins_available: false, ready_selected_capability_roots: Vec::new(), }), @@ -154,6 +158,8 @@ impl McpRuntime { async fn publish(&self, input: McpRuntimeInput, previous: Option<&McpConnectionSet>) { let (publish, publication_gate) = McpPublicationGate::pending(); let config = Arc::clone(&input.config); + let auth = input.auth.clone(); + let auth_token = auth.as_ref().and_then(|auth| auth.get_token().ok()); let plugins_available = input.plugins_available; let ready_selected_capability_roots = input.ready_selected_capability_roots.clone(); let connections = Arc::new( @@ -168,6 +174,8 @@ impl McpRuntime { self.current.store(Arc::new(PublishedMcpRuntime { connections, config: Some(config), + auth, + auth_token, plugins_available, ready_selected_capability_roots, })); @@ -186,6 +194,22 @@ impl McpRuntime { )) } + /// Returns whether the published snapshot still belongs to the current credentials. + pub fn current_auth_matches(&self, auth: Option<&CodexAuth>) -> bool { + let current = self.current.load(); + match (current.auth.as_ref(), auth) { + (Some(previous), Some(latest)) => { + previous == latest + && previous.get_account_id() == latest.get_account_id() + && previous.get_chatgpt_user_id() == latest.get_chatgpt_user_id() + && previous.is_fedramp_account() == latest.is_fedramp_account() + && current.auth_token == latest.get_token().ok() + } + (None, None) => true, + (Some(_), None) | (None, Some(_)) => false, + } + } + pub fn current_ready_selected_capability_roots(&self) -> Vec { self.current.load().ready_selected_capability_roots.clone() } diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 0a467a94fd..eff03d1287 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -152,6 +152,19 @@ impl Session { return; }; loop { + let auth = self.services.auth_manager.auth_cached(); + if self + .services + .plugins_manager + .set_auth_mode(auth.as_ref().map(CodexAuth::api_auth_mode)) + || !self + .services + .mcp_runtime + .current_auth_matches(auth.as_ref()) + { + self.mark_mcp_runtime_dirty(); + } + if !self .mcp_refresh_pending .swap(false, std::sync::atomic::Ordering::AcqRel) @@ -162,7 +175,11 @@ impl Session { pending: &self.mcp_refresh_pending, published: false, }; - let desired = self.latest_mcp_desired_state().await; + let auth = self.services.auth_manager.auth().await; + self.services + .plugins_manager + .set_auth_mode(auth.as_ref().map(CodexAuth::api_auth_mode)); + let desired = self.latest_mcp_desired_state(auth).await; let selected_capability_roots = self .resolve_selected_capability_roots_for_step(&desired.environments) .await; @@ -213,7 +230,11 @@ impl Session { .acquire() .await .map_err(|_| anyhow::anyhow!("MCP runtime refresh semaphore closed"))?; - let desired = self.latest_mcp_desired_state().await; + let auth = self.services.auth_manager.auth().await; + self.services + .plugins_manager + .set_auth_mode(auth.as_ref().map(CodexAuth::api_auth_mode)); + let desired = self.latest_mcp_desired_state(auth).await; let selected_capability_roots = self .resolve_selected_capability_roots_for_step(&desired.environments) .await; @@ -237,14 +258,12 @@ impl Session { executor_capability_discovery.as_deref(), ) .await; - let input = self - .build_mcp_runtime_input( - &desired, - mcp_projection, - &ready_selected_capability_roots, - Some(self.mcp_elicitation_reviewer()), - ) - .await; + let input = self.build_mcp_runtime_input( + &desired, + mcp_projection, + &ready_selected_capability_roots, + Some(self.mcp_elicitation_reviewer()), + ); anyhow::ensure!( input.mcp_servers.contains_key(CODEX_APPS_MCP_SERVER_NAME), "unknown MCP server '{CODEX_APPS_MCP_SERVER_NAME}'" @@ -516,6 +535,10 @@ impl Session { error!("MCP runtime refresh semaphore closed"); return; }; + let auth = self.services.auth_manager.auth().await; + self.services + .plugins_manager + .set_auth_mode(auth.as_ref().map(CodexAuth::api_auth_mode)); { let mut state = self.state.lock().await; let mut config = (*state.session_configuration.original_config_do_not_use).clone(); @@ -544,7 +567,7 @@ impl Session { executor_capability_discovery.as_deref(), ) .await; - let mut desired = self.latest_mcp_desired_state().await; + let mut desired = self.latest_mcp_desired_state(auth).await; desired.config = Arc::new(refresh_config.clone()); self.publish_mcp_runtime( &desired, diff --git a/codex-rs/core/src/session/mcp_runtime.rs b/codex-rs/core/src/session/mcp_runtime.rs index f8c4c08199..01f5ce4a26 100644 --- a/codex-rs/core/src/session/mcp_runtime.rs +++ b/codex-rs/core/src/session/mcp_runtime.rs @@ -12,6 +12,7 @@ use codex_protocol::capabilities::SelectedCapabilityRoot; pub(super) struct McpDesiredState { pub(super) config: Arc, + pub(super) auth: Option, pub(super) submit_id: String, pub(super) originator: String, pub(super) environments: TurnEnvironmentSnapshot, @@ -28,7 +29,10 @@ impl McpDesiredState { } impl Session { - pub(super) async fn latest_mcp_desired_state(&self) -> McpDesiredState { + pub(super) async fn latest_mcp_desired_state( + &self, + auth: Option, + ) -> McpDesiredState { let session_configuration = { let state = self.state.lock().await; state.session_configuration.clone() @@ -43,6 +47,7 @@ impl Session { McpDesiredState { config: Arc::new(config), + auth, submit_id: self.next_internal_sub_id(), originator: session_configuration.originator.clone(), environments, @@ -52,6 +57,7 @@ impl Session { pub(super) async fn install_initial_mcp_runtime( self: &Arc, session_configuration: &SessionConfiguration, + auth: Option, mcp_projection: McpRuntimeProjection, resolved_environments: &TurnEnvironmentSnapshot, local_stdio_fallback_cwd: PathBuf, @@ -62,6 +68,7 @@ impl Session { config.permissions.approval_policy = session_configuration.approval_policy.clone(); let desired = McpDesiredState { config: Arc::new(config), + auth, submit_id: INITIAL_SUBMIT_ID.to_owned(), originator: session_configuration.originator.clone(), environments: resolved_environments.clone(), @@ -89,25 +96,23 @@ impl Session { ready_selected_capability_roots: &[SelectedCapabilityRoot], elicitation_reviewer: Option, ) { - let input = self - .build_mcp_runtime_input( - desired, - mcp_projection, - ready_selected_capability_roots, - elicitation_reviewer, - ) - .await; + let input = self.build_mcp_runtime_input( + desired, + mcp_projection, + ready_selected_capability_roots, + elicitation_reviewer, + ); self.services.mcp_runtime.replace(input).await; } - pub(super) async fn build_mcp_runtime_input( + pub(super) fn build_mcp_runtime_input( &self, desired: &McpDesiredState, mcp_projection: McpRuntimeProjection, ready_selected_capability_roots: &[SelectedCapabilityRoot], elicitation_reviewer: Option, ) -> McpRuntimeInput { - let auth = self.services.auth_manager.auth().await; + let auth = desired.auth.clone(); let supports_openai_form_elicitation = self .services .supports_openai_form_elicitation diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index ac3e357f0b..2cfaf23e57 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -707,6 +707,7 @@ impl Session { session_init.ephemeral = config.ephemeral, )); + let mut mcp_auth_changes = auth_manager.auth_change_receiver(); let auth_manager_clone = Arc::clone(&auth_manager); let config_for_mcp = Arc::clone(&config); let mcp_manager_for_mcp = Arc::clone(&mcp_manager); @@ -1216,9 +1217,34 @@ impl Session { } turn_environments.start_connection_event_forwarding(tx_event.clone()); - let mcp_auth_changes = sess.services.auth_manager.auth_change_receiver(); + let startup_auth_changed = mcp_auth_changes.has_changed().unwrap_or(false); + if startup_auth_changed { + mcp_auth_changes.mark_unchanged(); + } + let latest_auth = sess.services.auth_manager.auth().await; + let mcp_projection = if startup_auth_changed + || mcp_auth_changes.has_changed().unwrap_or(false) + { + sess.services + .plugins_manager + .set_auth_mode(latest_auth.as_ref().map(CodexAuth::api_auth_mode)); + sess.services + .mcp_manager + .runtime_config_for_step( + config.as_ref(), + &sess.services.mcp_thread_init, + &sess.services.thread_extension_data, + &session_configuration.originator, + /*ready_selected_capability_roots*/ &[], + /*executor_capability_discovery*/ None, + ) + .await + } else { + mcp_projection + }; sess.install_initial_mcp_runtime( &session_configuration, + latest_auth, mcp_projection, &resolved_environments, mcp_runtime_cwd, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index edc39358a3..2a7e178701 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -7936,6 +7936,65 @@ async fn mcp_policy_changes_schedule_runtime_refresh() { ); } +#[tokio::test] +async fn mcp_refresh_updates_plugin_auth_mode_before_checking_pending_state() { + let codex_home = tempfile::tempdir().expect("create auth test directory"); + let (mut session, _turn_context) = make_session_and_context().await; + session.services.auth_manager = AuthManager::from_auth_for_testing_with_home( + CodexAuth::from_api_key("old-api-key"), + codex_home.path().to_path_buf(), + ); + let session = Arc::new(session); + let auth_mode = session.services.auth_manager.get_api_auth_mode(); + + assert_ne!(session.services.plugins_manager.auth_mode(), auth_mode); + session + .mcp_refresh_pending + .store(false, std::sync::atomic::Ordering::Release); + + session.refresh_mcp_if_dirty().await; + + assert_eq!(session.services.plugins_manager.auth_mode(), auth_mode); + assert!( + session + .services + .mcp_runtime + .current_binding() + .await + .is_some() + ); + + codex_login::login_with_api_key( + codex_home.path(), + "new-api-key", + codex_login::AuthCredentialsStoreMode::File, + codex_login::AuthKeyringBackendKind::default(), + ) + .expect("store replacement API key"); + session.services.auth_manager.reload().await; + assert_eq!( + session + .services + .auth_manager + .auth_cached() + .and_then(|auth| auth.get_token().ok()), + Some("new-api-key".to_string()) + ); + assert_eq!(session.services.plugins_manager.auth_mode(), auth_mode); + session + .mcp_refresh_pending + .store(false, std::sync::atomic::Ordering::Release); + + session.refresh_mcp_if_dirty().await; + + assert!( + session + .services + .mcp_runtime + .current_auth_matches(session.services.auth_manager.auth_cached().as_ref()) + ); +} + struct PendingNoiseConnectProvider; impl codex_exec_server::NoiseRendezvousConnectProvider for PendingNoiseConnectProvider {