From dd22460869a9649219900d7c099d8cd4e70124b2 Mon Sep 17 00:00:00 2001 From: jif Date: Mon, 10 Aug 2026 18:33:37 +0000 Subject: [PATCH] Add MCP OAuth credential contention regression tests (#37866) ## What changed - Cover non-blocking credential probes when the file or secrets store is locked, including retaining only a matching prior credential snapshot and recovering after the lock is released. - Extend the streamable HTTP OAuth round trip to verify that user turns continue during store contention, newly discovered servers recover after contention, and logged-out servers do not reuse authenticated connections. - Cover reconciliation of authentication failures and effective `Authorization` headers. GitOrigin-RevId: 8e377644001bd0be04dc9f79841080094d09d199 --- .../codex-mcp/src/connection_manager_tests.rs | 68 +++++++ codex-rs/core/tests/suite/rmcp_client.rs | 190 +++++++++++++++++- .../src/oauth/tests/store_lock_tests.rs | 69 +++++++ 3 files changed, 322 insertions(+), 5 deletions(-) diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 94565ace66..61678083b6 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -4075,6 +4075,74 @@ async fn reconciliation_reuses_an_unchanged_ready_server() { ); } +#[tokio::test] +async fn reconciliation_retries_non_oauth_authentication_failures() { + let runtime_context = reusable_server_runtime_context(); + let config = reusable_server_config("http://127.0.0.1:1"); + let mut previous = + manager_with_reusable_ready_server(&config, &runtime_context, Vec::new()).await; + let connection = + Arc::get_mut(&mut previous.servers.get_mut("docs").expect("server").connection) + .expect("test server has one connection owner"); + connection.client.client = futures::future::ready(Err(StartupOutcomeError::Failed { + error: "bearer token rejected".to_string(), + is_authentication_required: true, + })) + .boxed() + .shared(); + + let reconciled = reconcile_reusable_server(&previous, config, runtime_context).await; + + assert!(!previous.shares_test_connection_with(&reconciled, "docs")); +} + +#[test] +fn connection_identity_uses_effective_authorization_headers() { + let runtime_context = reusable_server_runtime_context(); + let missing_env_var = format!("CODEX_TEST_UNSET_MCP_AUTHORIZATION_{}", std::process::id()); + assert!(std::env::var_os(&missing_env_var).is_none()); + + for (static_header, environment_header, has_authorization) in [ + (Some("Bearer configured-token"), None, true), + (Some("invalid\nheader"), None, false), + (None, Some("PATH"), true), + (None, Some(missing_env_var.as_str()), false), + ] { + let mut config = reusable_server_config("http://127.0.0.1:1"); + config.transport = McpServerTransportConfig::StreamableHttp { + url: "http://127.0.0.1:1".to_string(), + bearer_token_env_var: None, + http_headers: static_header + .map(|value| HashMap::from([("aUtHoRiZaTiOn".to_string(), value.to_string())])), + env_http_headers: environment_header + .map(|value| HashMap::from([("aUtHoRiZaTiOn".to_string(), value.to_string())])), + }; + let server = EffectiveMcpServer::configured(config); + let identity = |keyring_backend_kind| { + McpServerConnectionIdentity::new( + "docs", + &server, + OAuthCredentialsStoreMode::File, + keyring_backend_kind, + &Ok(None), + &runtime_context, + /*runtime_auth_provider*/ None, + /*auth*/ None, + /*codex_apps_cache_identity*/ None, + ElicitationCapability::default(), + ClientMcpExtensions::default(), + /*previous_identity*/ None, + ) + }; + + assert_eq!( + identity(AuthKeyringBackendKind::Direct) + .has_same_connection_config(&identity(AuthKeyringBackendKind::Secrets)), + has_authorization, + ); + } +} + #[tokio::test] async fn reconciliation_reuses_legacy_stdio_server_with_existing_protocol_marker() { let runtime_context = McpRuntimeContext::new( diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index 7085cd29ca..d23db10004 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -3076,8 +3076,28 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> { let call_id = "call-789"; let server_name = "rmcp_http_oauth"; + let discovered_server_name = "rmcp_http_oauth_discovered"; let namespace = format!("mcp__{server_name}"); + mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_assistant_message("msg-contended", "OAuth sign-in is still pending."), + responses::ev_completed("resp-contended"), + ]), + ) + .await; + mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_assistant_message( + "msg-discovered", + "A newly discovered OAuth server is still starting.", + ), + responses::ev_completed("resp-discovered"), + ]), + ) + .await; let response_mock = mount_sse_once( &server, responses::sse(vec![ @@ -3121,6 +3141,11 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> { // server so the test does not share credentials with other suite cases. let temp_home = Arc::new(tempdir()?); let _codex_home_guard = EnvVarGuard::set("CODEX_HOME", temp_home.path().as_os_str()); + let unset_authorization_env_var = format!( + "CODEX_TEST_UNSET_MCP_OAUTH_AUTHORIZATION_{}", + std::process::id() + ); + assert!(std::env::var_os(&unset_authorization_env_var).is_none()); let environment_id = remote_aware_environment_id(); let credential_config: McpServerConfig = serde_json::from_value(json!({ "url": &server_url, @@ -3135,6 +3160,16 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> { refresh_token, OAuthCredentialExpiry::Expired, )?; + let discovered_credential_name = + credential_config.oauth_credential_name(discovered_server_name); + write_fallback_oauth_tokens( + discovered_credential_name.as_ref(), + &server_url, + client_id, + expected_token, + refresh_token, + OAuthCredentialExpiry::Valid, + )?; // Phase 4: configure Codex with the OAuth-backed Streamable HTTP MCP // server and build the fixture in the active local or remote-aware mode. @@ -3142,10 +3177,7 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> { .with_model_info_override("gpt-5.4", |model| model.supports_search_tool = false) .with_home(temp_home.clone()) .with_config(move |config| { - // Keep OAuth credentials isolated to this test home because Bazel - // runs the full core suite in one process. - config.mcp_oauth_credentials_store_mode = serde_json::from_value(json!("file")) - .expect("`file` should deserialize as OAuthCredentialsStoreMode"); + config.mcp_oauth_credentials_store_mode = OAuthCredentialsStoreMode::Auto; insert_mcp_server( config, server_name, @@ -3153,7 +3185,10 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> { url: server_url, bearer_token_env_var: None, http_headers: None, - env_http_headers: None, + env_http_headers: Some(HashMap::from([( + "Authorization".to_string(), + unset_authorization_env_var, + )])), }, TestMcpServerOptions { environment_id, @@ -3184,6 +3219,151 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> { failure_reason, Some(McpStartupFailureReason::ReauthenticationRequired), ); + + let store_lock = fs::OpenOptions::new() + .read(true) + .write(true) + .open(temp_home.path().join("mcp-oauth-locks/file-store.lock"))?; + store_lock.try_lock()?; + fixture + .codex + .submit(read_only_user_turn( + &fixture, + "continue while OAuth credentials are locked", + )) + .await?; + let (contended_turn, refreshed_failure_reason, refreshed_failed_servers) = + tokio::time::timeout(Duration::from_secs(5), async { + let mut contended_turn = None; + let mut refreshed_failure_reason = None; + let mut refreshed_failed_servers = None; + let mut refreshed_starting = false; + while contended_turn.is_none() + || refreshed_failure_reason.is_none() + || refreshed_failed_servers.is_none() + { + match fixture.codex.next_event().await?.msg { + EventMsg::McpStartupUpdate(update) if update.server == server_name => { + match update.status { + McpStartupStatus::Starting => refreshed_starting = true, + McpStartupStatus::Failed { reason, .. } => { + assert!(refreshed_starting); + refreshed_failure_reason = Some(reason); + } + McpStartupStatus::Ready | McpStartupStatus::Cancelled => {} + } + } + EventMsg::McpStartupComplete(startup) => { + refreshed_failed_servers = Some(startup.failed); + } + EventMsg::TurnComplete(turn) => contended_turn = Some(turn), + _ => {} + } + } + Ok::<_, anyhow::Error>(( + contended_turn.expect("turn completion was observed"), + refreshed_failure_reason.expect("failure status was observed"), + refreshed_failed_servers.expect("startup summary was observed"), + )) + }) + .await + .context("OAuth credential-store contention blocked the user turn or startup status")??; + assert_eq!( + refreshed_failure_reason, + Some(McpStartupFailureReason::ReauthenticationRequired), + ); + assert_eq!( + refreshed_failed_servers + .into_iter() + .map(|failure| failure.server) + .collect::>(), + vec![server_name.to_string()], + ); + assert!( + contended_turn.error.is_none(), + "the user turn failed while OAuth credentials were locked: {:?}", + contended_turn.error, + ); + assert_eq!( + contended_turn.last_agent_message.as_deref(), + Some("OAuth sign-in is still pending."), + ); + + let mut refreshed_config = fixture.config.clone(); + let mut refreshed_servers = refreshed_config.mcp_servers.get().clone(); + let discovered_server = refreshed_servers + .get(server_name) + .cloned() + .expect("the initial OAuth server should be configured"); + refreshed_servers.insert(discovered_server_name.to_string(), discovered_server); + refreshed_config + .mcp_servers + .set(refreshed_servers) + .expect("test MCP servers should accept the discovered OAuth server"); + let discovered_turn = tokio::time::timeout(Duration::from_secs(5), async { + fixture + .codex + .refresh_runtime_config(refreshed_config.clone()) + .await; + fixture + .codex + .submit(read_only_user_turn( + &fixture, + "continue while a newly discovered OAuth server is starting", + )) + .await?; + loop { + if let EventMsg::TurnComplete(turn) = fixture.codex.next_event().await?.msg { + return Ok::<_, anyhow::Error>(turn); + } + } + }) + .await + .context("a newly discovered OAuth server blocked the user turn during store contention")??; + assert!( + discovered_turn.error.is_none(), + "the user turn failed while a newly discovered OAuth server was locked: {:?}", + discovered_turn.error, + ); + assert_eq!( + discovered_turn.last_agent_message.as_deref(), + Some("A newly discovered OAuth server is still starting."), + ); + drop(store_lock); + + tokio::time::timeout( + Duration::from_secs(5), + wait_for_mcp_server(&fixture.codex, discovered_server_name), + ) + .await + .context("the newly discovered OAuth server did not recover after its store was unlocked")??; + + assert!(codex_rmcp_client::delete_oauth_tokens( + discovered_credential_name.as_ref(), + http_server.url(), + OAuthCredentialsStoreMode::File, + codex_config::types::AuthKeyringBackendKind::default(), + )?); + fixture.codex.refresh_runtime_config(refreshed_config).await; + let logged_out_startup = tokio::time::timeout( + Duration::from_secs(5), + wait_for_event(&fixture.codex, |event| { + matches!(event, EventMsg::McpStartupComplete(_)) + }), + ) + .await + .context("MCP startup did not finish after the discovered OAuth server was logged out")?; + let EventMsg::McpStartupComplete(logged_out_startup) = logged_out_startup else { + unreachable!("event guard guarantees McpStartupComplete"); + }; + assert!( + logged_out_startup + .failed + .iter() + .any(|failure| failure.server == discovered_server_name), + "the logged-out MCP server reused its authenticated connection: {logged_out_startup:?}", + ); + write_fallback_oauth_tokens( credential_name.as_ref(), http_server.url(), diff --git a/codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs b/codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs index 037f05cff4..a78a47b1c6 100644 --- a/codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs +++ b/codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs @@ -28,6 +28,7 @@ use super::OAuthStore; use super::OAuthStoreLock; use super::OAuthStoreLockFailure; use super::OAuthStoreLockMode; +use crate::oauth::StoredOAuthCredentialSnapshot; use crate::oauth::StoredOAuthTokens; use crate::oauth::WrappedOAuthTokenResponse; use crate::oauth::fallback_file_path; @@ -39,6 +40,7 @@ use crate::oauth::save_oauth_tokens_to_file_with_lock_held; use crate::oauth::save_oauth_tokens_to_secrets_keyring_with_lock_held; use crate::oauth::save_oauth_tokens_with_keyring; use crate::oauth::save_oauth_tokens_with_keyring_with_fallback_to_file; +use crate::oauth::stored_oauth_credential_snapshot; use crate::oauth::test_support::TempCodexHome; use codex_config::types::OAuthCredentialsStoreMode; @@ -279,6 +281,73 @@ fn auto_load_secrets_lock_failure_does_not_fall_back_to_file() -> Result<()> { Ok(()) } +#[test] +fn oauth_credential_probes_skip_contended_file_and_secrets_stores() -> Result<()> { + let _env = TempCodexHome::new(); + let tokens = sample_tokens(); + let file = OAuthCredentialsStoreMode::File; + let auto = OAuthCredentialsStoreMode::Auto; + let direct = AuthKeyringBackendKind::Direct; + let secrets = AuthKeyringBackendKind::Secrets; + save_oauth_tokens_to_file(&tokens)?; + let snapshot = + stored_oauth_credential_snapshot(&tokens.server_name, &tokens.url, file, direct)? + .expect("fallback credentials should exist before contention"); + let reload = |store_mode, keyring_backend_kind| { + snapshot.reload( + &tokens.server_name, + &tokens.url, + store_mode, + keyring_backend_kind, + ) + }; + let probe = |store_mode, keyring_backend_kind| { + StoredOAuthCredentialSnapshot::for_runtime_refresh( + /*previous*/ None, + &tokens.server_name, + &tokens.url, + store_mode, + keyring_backend_kind, + ) + }; + let refresh = |url| { + StoredOAuthCredentialSnapshot::for_runtime_refresh( + Some(&snapshot), + &tokens.server_name, + url, + file, + direct, + ) + }; + + { + let _lock = OAuthStoreLock::acquire_for_write(OAuthStore::File)?; + assert_eq!(reload(file, direct)?, None); + assert_eq!(reload(auto, secrets)?, None); + assert_eq!(probe(file, direct)?, None); + let contended = refresh(&tokens.url)?.expect("the previous snapshot should be retained"); + assert!(contended.store_was_contended()); + assert_eq!(contended, snapshot); + assert_eq!( + refresh("https://another.example.test")?, + None, + "store contention must not replay credentials for another endpoint", + ); + } + + assert_eq!(reload(file, direct)?, Some(snapshot.credentials().clone())); + let refreshed = refresh(&tokens.url)?.expect("the unlocked store should still contain tokens"); + assert!(!refreshed.store_was_contended()); + + let _lock = OAuthStoreLock::acquire_for_write(OAuthStore::Secrets)?; + assert_eq!( + probe(auto, secrets)?, + None, + "a locked Secrets authority must not fall back to the stale File entry", + ); + Ok(()) +} + struct LockContentionSubscriber { contended_tx: mpsc::Sender<()>, }