Refresh MCP runtimes when session auth changes (#35146)

## What changed

- Track the authentication state used by each published MCP runtime and mark the runtime dirty when the auth mode, account details, or token changes.
- Refresh plugin auth mode before evaluating pending MCP work, and rebuild the initial MCP projection if authentication changes during session startup.
- Carry one auth snapshot through runtime construction so plugin discovery and MCP connections use consistent credentials.

## Testing

- Add a regression test that replaces an API key, reloads authentication, and verifies that MCP refresh publishes a runtime using the new credentials.

GitOrigin-RevId: 219c2f7a157f53b5467116e18bfe70a8c3dca695

codex-rs/codex-mcp/src/runtime.rs

@@ -76,6 +76,8 @@ pub struct McpRuntime {
 struct PublishedMcpRuntime {
     connections: Arc<McpConnectionSet>,
     config: Option<Arc<McpConfig>>,
+    auth: Option<CodexAuth>,
+    auth_token: Option<String>,
     plugins_available: bool,
     ready_selected_capability_roots: Vec<SelectedCapabilityRoot>,
 }
@@ -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<SelectedCapabilityRoot> {
         self.current.load().ready_selected_capability_roots.clone()
     }

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,

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<Config>,
+    pub(super) auth: Option<CodexAuth>,
     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<CodexAuth>,
+    ) -> 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<Self>,
         session_configuration: &SessionConfiguration,
+        auth: Option<CodexAuth>,
         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<ElicitationReviewerHandle>,
     ) {
-        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<ElicitationReviewerHandle>,
     ) -> 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

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,

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 {