Prepare MCP connections for coordinated OAuth refresh (#42128)

## What changed

- Add the under-development `mcp_oauth_refresh_coordination` feature and pin its selected refresh mode to each MCP connection.
- Treat the refresh mode as part of OAuth connection identity so runtime and MCP configuration refreshes reconnect when the mode changes.
- Keep legacy refresh behavior in builds where coordination is unavailable and emit a warning when coordinated mode is selected.

## Testing

- Cover connection replacement when the feature is toggled through both runtime and MCP configuration refresh paths.

GitOrigin-RevId: 89596150f0338e6ede0e712badecea4450831766
This commit is contained in:
Matthew Zeng
2026-09-01 20:26:57 +00:00
committed by copyberry
parent 7730edf62c
commit bd89ce67cc
15 changed files with 232 additions and 3 deletions
@@ -231,6 +231,7 @@ impl McpConnectionSet {
} = input;
let store_mode = config.mcp_oauth_credentials_store_mode;
let keyring_backend_kind = config.auth_keyring_backend_kind;
let oauth_refresh_mode = config.oauth_refresh_mode;
let codex_home = config.codex_home.clone();
let prefix_mcp_tool_names = config.prefix_mcp_tool_names;
let non_prefixed_mcp_tool_servers = config.non_prefixed_mcp_tool_servers.clone();
@@ -362,6 +363,7 @@ impl McpConnectionSet {
server: server.clone(),
store_mode,
keyring_backend_kind,
oauth_refresh_mode,
runtime_context: runtime_context.clone(),
resolved_environment: resolved_environment.clone(),
auth_provider: runtime_auth_provider.clone(),
@@ -377,6 +379,7 @@ impl McpConnectionSet {
host_plugin_root,
store_mode,
keyring_backend_kind,
oauth_refresh_mode,
&resolved_environment,
&runtime_context,
runtime_auth_provider.as_ref(),
@@ -539,6 +542,7 @@ impl McpConnectionSet {
server,
store_mode,
keyring_backend_kind,
oauth_refresh_mode,
cancel_token.clone(),
tx_event.clone(),
elicitation_requests.clone(),
@@ -57,6 +57,7 @@ use codex_rmcp_client::ElicitationResponse;
use codex_rmcp_client::InProcessTransportFactory;
use codex_rmcp_client::McpAuthState;
use codex_rmcp_client::McpLoginRequirement;
use codex_rmcp_client::McpOAuthRefreshMode;
use codex_rmcp_client::RmcpClient;
use codex_utils_path_uri::PathUri;
use futures::FutureExt;
@@ -4153,6 +4154,7 @@ async fn executor_owned_chatgpt_mcp_accepts_only_safe_explicit_authorization() -
/*host_plugin_root*/ None,
OAuthCredentialsStoreMode::File,
keyring_backend_kind,
McpOAuthRefreshMode::Legacy,
&resolved_environment,
&runtime_context,
/*runtime_auth_provider*/ None,
@@ -4675,6 +4677,7 @@ fn reusable_server_identity(
/*host_plugin_root*/ None,
OAuthCredentialsStoreMode::default(),
AuthKeyringBackendKind::default(),
McpOAuthRefreshMode::Legacy,
&resolved_environment,
runtime_context,
/*runtime_auth_provider*/ None,
@@ -5054,6 +5057,7 @@ fn connection_identity_uses_effective_authorization_headers() {
(Some("invalid\nheader"), None, false),
(None, Some("PATH"), true),
(None, Some(missing_env_var.as_str()), false),
(None, None, false),
] {
let mut config = reusable_server_config("http://127.0.0.1:1");
config.transport = McpServerTransportConfig::StreamableHttp {
@@ -5066,13 +5070,14 @@ fn connection_identity_uses_effective_authorization_headers() {
http_headers_helper: None,
};
let server = EffectiveMcpServer::configured(config);
let identity = |keyring_backend_kind| {
let identity = |keyring_backend_kind, oauth_refresh_mode| {
McpServerConnectionIdentity::new(
"docs",
&server,
/*host_plugin_root*/ None,
OAuthCredentialsStoreMode::File,
keyring_backend_kind,
oauth_refresh_mode,
&Ok(None),
&runtime_context,
/*runtime_auth_provider*/ None,
@@ -5085,8 +5090,19 @@ fn connection_identity_uses_effective_authorization_headers() {
};
assert_eq!(
identity(AuthKeyringBackendKind::Direct)
.has_same_connection_config(&identity(AuthKeyringBackendKind::Secrets)),
identity(AuthKeyringBackendKind::Direct, McpOAuthRefreshMode::Legacy)
.has_same_connection_config(&identity(
AuthKeyringBackendKind::Secrets,
McpOAuthRefreshMode::Legacy,
)),
has_authorization,
);
assert_eq!(
identity(AuthKeyringBackendKind::Direct, McpOAuthRefreshMode::Legacy)
.has_same_connection_config(&identity(
AuthKeyringBackendKind::Direct,
McpOAuthRefreshMode::Coordinated,
)),
has_authorization,
);
}
@@ -5472,6 +5488,7 @@ async fn reconciliation_reconnects_when_host_plugin_root_changes() {
Some(&original_root),
OAuthCredentialsStoreMode::default(),
AuthKeyringBackendKind::default(),
McpOAuthRefreshMode::Legacy,
&resolved_environment,
&runtime_context,
/*runtime_auth_provider*/ None,
@@ -5550,6 +5567,7 @@ async fn connection_identity_distinguishes_accounts_with_the_same_token() -> any
/*host_plugin_root*/ None,
OAuthCredentialsStoreMode::default(),
AuthKeyringBackendKind::default(),
McpOAuthRefreshMode::Legacy,
&Ok(None),
&runtime_context,
Some(&provider),
@@ -5603,6 +5621,7 @@ async fn connection_identity_distinguishes_agent_account_runtime_and_task() -> a
/*host_plugin_root*/ None,
OAuthCredentialsStoreMode::default(),
AuthKeyringBackendKind::default(),
McpOAuthRefreshMode::Legacy,
&Ok(None),
&runtime_context,
Some(&provider),
+3
View File
@@ -12,6 +12,7 @@ use codex_login::AuthManager;
use codex_login::CodexAuth;
use codex_protocol::mcp::ClientMcpExtensions;
use codex_rmcp_client::ElicitationResponse;
use codex_rmcp_client::McpOAuthRefreshMode;
use rmcp::model::ElicitationAction;
use rmcp::model::ElicitationCapability;
use serde_json::Map;
@@ -31,6 +32,7 @@ pub(crate) struct EventStreamConnectionSettings {
pub server: EffectiveMcpServer,
pub store_mode: OAuthCredentialsStoreMode,
pub keyring_backend_kind: AuthKeyringBackendKind,
pub oauth_refresh_mode: McpOAuthRefreshMode,
pub runtime_context: McpRuntimeContext,
pub resolved_environment: std::result::Result<Option<Arc<Environment>>, String>,
pub auth_provider: Option<SharedAuthProvider>,
@@ -81,6 +83,7 @@ impl McpEventStreamOpener {
connection.server.clone(),
connection.store_mode,
connection.keyring_backend_kind,
connection.oauth_refresh_mode,
connection.runtime_context.clone(),
connection.resolved_environment.clone(),
connection.auth_provider.clone(),
+3
View File
@@ -41,6 +41,7 @@ use codex_protocol::mcp::Tool;
use codex_protocol::models::PermissionProfile;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::McpAuthStatus;
use codex_rmcp_client::McpOAuthRefreshMode;
use codex_utils_path_uri::PathUri;
use rmcp::model::ElicitationCapability;
use rmcp::model::ReadResourceRequestParams;
@@ -128,6 +129,8 @@ pub struct McpConfig {
pub codex_home: PathBuf,
/// Preferred credential store for MCP OAuth tokens.
pub mcp_oauth_credentials_store_mode: OAuthCredentialsStoreMode,
/// OAuth refresh ownership selected for new MCP connections.
pub oauth_refresh_mode: McpOAuthRefreshMode,
/// Backend used when MCP OAuth storage is configured for keyring-backed persistence.
pub auth_keyring_backend_kind: AuthKeyringBackendKind,
/// Optional fixed localhost callback port for MCP OAuth login.
+1
View File
@@ -24,6 +24,7 @@ pub(crate) fn test_mcp_config(codex_home: PathBuf) -> McpConfig {
apps_mcp_product_sku: None,
codex_home,
mcp_oauth_credentials_store_mode: OAuthCredentialsStoreMode::default(),
oauth_refresh_mode: McpOAuthRefreshMode::Legacy,
auth_keyring_backend_kind: AuthKeyringBackendKind::default(),
mcp_oauth_callback_port: None,
mcp_oauth_callback_url: None,
+12
View File
@@ -60,6 +60,7 @@ use codex_protocol::protocol::McpStartupStatus;
use codex_protocol::protocol::McpStartupUpdateEvent;
use codex_rmcp_client::ExecutorStdioServerLauncher;
use codex_rmcp_client::LocalStdioServerLauncher;
use codex_rmcp_client::McpOAuthRefreshMode;
use codex_rmcp_client::McpProtocolMode;
use codex_rmcp_client::RmcpClient;
use codex_rmcp_client::StdioServerLauncher;
@@ -283,6 +284,7 @@ struct ManagedClientStartup {
server: EffectiveMcpServer,
store_mode: OAuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
oauth_refresh_mode: McpOAuthRefreshMode,
tx_event: Option<Sender<Event>>,
elicitation_requests: ElicitationRequestManager,
codex_apps_tools_cache_context: Option<ConnectorRuntimeContext<ToolInfo>>,
@@ -305,6 +307,7 @@ impl ManagedClientStartup {
server,
store_mode,
keyring_backend_kind,
oauth_refresh_mode,
tx_event,
elicitation_requests,
codex_apps_tools_cache_context,
@@ -342,6 +345,7 @@ impl ManagedClientStartup {
server.clone(),
store_mode,
keyring_backend_kind,
oauth_refresh_mode,
runtime_context,
resolved_environment,
runtime_auth_provider,
@@ -423,6 +427,7 @@ impl AsyncManagedClient {
server: EffectiveMcpServer,
store_mode: OAuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
oauth_refresh_mode: McpOAuthRefreshMode,
cancel_token: CancellationToken,
tx_event: Option<Sender<Event>>,
elicitation_requests: ElicitationRequestManager,
@@ -452,6 +457,7 @@ impl AsyncManagedClient {
server,
store_mode,
keyring_backend_kind,
oauth_refresh_mode,
tx_event,
elicitation_requests,
codex_apps_tools_cache_context: codex_apps_tools_cache_context.clone(),
@@ -1072,11 +1078,17 @@ pub(crate) async fn make_rmcp_client(
server: EffectiveMcpServer,
store_mode: OAuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
oauth_refresh_mode: McpOAuthRefreshMode,
runtime_context: McpRuntimeContext,
resolved_environment: std::result::Result<Option<Arc<Environment>>, String>,
runtime_auth_provider: Option<SharedAuthProvider>,
protocol_mode: McpProtocolMode,
) -> Result<RmcpClient, StartupOutcomeError> {
if oauth_refresh_mode == McpOAuthRefreshMode::Coordinated {
warn!(
"MCP OAuth refresh coordination is not available in this build; using legacy refresh"
);
}
let config = server.config().clone();
if matches!(config.auth, McpServerAuth::ChatGpt)
&& !config.is_local_environment()
+5
View File
@@ -15,6 +15,7 @@ use codex_connectors::ConnectorRuntimeContextKey;
use codex_exec_server::Environment;
use codex_login::CodexAuth;
use codex_protocol::mcp::ClientMcpExtensions;
use codex_rmcp_client::McpOAuthRefreshMode;
use codex_rmcp_client::StoredOAuthCredentialSnapshot;
use codex_rmcp_client::StoredOAuthTokens;
use codex_utils_path_uri::PathUri;
@@ -99,6 +100,7 @@ pub(crate) struct McpServerConnectionIdentity {
environment_id: String,
host_plugin_root: Option<PathUri>,
oauth_store: Option<(OAuthCredentialsStoreMode, AuthKeyringBackendKind)>,
oauth_refresh_mode: Option<McpOAuthRefreshMode>,
oauth_credentials: Result<Option<StoredOAuthCredentialSnapshot>, String>,
pub(crate) oauth_store_was_contended: bool,
resolved_environment: Result<Option<Arc<Environment>>, String>,
@@ -120,6 +122,7 @@ impl McpServerConnectionIdentity {
host_plugin_root: Option<&PathUri>,
store_mode: OAuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
oauth_refresh_mode: McpOAuthRefreshMode,
resolved_environment: &Result<Option<Arc<Environment>>, String>,
runtime_context: &McpRuntimeContext,
runtime_auth_provider: Option<&SharedAuthProvider>,
@@ -217,6 +220,7 @@ impl McpServerConnectionIdentity {
oauth_store: stored_oauth_url
.is_some()
.then_some((store_mode, keyring_backend_kind)),
oauth_refresh_mode: stored_oauth_url.is_some().then_some(oauth_refresh_mode),
oauth_credentials,
oauth_store_was_contended,
resolved_environment: resolved_environment.clone(),
@@ -250,6 +254,7 @@ impl McpServerConnectionIdentity {
&& self.environment_id == other.environment_id
&& self.host_plugin_root == other.host_plugin_root
&& self.oauth_store == other.oauth_store
&& self.oauth_refresh_mode == other.oauth_refresh_mode
&& same_resolved_environment(&self.resolved_environment, &other.resolved_environment)
&& self.local_stdio_fallback_cwd == other.local_stdio_fallback_cwd
&& self.referenced_environment_variables == other.referenced_environment_variables