diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index 138cea748c..f4456646ac 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -48,9 +48,8 @@ use tracing::warn; use codex_keyring_store::DefaultKeyringStore; use codex_keyring_store::KeyringStore; use rmcp::transport::auth::AuthorizationManager; -use rmcp::transport::auth::CredentialStore; use rmcp::transport::auth::InMemoryCredentialStore; -use rmcp::transport::auth::StoredCredentials; +use rmcp::transport::auth::OAuthState; use tokio::sync::Mutex; use tokio::sync::OwnedMutexGuard; @@ -140,8 +139,14 @@ fn load_oauth_tokens_from_keyring_with_fallback_to_file( Ok(None) => load_oauth_tokens_from_file(server_name, url), Err(error) => { warn!("failed to read OAuth tokens from keyring: {error}"); - load_oauth_tokens_from_file(server_name, url) - .with_context(|| format!("failed to read OAuth tokens from keyring: {error}")) + match load_oauth_tokens_from_file(server_name, url) { + Ok(Some(tokens)) => Ok(Some(tokens)), + Ok(None) => Err(error).context( + "failed to read OAuth tokens from keyring and no fallback credentials exist", + ), + Err(file_error) => Err(file_error) + .with_context(|| format!("failed to read OAuth tokens from keyring: {error}")), + } } } } @@ -297,6 +302,7 @@ struct OAuthPersistorInner { server_name: String, url: String, authorization_manager: Arc>, + oauth_metadata_client: reqwest::Client, store_mode: OAuthCredentialsStoreMode, last_credentials: Mutex>, } @@ -307,11 +313,30 @@ enum CredentialReload { Removed, } +pub(crate) async fn authorization_manager_from_tokens( + url: &str, + oauth_metadata_client: reqwest::Client, + tokens: &StoredOAuthTokens, +) -> Result { + let mut oauth_state = OAuthState::new(url.to_string(), Some(oauth_metadata_client)).await?; + oauth_state + .set_credentials(&tokens.client_id, tokens.token_response.0.clone()) + .await?; + + match oauth_state { + OAuthState::Authorized(manager) | OAuthState::Unauthorized(manager) => Ok(manager), + _ => Err(anyhow::anyhow!( + "unexpected OAuth state while replacing credentials" + )), + } +} + impl OAuthPersistor { pub(crate) fn new( server_name: String, url: String, authorization_manager: Arc>, + oauth_metadata_client: reqwest::Client, store_mode: OAuthCredentialsStoreMode, initial_credentials: Option, ) -> Self { @@ -320,6 +345,7 @@ impl OAuthPersistor { server_name, url, authorization_manager, + oauth_metadata_client, store_mode, last_credentials: Mutex::new(initial_credentials), }), @@ -481,29 +507,15 @@ impl OAuthPersistor { } async fn replace_manager_credentials(&self, tokens: &StoredOAuthTokens) -> Result<()> { - let token_response = tokens.token_response.0.clone(); - let store = InMemoryCredentialStore::new(); - store - .save(StoredCredentials::new( - tokens.client_id.clone(), - Some(token_response.clone()), - token_response - .scopes() - .map(|scopes| scopes.iter().map(|scope| scope.to_string()).collect()) - .unwrap_or_default(), - Some( - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_else(|_| Duration::from_secs(0)) - .as_secs(), - ), - )) - .await?; - + let replacement = authorization_manager_from_tokens( + &self.inner.url, + self.inner.oauth_metadata_client.clone(), + tokens, + ) + .await?; let manager = self.inner.authorization_manager.clone(); let mut guard = manager.lock().await; - guard.configure_client_id(&tokens.client_id)?; - guard.set_credential_store(store); + *guard = replacement; Ok(()) } @@ -944,17 +956,28 @@ mod tests { use keyring::Error as KeyringError; use pretty_assertions::assert_eq; use std::ffi::OsString; - use std::sync::Barrier; + use std::process::Child; + use std::process::Command; use std::sync::Mutex; use std::sync::MutexGuard; use std::sync::OnceLock; use std::sync::PoisonError; - use std::sync::mpsc; use std::thread; + use std::time::Instant; use tempfile::tempdir; + use wiremock::Mock; + use wiremock::MockServer; + use wiremock::ResponseTemplate; + use wiremock::matchers::body_string_contains; + use wiremock::matchers::method; + use wiremock::matchers::path; use codex_keyring_store::tests::MockKeyringStore; + const FALLBACK_WRITER_VARIANT_ENV: &str = "CODEX_TEST_FALLBACK_WRITER_VARIANT"; + const FALLBACK_WRITER_ATTEMPT_ENV: &str = "CODEX_TEST_FALLBACK_WRITER_ATTEMPT"; + const FALLBACK_WRITER_DONE_ENV: &str = "CODEX_TEST_FALLBACK_WRITER_DONE"; + struct TempCodexHome { _guard: MutexGuard<'static, ()>, _dir: tempfile::TempDir, @@ -1071,6 +1094,99 @@ mod tests { Ok(()) } + #[test] + fn load_oauth_tokens_preserves_keyring_error_without_fallback() -> Result<()> { + let _env = TempCodexHome::new(); + let store = MockKeyringStore::default(); + let tokens = sample_tokens(); + let key = super::compute_store_key(&tokens.server_name, &tokens.url)?; + store.set_error(&key, KeyringError::Invalid("error".into(), "load".into())); + + let error = super::load_oauth_tokens_from_keyring_with_fallback_to_file( + &store, + &tokens.server_name, + &tokens.url, + ) + .expect_err("missing fallback must not turn a keyring failure into logout"); + + assert!(error.to_string().contains( + "failed to read OAuth tokens from keyring and no fallback credentials exist" + )); + Ok(()) + } + + #[tokio::test] + #[expect( + clippy::await_holding_invalid_type, + reason = "the test must inspect and refresh the same authorization manager instance" + )] + async fn replacement_updates_current_scopes_used_by_scope_less_refresh() -> Result<()> { + let server = MockServer::start().await; + let server_url = format!("{}/mcp", server.uri()); + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": ["external-scope"], + }))) + .expect(2) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("refresh_token=refresh-token")) + .and(body_string_contains("scope=external-scope")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "refreshed-access-token", + "token_type": "Bearer", + "expires_in": 7200, + }))) + .expect(1) + .mount(&server) + .await; + + let mut initial = sample_tokens(); + initial.url = server_url.clone(); + let oauth_metadata_client = reqwest::Client::new(); + let manager = Arc::new(tokio::sync::Mutex::new( + super::authorization_manager_from_tokens( + &server_url, + oauth_metadata_client.clone(), + &initial, + ) + .await?, + )); + let runtime = OAuthPersistor::new( + initial.server_name.clone(), + server_url, + Arc::clone(&manager), + oauth_metadata_client, + OAuthCredentialsStoreMode::File, + Some(initial.clone()), + ); + let mut replacement = initial; + replacement + .token_response + .0 + .set_scopes(Some(vec![Scope::new("external-scope".to_string())])); + + runtime.replace_manager_credentials(&replacement).await?; + let guard = manager.lock().await; + assert_eq!( + guard.get_current_scopes().await, + vec!["external-scope".to_string()] + ); + guard.refresh_token().await?; + assert_eq!( + guard.get_current_scopes().await, + vec!["external-scope".to_string()] + ); + drop(guard); + server.verify().await; + Ok(()) + } + #[test] fn save_oauth_tokens_prefers_keyring_when_available() -> Result<()> { let _env = TempCodexHome::new(); @@ -1319,7 +1435,7 @@ mod tests { } #[test] - fn fallback_store_serializes_writes_for_different_servers() -> Result<()> { + fn fallback_store_serializes_cross_process_writes_for_different_servers() -> Result<()> { let _env = TempCodexHome::new(); let first = sample_tokens(); let mut second = sample_tokens(); @@ -1332,32 +1448,45 @@ mod tests { ); let fallback_lock = super::acquire_fallback_store_lock()?; - let barrier = Arc::new(Barrier::new(3)); - let (tx, rx) = mpsc::channel(); - let handles = [first.clone(), second.clone()] + let codex_home = find_codex_home()?; + let markers = tempdir()?; + let mut writers = ["first", "second"] .into_iter() - .map(|tokens| { - let barrier = Arc::clone(&barrier); - let tx = tx.clone(); - thread::spawn(move || { - barrier.wait(); - tx.send(super::save_oauth_tokens_to_file(&tokens)) - .expect("send save result"); - }) + .map(|variant| { + let attempt = markers.path().join(format!("{variant}-attempt")); + let done = markers.path().join(format!("{variant}-done")); + let child = Command::new(std::env::current_exe()?) + .args(["fallback_store_writer_child", "--ignored", "--nocapture"]) + .env("CODEX_HOME", codex_home.as_path()) + .env(FALLBACK_WRITER_VARIANT_ENV, variant) + .env(FALLBACK_WRITER_ATTEMPT_ENV, &attempt) + .env(FALLBACK_WRITER_DONE_ENV, &done) + .spawn()?; + Ok((child, attempt, done)) }) - .collect::>(); - drop(tx); - - barrier.wait(); - thread::sleep(Duration::from_millis(50)); - assert!(matches!(rx.try_recv(), Err(mpsc::TryRecvError::Empty))); + .collect::>>()?; + wait_for_paths( + &writers + .iter() + .map(|(_, attempt, _)| attempt.clone()) + .collect::>(), + Duration::from_secs(5), + )?; + thread::sleep(Duration::from_millis(100)); + for (child, _, done) in &mut writers { + assert!( + child.try_wait()?.is_none(), + "fallback writer should block on the cross-process file lock" + ); + assert!( + !done.exists(), + "fallback writer completed while another process held the file lock" + ); + } drop(fallback_lock); - for _ in 0..2 { - rx.recv_timeout(Duration::from_secs(5))??; - } - for handle in handles { - handle.join().expect("fallback writer should not panic"); + for (child, _, _) in &mut writers { + wait_for_child_success(child, Duration::from_secs(5))?; } let store = super::read_fallback_file()?.expect("fallback file should load"); @@ -1373,6 +1502,38 @@ mod tests { Ok(()) } + #[test] + #[ignore = "spawned by fallback_store_serializes_cross_process_writes_for_different_servers"] + fn fallback_store_writer_child() -> Result<()> { + let mut tokens = sample_tokens(); + match std::env::var(FALLBACK_WRITER_VARIANT_ENV)?.as_str() { + "first" => {} + "second" => { + tokens.server_name = "other-server".to_string(); + tokens.url = "https://other.example.test".to_string(); + tokens.token_response.0 = OAuthTokenResponse::new( + AccessToken::new("other-access-token".to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + } + variant => anyhow::bail!("unexpected fallback writer variant: {variant}"), + } + + let attempt = PathBuf::from( + std::env::var_os(FALLBACK_WRITER_ATTEMPT_ENV) + .context("fallback writer attempt path should be set")?, + ); + let done = PathBuf::from( + std::env::var_os(FALLBACK_WRITER_DONE_ENV) + .context("fallback writer done path should be set")?, + ); + fs::write(attempt, b"attempting")?; + super::save_oauth_tokens_to_file(&tokens)?; + fs::write(done, b"done")?; + Ok(()) + } + #[cfg(any(unix, windows))] #[test] fn server_generation_lock_is_stable_across_read_only_homes_and_tmpdirs() -> Result<()> { @@ -1460,6 +1621,31 @@ mod tests { ); } + fn wait_for_paths(paths: &[PathBuf], timeout: Duration) -> Result<()> { + let deadline = Instant::now() + timeout; + while !paths.iter().all(|path| path.exists()) { + if Instant::now() >= deadline { + anyhow::bail!("timed out waiting for fallback writer readiness"); + } + thread::sleep(Duration::from_millis(10)); + } + Ok(()) + } + + fn wait_for_child_success(child: &mut Child, timeout: Duration) -> Result<()> { + let deadline = Instant::now() + timeout; + loop { + if let Some(status) = child.try_wait()? { + anyhow::ensure!(status.success(), "fallback writer failed: {status}"); + return Ok(()); + } + if Instant::now() >= deadline { + anyhow::bail!("timed out waiting for fallback writer"); + } + thread::sleep(Duration::from_millis(10)); + } + } + fn sample_tokens() -> StoredOAuthTokens { let mut response = OAuthTokenResponse::new( AccessToken::new("access-token".to_string()), diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 90b09d724c..4a035f9fa2 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -47,7 +47,6 @@ use rmcp::service::{self}; use rmcp::transport::StreamableHttpClientTransport; use rmcp::transport::auth::AuthClient; use rmcp::transport::auth::AuthError; -use rmcp::transport::auth::OAuthState; use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig; use rmcp::transport::streamable_http_client::StreamableHttpError; use serde::Deserialize; @@ -66,6 +65,7 @@ use crate::in_process_transport::InProcessTransportFactory; use crate::load_oauth_tokens; use crate::oauth::OAuthPersistor; use crate::oauth::StoredOAuthTokens; +use crate::oauth::authorization_manager_from_tokens; use crate::stdio_server_launcher::StdioServerCommand; use crate::stdio_server_launcher::StdioServerLauncher; use crate::stdio_server_launcher::StdioServerProcessHandle; @@ -439,7 +439,6 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; let result = self .run_service_operation("tools/list", timeout, move |service| { let params = params.clone(); @@ -455,7 +454,6 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; let result = self .run_service_operation("tools/list", timeout, move |service| { let params = params.clone(); @@ -500,7 +498,6 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; let result = self .run_service_operation("resources/list", timeout, move |service| { let params = params.clone(); @@ -516,7 +513,6 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; let result = self .run_service_operation("resources/templates/list", timeout, move |service| { let params = params.clone(); @@ -532,7 +528,6 @@ impl RmcpClient { params: ReadResourceRequestParams, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; let result = self .run_service_operation("resources/read", timeout, move |service| { let params = params.clone(); @@ -550,7 +545,6 @@ impl RmcpClient { meta: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; let arguments = match arguments { Some(Value::Object(map)) => Some(map), Some(other) => { @@ -606,7 +600,6 @@ impl RmcpClient { method: &str, params: Option, ) -> Result<()> { - self.refresh_oauth_if_needed().await; self.run_service_operation( "notifications/custom", /*timeout*/ None, @@ -636,7 +629,6 @@ impl RmcpClient { method: &str, params: Option, ) -> Result { - self.refresh_oauth_if_needed().await; let response = self .run_service_operation("requests/custom", /*timeout*/ None, move |service| { let params = params.clone(); @@ -869,20 +861,21 @@ impl RmcpClient { Fut: std::future::Future>, { let service = self.service().await?; - match Self::run_service_operation_once( - Arc::clone(&service), - label, - timeout, - self.elicitation_pause_state.clone(), - &operation, - ) - .await + match self + .run_service_operation_once( + Arc::clone(&service), + label, + timeout, + self.elicitation_pause_state.clone(), + &operation, + ) + .await { Ok(result) => Ok(result), Err(error) if Self::is_session_expired_404(&error) => { self.reinitialize_after_session_expiry(&service).await?; let recovered_service = self.service().await?; - Self::run_service_operation_once( + self.run_service_operation_once( recovered_service, label, timeout, @@ -897,6 +890,7 @@ impl RmcpClient { } async fn run_service_operation_once( + &self, service: Arc>, label: &str, timeout: Option, @@ -907,17 +901,19 @@ impl RmcpClient { F: Fn(Arc>) -> Fut, Fut: std::future::Future>, { + let operation = async { + self.refresh_oauth_if_needed().await; + operation(service).await + }; match timeout { - Some(duration) => { - active_time_timeout(duration, pause_state.subscribe(), operation(service)) - .await - .map_err(|_| ClientOperationError::Timeout { - label: label.to_string(), - duration, - })? - .map_err(ClientOperationError::from) - } - None => operation(service).await.map_err(ClientOperationError::from), + Some(duration) => active_time_timeout(duration, pause_state.subscribe(), operation) + .await + .map_err(|_| ClientOperationError::Timeout { + label: label.to_string(), + duration, + })? + .map_err(ClientOperationError::from), + None => operation.await.map_err(ClientOperationError::from), } } @@ -1021,23 +1017,9 @@ async fn create_oauth_transport_and_runtime( // TODO(aibrahim): teach OAuth bootstrap and refresh to use the same // shared HTTP client abstraction instead of always creating the local // reqwest metadata client here. - let mut oauth_state = - OAuthState::new(url.to_string(), Some(oauth_metadata_client.clone())).await?; - - oauth_state - .set_credentials( - &initial_tokens.client_id, - initial_tokens.token_response.0.clone(), - ) - .await?; - - let manager = match oauth_state { - OAuthState::Authorized(manager) => manager, - OAuthState::Unauthorized(manager) => manager, - _ => { - return Err(anyhow!("unexpected OAuth state during client setup")); - } - }; + let manager = + authorization_manager_from_tokens(url, oauth_metadata_client.clone(), &initial_tokens) + .await?; let auth_client = AuthClient::new( StreamableHttpClientAdapter::new(http_client, default_headers, /*auth_provider*/ None), @@ -1054,6 +1036,7 @@ async fn create_oauth_transport_and_runtime( server_name.to_string(), url.to_string(), auth_manager, + oauth_metadata_client, credentials_store, Some(initial_tokens), ); diff --git a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs index b81613dbae..938d67867f 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs @@ -1,6 +1,7 @@ mod streamable_http_test_support; use std::collections::BTreeMap; +use std::sync::Arc; use std::time::Duration; use std::time::SystemTime; use std::time::UNIX_EPOCH; @@ -28,28 +29,218 @@ use wiremock::MockServer; use wiremock::Request; use wiremock::ResponseTemplate; use wiremock::matchers::body_string_contains; +use wiremock::matchers::header; use wiremock::matchers::method; use wiremock::matchers::path; use streamable_http_test_support::initialize_client; const SERVER_NAME: &str = "test-streamable-http-oauth-startup"; +const STARTUP_REFRESH_CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_STARTUP_REFRESH_SERVER_URL"; const EXTERNAL_UPDATE_CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_EXTERNAL_UPDATE_SERVER_URL"; +const OPERATION_TIMEOUT_CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_TIMEOUT_SERVER_URL"; +const OPERATION_TIMEOUT_REFRESH_STARTED_ENV: &str = "MCP_TEST_OAUTH_TIMEOUT_REFRESH_STARTED"; const REFRESH_CONTENTION_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_REFRESH_CONTENTION_SERVER_URL"; const REFRESH_CONTENTION_ATTEMPT_ENV: &str = "MCP_TEST_OAUTH_REFRESH_CONTENTION_ATTEMPT"; const REFRESH_CONTENTION_DONE_ENV: &str = "MCP_TEST_OAUTH_REFRESH_CONTENTION_DONE"; +const EXPIRED_ACCESS_TOKEN: &str = "expired-access-token"; +const VALID_REFRESH_TOKEN: &str = "valid-refresh-token"; +const REFRESHED_ACCESS_TOKEN: &str = "refreshed-access-token"; const OLD_ACCESS_TOKEN: &str = "old-access-token"; const OLD_REFRESH_TOKEN: &str = "old-refresh-token"; const EXTERNAL_ACCESS_TOKEN: &str = "external-access-token"; const EXTERNAL_REFRESH_TOKEN: &str = "external-refresh-token"; const EXTERNAL_EXPIRES_AT: u64 = u64::MAX; +const TIMEOUT_REFRESHED_ACCESS_TOKEN: &str = "timeout-refreshed-access-token"; const STALE_REFRESHED_ACCESS_TOKEN: &str = "stale-refreshed-access-token"; const STALE_ROTATED_REFRESH_TOKEN: &str = "stale-rotated-refresh-token"; +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result<()> { + let server = MockServer::start().await; + mount_oauth_metadata(&server, /*expected_calls*/ 1).await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={VALID_REFRESH_TOKEN}" + ))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": REFRESHED_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": VALID_REFRESH_TOKEN, + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {REFRESHED_ACCESS_TOKEN}"), + )) + .respond_with(|request: &Request| oauth_mcp_response(request, REFRESHED_ACCESS_TOKEN)) + .expect(2) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + let server_url = format!("{}/mcp", server.uri()); + let status = Command::new(std::env::current_exe()?) + .args([ + "oauth_startup_refresh_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(STARTUP_REFRESH_CHILD_SERVER_URL_ENV, server_url) + .status() + .await?; + assert!(status.success(), "OAuth startup child failed: {status}"); + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by refreshes_expired_persisted_token_before_initialize"] +async fn oauth_startup_refresh_child() -> anyhow::Result<()> { + let server_url = std::env::var(STARTUP_REFRESH_CHILD_SERVER_URL_ENV)?; + save_test_tokens( + &server_url, + EXPIRED_ACCESS_TOKEN, + VALID_REFRESH_TOKEN, + /*expires_at*/ 0, + )?; + initialized_oauth_client(&server_url).await?; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn operation_timeout_includes_oauth_generation_lock_wait() -> anyhow::Result<()> { + let server = MockServer::start().await; + mount_oauth_metadata(&server, /*expected_calls*/ 3).await; + let codex_home = TempDir::new()?; + let refresh_started = codex_home.path().join("refresh-started"); + let refresh_started_for_mock = refresh_started.clone(); + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={EXTERNAL_REFRESH_TOKEN}" + ))) + .respond_with(move |_request: &Request| { + std::fs::write(&refresh_started_for_mock, b"started").expect("write refresh marker"); + ResponseTemplate::new(200) + .set_delay(Duration::from_secs(1)) + .set_body_json(json!({ + "access_token": TIMEOUT_REFRESHED_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": EXTERNAL_REFRESH_TOKEN, + })) + }) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(|request: &Request| { + let Ok(body) = request.body_json::() else { + return ResponseTemplate::new(400); + }; + let expected_access_token = match ( + body.get("method").and_then(Value::as_str), + body.pointer("/params/name").and_then(Value::as_str), + ) { + (Some("tools/call"), Some("hold-generation-lock")) => { + TIMEOUT_REFRESHED_ACCESS_TOKEN + } + _ => OLD_ACCESS_TOKEN, + }; + oauth_mcp_response(request, expected_access_token) + }) + .expect(5) + .mount(&server) + .await; + + let server_url = format!("{}/mcp", server.uri()); + let status = Command::new(std::env::current_exe()?) + .args([ + "oauth_operation_timeout_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(OPERATION_TIMEOUT_CHILD_SERVER_URL_ENV, server_url) + .env(OPERATION_TIMEOUT_REFRESH_STARTED_ENV, refresh_started) + .status() + .await?; + assert!( + status.success(), + "OAuth operation timeout child failed: {status}" + ); + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by operation_timeout_includes_oauth_generation_lock_wait"] +async fn oauth_operation_timeout_child() -> anyhow::Result<()> { + let server_url = std::env::var(OPERATION_TIMEOUT_CHILD_SERVER_URL_ENV)?; + let refresh_started = + std::path::PathBuf::from(std::env::var(OPERATION_TIMEOUT_REFRESH_STARTED_ENV)?); + save_test_tokens( + &server_url, + OLD_ACCESS_TOKEN, + OLD_REFRESH_TOKEN, + future_expiry()?, + )?; + let lock_holder = Arc::new(initialized_oauth_client(&server_url).await?); + let timed_client = initialized_oauth_client(&server_url).await?; + save_test_tokens( + &server_url, + EXTERNAL_ACCESS_TOKEN, + EXTERNAL_REFRESH_TOKEN, + /*expires_at*/ 0, + )?; + + let lock_holder_call = tokio::spawn(async move { + lock_holder + .call_tool( + "hold-generation-lock".to_string(), + /*arguments*/ None, + /*meta*/ None, + Some(Duration::from_secs(5)), + ) + .await + }); + wait_for_path(&refresh_started).await?; + let error = timed_client + .call_tool( + "must-time-out-on-generation-lock".to_string(), + /*arguments*/ None, + /*meta*/ None, + Some(Duration::from_millis(100)), + ) + .await + .expect_err("generation lock wait must consume the operation timeout"); + assert!( + error + .to_string() + .contains("timed out awaiting tools/call after 100ms"), + "unexpected timeout error: {error}" + ); + lock_holder_call.await??; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn running_client_adopts_external_oauth_update_and_deletion() -> anyhow::Result<()> { let server = MockServer::start().await; - mount_oauth_metadata(&server).await; + mount_oauth_metadata(&server, /*expected_calls*/ 2).await; Mock::given(method("POST")) .and(path("/mcp")) @@ -168,7 +359,7 @@ async fn oauth_external_delete_writer_child() -> anyhow::Result<()> { async fn refresh_holds_generation_lock_against_external_oauth_update() -> anyhow::Result<()> { let server = MockServer::start().await; let server_url = format!("{}/mcp", server.uri()); - mount_oauth_metadata(&server).await; + mount_oauth_metadata(&server, /*expected_calls*/ 3).await; Mock::given(method("POST")) .and(path("/oauth/token")) .and(body_string_contains("grant_type=refresh_token")) @@ -322,7 +513,7 @@ async fn oauth_external_update_writer_child() -> anyhow::Result<()> { Ok(()) } -async fn mount_oauth_metadata(server: &MockServer) { +async fn mount_oauth_metadata(server: &MockServer, expected_calls: u64) { Mock::given(method("GET")) .and(path("/.well-known/oauth-authorization-server/mcp")) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ @@ -330,7 +521,7 @@ async fn mount_oauth_metadata(server: &MockServer) { "token_endpoint": format!("{}/oauth/token", server.uri()), "scopes_supported": [""], }))) - .expect(1) + .expect(expected_calls) .mount(server) .await; }