diff --git a/codex-rs/rmcp-client/src/http_client_adapter.rs b/codex-rs/rmcp-client/src/http_client_adapter.rs index 19befb62355e..75bb42a5e3e8 100644 --- a/codex-rs/rmcp-client/src/http_client_adapter.rs +++ b/codex-rs/rmcp-client/src/http_client_adapter.rs @@ -22,6 +22,7 @@ use codex_exec_server::HttpResponseBodyStream; use futures::StreamExt; use futures::stream; use futures::stream::BoxStream; +use oauth2::AccessToken; use reqwest::StatusCode; use reqwest::header::ACCEPT; use reqwest::header::AUTHORIZATION; @@ -55,12 +56,15 @@ pub(crate) struct StreamableHttpClientAdapter { http_client: Arc, default_headers: HeaderMap, auth_provider: Option, + attribute_rejected_access_token: bool, } #[derive(Debug, thiserror::Error)] pub(crate) enum StreamableHttpClientAdapterError { #[error("streamable HTTP session expired with 404 Not Found")] SessionExpired404, + #[error("MCP server rejected the access token with HTTP 401 Unauthorized")] + AccessTokenRejected { rejected_access_token: AccessToken }, #[error(transparent)] HttpRequest(#[from] ExecServerError), #[error("invalid HTTP header: {0}")] @@ -77,8 +81,15 @@ impl StreamableHttpClientAdapter { http_client, default_headers, auth_provider, + attribute_rejected_access_token: false, } } + + /// Preserves the access token associated with a 401 for Codex-owned OAuth recovery. + pub(crate) fn with_rejected_token_attribution(mut self) -> Self { + self.attribute_rejected_access_token = true; + self + } } impl StreamableHttpClient for StreamableHttpClientAdapter { @@ -109,7 +120,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { JSON_MIME_TYPE.to_string(), StreamableHttpClientAdapterError::Header, )?; - if let Some(auth_token) = auth_token { + if let Some(auth_token) = auth_token.as_deref() { insert_header( &mut headers, AUTHORIZATION, @@ -162,6 +173,12 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { StreamableHttpClientAdapterError::SessionExpired404, )); } + if self.attribute_rejected_access_token + && response.status == StatusCode::UNAUTHORIZED.as_u16() + && let Some(error) = access_token_rejected(auth_token.as_deref()) + { + return Err(error); + } if response.status == StatusCode::UNAUTHORIZED.as_u16() && let Some(header) = response_header(&response.headers, reqwest::header::WWW_AUTHENTICATE) @@ -240,7 +257,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { let mut headers = self.default_headers.clone(); headers.extend(custom_headers); self.add_auth_headers(&mut headers); - if let Some(auth_token) = auth_token { + if let Some(auth_token) = auth_token.as_deref() { insert_header( &mut headers, AUTHORIZATION, @@ -274,6 +291,12 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { if response.status == StatusCode::METHOD_NOT_ALLOWED.as_u16() { return Ok(()); } + if self.attribute_rejected_access_token + && response.status == StatusCode::UNAUTHORIZED.as_u16() + && let Some(error) = access_token_rejected(auth_token.as_deref()) + { + return Err(error); + } if !status_is_success(response.status) { return Err(StreamableHttpError::UnexpectedServerResponse( format!("DELETE returned HTTP {}", response.status).into(), @@ -316,7 +339,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { StreamableHttpClientAdapterError::Header, )?; } - if let Some(auth_token) = auth_token { + if let Some(auth_token) = auth_token.as_deref() { insert_header( &mut headers, AUTHORIZATION, @@ -349,6 +372,12 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { StreamableHttpClientAdapterError::SessionExpired404, )); } + if self.attribute_rejected_access_token + && response.status == StatusCode::UNAUTHORIZED.as_u16() + && let Some(error) = access_token_rejected(auth_token.as_deref()) + { + return Err(error); + } if !status_is_success(response.status) { return Err(StreamableHttpError::UnexpectedServerResponse( format!("GET returned HTTP {}", response.status).into(), @@ -371,6 +400,20 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { } } +fn access_token_rejected( + auth_token: Option<&str>, +) -> Option> { + // Preserve the token associated with this response. Reading the current credential after a + // delayed 401 is racy: another concurrent request may already have refreshed A to B, in which + // case recovery must retry B rather than refresh B a second time. AccessToken's Debug + // implementation redacts the secret if this error is logged. + auth_token.map(|rejected_access_token| { + StreamableHttpError::Client(StreamableHttpClientAdapterError::AccessTokenRejected { + rejected_access_token: AccessToken::new(rejected_access_token.to_string()), + }) + }) +} + impl StreamableHttpClientAdapter { fn add_auth_headers(&self, headers: &mut HeaderMap) { if let Some(auth_provider) = &self.auth_provider { diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index 74dd1697bebd..84eba64fed3f 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -6,6 +6,7 @@ mod in_process_transport; mod logging_client_handler; mod oauth; mod oauth_http_client; +mod oauth_transport; mod perform_oauth_login; mod program_resolver; mod rmcp_client; diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index b8736168856c..6676ad999149 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -69,6 +69,7 @@ use tokio::sync::Mutex; use codex_utils_home_dir::find_codex_home; +pub(crate) use self::refresh_transaction::request_oauth_token_response; pub(crate) use self::resolved_store::ResolvedOAuthCredentialStore; pub(crate) use self::resolved_store::ResolvedOAuthTokens; pub(crate) use self::resolved_store::resolve_oauth_tokens_from_store_policy; @@ -509,72 +510,7 @@ impl OAuthPersistor { }), } } - - /// Persists RMCP-managed credential changes back to this client's resolved authority. - #[expect( - clippy::await_holding_invalid_type, - reason = "AuthorizationManager async access must be serialized through its mutex" - )] - pub(crate) async fn persist_if_needed(&self) -> Result<()> { - let (client_id, maybe_credentials) = { - let manager = self.inner.authorization_manager.clone(); - let guard = manager.lock().await; - guard.get_credentials().await - }?; - - match maybe_credentials { - Some(credentials) => { - let mut last_credentials = self.inner.last_credentials.lock().await; - let new_token_response = WrappedOAuthTokenResponse(credentials.clone()); - let same_token = last_credentials - .as_ref() - .map(|previous| previous.token_response == new_token_response) - .unwrap_or(false); - let expires_at = if same_token { - last_credentials - .as_ref() - .and_then(|previous| previous.expires_at) - } else { - compute_expires_at_millis(&credentials) - }; - let stored = StoredOAuthTokens { - server_name: self.inner.server_name.clone(), - url: self.inner.url.clone(), - client_id, - token_response: new_token_response, - expires_at, - }; - if last_credentials.as_ref() != Some(&stored) { - self.inner.credential_store.save( - &DefaultKeyringStore, - &self.inner.server_name, - &stored, - )?; - *last_credentials = Some(stored); - } - } - None => { - let mut last_credentials = self.inner.last_credentials.lock().await; - if last_credentials.take().is_some() - && let Err(error) = self.inner.credential_store.delete( - &DefaultKeyringStore, - &self.inner.server_name, - &self.inner.url, - ) - { - warn!( - server_name = %self.inner.server_name, - error = %error, - "failed to remove MCP OAuth credentials from the resolved store" - ); - } - } - } - - Ok(()) - } } - const FALLBACK_FILENAME: &str = ".credentials.json"; const MCP_SERVER_TYPE: &str = "http"; diff --git a/codex-rs/rmcp-client/src/oauth/refresh_transaction.rs b/codex-rs/rmcp-client/src/oauth/refresh_transaction.rs index 9a656ed3cf5f..8420b63ad181 100644 --- a/codex-rs/rmcp-client/src/oauth/refresh_transaction.rs +++ b/codex-rs/rmcp-client/src/oauth/refresh_transaction.rs @@ -8,6 +8,7 @@ use anyhow::Context; use anyhow::Result; use codex_keyring_store::DefaultKeyringStore; use codex_keyring_store::KeyringStore; +use oauth2::AccessToken; use oauth2::TokenResponse; use rmcp::transport::auth::AuthError; use rmcp::transport::auth::AuthorizationManager; @@ -31,42 +32,103 @@ const REFRESH_REQUEST_TIMEOUT: Duration = Duration::from_secs(45); impl OAuthPersistor { pub(crate) async fn refresh_if_needed(&self) -> Result<()> { - self.refresh_if_needed_in(&DefaultKeyringStore, REFRESH_REQUEST_TIMEOUT) - .await + self.refresh_in( + DefaultKeyringStore, + RefreshReason::Expiry, + REFRESH_REQUEST_TIMEOUT, + ) + .await } - /// Injects the credential backend and provider timeout for deterministic failure-path tests. - pub(super) async fn refresh_if_needed_in( + pub(crate) async fn refresh_after_unauthorized( &self, - keyring_store: &K, - refresh_request_timeout: Duration, + rejected_access_token: AccessToken, ) -> Result<()> { - let expires_at = { - let guard = self.inner.last_credentials.lock().await; - guard.as_ref().and_then(|tokens| tokens.expires_at) + self.refresh_in( + DefaultKeyringStore, + RefreshReason::Unauthorized { + rejected_access_token, + }, + REFRESH_REQUEST_TIMEOUT, + ) + .await + } + + /// Adopts a credential committed by another serialized refresh after this caller already used + /// its one provider refresh. A delayed 401 for B must not force reauthorization if C is now + /// authoritative, but this path must never rotate B a second time. + pub(crate) async fn adopt_newer_credentials_after_unauthorized( + &self, + rejected_access_token: &AccessToken, + ) -> Result { + let _lock = + RefreshCredentialLock::acquire_for_server(&self.inner.server_name, &self.inner.url) + .await?; + let Some(latest) = self.inner.credential_store.load( + &DefaultKeyringStore, + &self.inner.server_name, + &self.inner.url, + )? + else { + let manager = self.inner.authorization_manager.clone(); + manager + .lock() + .await + .set_credential_store(InMemoryCredentialStore::new()); + *self.inner.last_credentials.lock().await = None; + return Err(AuthError::AuthorizationRequired).with_context(|| { + format!( + "OAuth tokens for server {} were removed before recovery; authorization required", + self.inner.server_name + ) + }); }; + if latest.token_response.0.access_token().secret() == rejected_access_token.secret() { + return Ok(false); + } - if !token_needs_refresh(expires_at) { - return Ok(()); + debug!("adopting newer MCP OAuth credentials after a delayed unauthorized response"); + let manager = self.inner.authorization_manager.clone(); + let mut guard = manager.lock().await; + install_tokens_in_manager_guard(&mut guard, &latest, CredentialExposure::Request).await?; + *self.inner.last_credentials.lock().await = Some(latest); + Ok(true) + } + + /// Injects the credential backend and provider timeout for deterministic failure-path tests. + pub(super) async fn refresh_in( + &self, + keyring_store: K, + reason: RefreshReason, + refresh_request_timeout: Duration, + ) -> Result<()> { + if matches!(reason, RefreshReason::Expiry) { + let expires_at = { + let guard = self.inner.last_credentials.lock().await; + guard.as_ref().and_then(|tokens| tokens.expires_at) + }; + if !token_needs_refresh(expires_at) { + return Ok(()); + } } let persistor = self.clone(); - let keyring_store = keyring_store.clone(); // Once the provider can consume a rotating token, caller cancellation must not cancel // persistence. The owned task continues with independently bounded lock and request waits. // A provider timeout leaves the outcome unknown and permits a later serialized retry: // provider grace may recover, otherwise reauthorization is unavoidable. This residual // risk is preferred to holding the credential lock indefinitely. + let refresh_reason = reason.as_str(); let transaction_task = tokio::spawn(async move { let result = persistor - .refresh_transaction(&keyring_store, refresh_request_timeout) + .refresh_transaction(&keyring_store, reason, refresh_request_timeout) .await; // Keep this summary inside the owned task so caller cancellation cannot suppress it. if let Err(error) = &result { warn!( server_name = %persistor.inner.server_name, - refresh_reason = "expiry", + refresh_reason, error = %error, "MCP OAuth refresh transaction failed" ); @@ -91,13 +153,14 @@ impl OAuthPersistor { skip_all, fields( server_name = %self.inner.server_name, - refresh_reason = "expiry", + refresh_reason = reason.as_str(), ), err )] async fn refresh_transaction( &self, keyring_store: &K, + reason: RefreshReason, refresh_request_timeout: Duration, ) -> Result<()> { debug!("waiting for the MCP OAuth credential transaction lock"); @@ -132,11 +195,22 @@ impl OAuthPersistor { }); }; - if !token_needs_refresh(latest.expires_at) { + let latest_access_token = latest.token_response.0.access_token().secret(); + // A 401 belongs to the token sent by that request. If another request already refreshed A + // to B, adopt B and retry rather than rotating B because a delayed response rejected A. + let should_adopt = !token_needs_refresh(latest.expires_at) + && match &reason { + RefreshReason::Expiry => true, + RefreshReason::Unauthorized { + rejected_access_token, + } => rejected_access_token.secret() != latest_access_token, + }; + if should_adopt { debug!("adopting newer MCP OAuth credentials without contacting the provider"); let manager = self.inner.authorization_manager.clone(); let mut guard = manager.lock().await; - install_tokens_in_manager_guard(&mut guard, &latest).await?; + install_tokens_in_manager_guard(&mut guard, &latest, CredentialExposure::Request) + .await?; *self.inner.last_credentials.lock().await = Some(latest); return Ok(()); } @@ -161,19 +235,24 @@ impl OAuthPersistor { // The provider uses a separate HTTP client and cannot re-enter `AuthClient`. Retain this // async guard so requests cannot observe credentials while they are staged and committed. let mut guard = manager.lock().await; - install_tokens_in_manager_guard(&mut guard, &latest) - .await - .context("failed to stage OAuth credentials for refresh")?; + if let Err(error) = + install_tokens_in_manager_guard(&mut guard, &latest, CredentialExposure::Refresh).await + { + install_tokens_in_manager_guard(&mut guard, &latest, CredentialExposure::Request) + .await + .context("failed to restore request-only OAuth credentials")?; + return Err(error).context("failed to stage OAuth credentials for refresh"); + } // The owned task prevents caller deadlines from canceling after possible token rotation; // this timeout independently bounds the provider request. debug!( timeout_ms = refresh_request_timeout.as_millis(), "requesting refreshed MCP OAuth credentials from the provider" ); - let refreshed = match timeout(refresh_request_timeout, guard.refresh_token()).await { + let refresh_result = match timeout(refresh_request_timeout, guard.refresh_token()).await { Ok(Ok(token_response)) => { debug!("received refreshed MCP OAuth credentials from the provider"); - refreshed_tokens(token_response, &latest, &self.inner) + Ok(refreshed_tokens(token_response, &latest, &self.inner)) } Ok(Err(error @ AuthError::TokenRefreshFailed(_))) => { // RMCP 1.8 collapses definitive OAuth rejection (for example, @@ -199,24 +278,29 @@ impl OAuthPersistor { error = %error, "MCP OAuth provider refresh failed" ); - return Err(error).with_context(|| { + Err(error).with_context(|| { format!( "failed to refresh OAuth tokens for server {}", self.inner.server_name ) - }); + }) } Err(_) => { warn!( timeout_ms = refresh_request_timeout.as_millis(), "MCP OAuth provider refresh timed out; the outcome is unknown and a later serialized retry is permitted" ); - anyhow::bail!( + Err(anyhow::anyhow!( "timed out after {refresh_request_timeout:?} refreshing OAuth tokens for server {}", self.inner.server_name - ); + )) } }; + let request_tokens = refresh_result.as_ref().unwrap_or(&latest); + install_tokens_in_manager_guard(&mut guard, request_tokens, CredentialExposure::Request) + .await + .context("failed to restore request-only OAuth credentials")?; + let refreshed = refresh_result?; // Persist to the pinned source before exposing the refreshed token. On failure, restore // the prior in-process credential and return the error; serving an unpersisted token would @@ -235,18 +319,17 @@ impl OAuthPersistor { error = %error, "failed to persist refreshed MCP OAuth credentials; returning the error and restoring the previous in-process credentials" ); - install_tokens_in_manager_guard(&mut guard, &latest) + install_tokens_in_manager_guard(&mut guard, &latest, CredentialExposure::Request) .await .context( - "failed to restore previous OAuth credentials after refresh persistence failed", + "failed to restore previous request-only OAuth credentials after refresh persistence failed", )?; return Err(error); } - // This layer retains RMCP's legacy persistence hook. Install the same merged response - // (including carried-forward refresh token/scopes) so that hook cannot overwrite durable - // credentials with the provider's partial response. - install_tokens_in_manager_guard(&mut guard, &refreshed) + // Commit the merged response, then expose only its request-safe form so RMCP cannot + // refresh independently outside this transaction. + install_tokens_in_manager_guard(&mut guard, &refreshed, CredentialExposure::Request) .await .context( "refreshed OAuth tokens were persisted but could not be installed in the authorization manager", @@ -258,20 +341,50 @@ impl OAuthPersistor { } } +pub(super) enum RefreshReason { + Expiry, + Unauthorized { rejected_access_token: AccessToken }, +} + +impl RefreshReason { + fn as_str(&self) -> &'static str { + match self { + Self::Expiry => "expiry", + Self::Unauthorized { .. } => "unauthorized", + } + } +} + +/// Ordinary requests receive neither refresh token nor expiry metadata, so RMCP cannot refresh +/// outside Codex's transaction. Full credentials are exposed only while both transaction locks +/// are held. +#[derive(Clone, Copy)] +enum CredentialExposure { + Request, + Refresh, +} + async fn install_tokens_in_manager_guard( authorization_manager: &mut AuthorizationManager, tokens: &StoredOAuthTokens, + exposure: CredentialExposure, ) -> Result<()> { let store = InMemoryCredentialStore::new(); - let token_response = tokens.token_response.0.clone(); + let token_response = match exposure { + CredentialExposure::Request => request_oauth_token_response(tokens), + CredentialExposure::Refresh => tokens.token_response.0.clone(), + }; let granted_scopes = token_response .scopes() .map(|scopes| scopes.iter().map(|scope| scope.to_string()).collect()) .unwrap_or_default(); - let token_received_at = SystemTime::now() - .duration_since(UNIX_EPOCH) - .ok() - .map(|duration| duration.as_secs()); + let token_received_at = match exposure { + CredentialExposure::Request => None, + CredentialExposure::Refresh => SystemTime::now() + .duration_since(UNIX_EPOCH) + .ok() + .map(|duration| duration.as_secs()), + }; store .save(StoredCredentials::new( tokens.client_id.clone(), @@ -292,6 +405,13 @@ async fn install_tokens_in_manager_guard( Ok(()) } +pub(crate) fn request_oauth_token_response(tokens: &StoredOAuthTokens) -> OAuthTokenResponse { + let mut token_response = tokens.token_response.0.clone(); + token_response.set_refresh_token(None); + token_response.set_expires_in(None); + token_response +} + fn refreshed_tokens( mut token_response: OAuthTokenResponse, previous: &StoredOAuthTokens, diff --git a/codex-rs/rmcp-client/src/oauth/resolved_store.rs b/codex-rs/rmcp-client/src/oauth/resolved_store.rs index 6a688ea53938..937fa51e6df2 100644 --- a/codex-rs/rmcp-client/src/oauth/resolved_store.rs +++ b/codex-rs/rmcp-client/src/oauth/resolved_store.rs @@ -9,10 +9,6 @@ use tracing::warn; use super::OAuthKeyringLoadError; use super::StoredOAuthTokens; -use super::compute_store_key; -use super::delete_oauth_tokens_from_direct_keyring; -use super::delete_oauth_tokens_from_file; -use super::delete_oauth_tokens_from_secrets_keyring; use super::load_oauth_tokens_from_file; use super::load_oauth_tokens_from_keyring; use super::save_oauth_tokens_to_file; @@ -74,27 +70,6 @@ impl ResolvedOAuthCredentialStore { ), } } - - /// Deletes credentials only from this already-resolved authority. - pub(crate) fn delete( - self, - keyring_store: &K, - server_name: &str, - url: &str, - ) -> Result { - match self { - Self::File => { - let key = compute_store_key(server_name, url)?; - delete_oauth_tokens_from_file(&key) - } - Self::Keyring(AuthKeyringBackendKind::Direct) => { - delete_oauth_tokens_from_direct_keyring(keyring_store, server_name, url) - } - Self::Keyring(AuthKeyringBackendKind::Secrets) => { - delete_oauth_tokens_from_secrets_keyring(keyring_store, server_name, url) - } - } - } } #[derive(Debug)] diff --git a/codex-rs/rmcp-client/src/oauth/tests/persistor_tests.rs b/codex-rs/rmcp-client/src/oauth/tests/persistor_tests.rs index f32c2b691c2b..ff9121c628f8 100644 --- a/codex-rs/rmcp-client/src/oauth/tests/persistor_tests.rs +++ b/codex-rs/rmcp-client/src/oauth/tests/persistor_tests.rs @@ -7,6 +7,7 @@ use anyhow::Result; use codex_config::types::AuthKeyringBackendKind; use keyring::Error as KeyringError; use oauth2::AccessToken; +use oauth2::RefreshToken; use oauth2::TokenResponse; use pretty_assertions::assert_eq; use rmcp::transport::auth::AuthError; @@ -38,6 +39,7 @@ use crate::oauth::WrappedOAuthTokenResponse; use crate::oauth::compute_store_key; use crate::oauth::load_oauth_tokens_from_file; use crate::oauth::refresh_lock::RefreshCredentialLock; +use crate::oauth::refresh_transaction::RefreshReason; use crate::oauth::save_oauth_tokens_to_file; use crate::startup_error::is_authentication_required_error; @@ -129,10 +131,6 @@ async fn concurrent_refreshes_call_provider_once_and_carry_omitted_fields() -> R second_task.await??; server.verify().await; - // Layer 2 still invokes the legacy RMCP persistence hook after operations. Exercise that hook - // so a raw provider response that omitted refresh token/scopes cannot overwrite the merged - // authoritative credential. - first.persist_if_needed().await?; let stored = load_oauth_tokens_from_file(&initial.server_name, &initial.url)? .expect("refreshed credentials should be stored"); let mut expected_response = initial.token_response.0.clone(); @@ -148,6 +146,145 @@ async fn concurrent_refreshes_call_provider_once_and_carry_omitted_fields() -> R Ok(()) } +#[expect( + clippy::await_holding_invalid_type, + reason = "AuthorizationManager async access must be serialized through its Tokio mutex" +)] +#[tokio::test(flavor = "current_thread")] +async fn delayed_unauthorized_retries_adopt_the_winning_token() -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains("refresh_token=refresh-token")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "refreshed-access-token", + "token_type": "Bearer", + "expires_in": 3600, + }))) + .expect(1) + .mount(&server) + .await; + let mut initial = expired_tokens(&format!("{}/mcp", server.uri())); + initial.expires_at = None; + initial.token_response.0.set_expires_in(None); + save_oauth_tokens_to_file(&initial)?; + + let first = persistor_for(&initial).await?; + let second_manager = authorization_manager_for(&initial).await?; + let second = OAuthPersistor::new( + initial.server_name.clone(), + initial.url.clone(), + Arc::clone(&second_manager), + ResolvedOAuthCredentialStore::File, + Some(initial.clone()), + ); + let rejected_access_token = initial.token_response.0.access_token().clone(); + + first + .refresh_after_unauthorized(rejected_access_token.clone()) + .await?; + // Both calls model requests that left with A. Once the first 401 rotates A to B, later 401s + // must adopt B and retry their requests instead of rotating B again. + first + .refresh_after_unauthorized(rejected_access_token.clone()) + .await?; + second + .refresh_after_unauthorized(rejected_access_token) + .await?; + + server.verify().await; + let stored = load_oauth_tokens_from_file(&initial.server_name, &initial.url)? + .expect("the winning refresh should be persisted"); + assert_eq!( + stored.token_response.0.access_token().secret(), + "refreshed-access-token" + ); + let guard = second_manager.lock().await; + let (_client_id, adopted) = guard.get_credentials().await?; + let adopted = adopted.expect("second manager should adopt the winning token"); + assert_eq!(adopted.access_token().secret(), "refreshed-access-token"); + assert!(adopted.refresh_token().is_none()); + Ok(()) +} + +#[expect( + clippy::await_holding_invalid_type, + reason = "AuthorizationManager async access must be serialized through its Tokio mutex" +)] +#[tokio::test(flavor = "current_thread")] +async fn second_unauthorized_retry_adopts_newer_credentials_without_refreshing() -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains("refresh_token=refresh-token-a")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "access-token-b", + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": "refresh-token-b", + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains("refresh_token=refresh-token-b")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "access-token-c", + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": "refresh-token-c", + }))) + .expect(1) + .mount(&server) + .await; + let mut initial = expired_tokens(&format!("{}/mcp", server.uri())); + initial.expires_at = None; + initial.token_response.0.set_expires_in(None); + initial + .token_response + .0 + .set_refresh_token(Some(RefreshToken::new("refresh-token-a".to_string()))); + save_oauth_tokens_to_file(&initial)?; + + let first_manager = authorization_manager_for(&initial).await?; + let first = OAuthPersistor::new( + initial.server_name.clone(), + initial.url.clone(), + Arc::clone(&first_manager), + ResolvedOAuthCredentialStore::File, + Some(initial.clone()), + ); + let second = persistor_for(&initial).await?; + let access_token_a = initial.token_response.0.access_token().clone(); + first.refresh_after_unauthorized(access_token_a).await?; + + let access_token_b = AccessToken::new("access-token-b".to_string()); + second + .refresh_after_unauthorized(access_token_b.clone()) + .await?; + assert!( + first + .adopt_newer_credentials_after_unauthorized(&access_token_b) + .await? + ); + + server.verify().await; + let guard = first_manager.lock().await; + let (_client_id, adopted) = guard.get_credentials().await?; + let adopted = adopted.expect("first manager should adopt the newest token"); + assert_eq!(adopted.access_token().secret(), "access-token-c"); + assert!(adopted.refresh_token().is_none()); + Ok(()) +} + #[expect( clippy::await_holding_invalid_type, reason = "AuthorizationManager async access must be serialized through its Tokio mutex" @@ -168,7 +305,11 @@ async fn resolved_keyring_read_error_preserves_in_memory_credentials() -> Result ); let error = persistor - .refresh_if_needed_in(&keyring_store, Duration::from_secs(/*secs*/ 45)) + .refresh_in( + keyring_store, + RefreshReason::Expiry, + Duration::from_secs(/*secs*/ 45), + ) .await .expect_err("the resolved keyring read error should abort refresh"); assert!( @@ -310,8 +451,9 @@ async fn provider_timeout_releases_lock_and_preserves_durable_credentials() -> R let persistor = persistor_for(&initial).await?; let error = persistor - .refresh_if_needed_in( - &MockKeyringStore::default(), + .refresh_in( + MockKeyringStore::default(), + RefreshReason::Expiry, Duration::from_millis(/*millis*/ 50), ) .await diff --git a/codex-rs/rmcp-client/src/oauth_transport.rs b/codex-rs/rmcp-client/src/oauth_transport.rs new file mode 100644 index 000000000000..ad4a9fadb941 --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth_transport.rs @@ -0,0 +1,220 @@ +//! Codex-owned OAuth policy for RMCP Streamable HTTP traffic. +//! +//! RMCP remains responsible for transport mechanics and bearer-token injection. Its authorization +//! manager receives only request-safe credentials, so it cannot independently refresh outside +//! Codex's serialized transaction. +//! +//! Client-originated requests retain their outer `RmcpClient` recovery, which owns caller +//! deadlines and replay decisions. RMCP-owned responses, SSE GET/reconnects, and session DELETEs +//! have no public caller; this transport reports their exact rejected token to the parent +//! `RmcpClient` and stops RMCP's unbounded SSE reconnect loop. The parent then owns any refresh +//! and session rebuild before the next public operation. + +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::PoisonError; +use std::time::Duration; + +use oauth2::AccessToken; +use reqwest::header::HeaderName; +use reqwest::header::HeaderValue; +use rmcp::model::ClientJsonRpcMessage; +use rmcp::model::JsonRpcMessage; +use rmcp::transport::auth::AuthClient; +use rmcp::transport::common::client_side_sse::ExponentialBackoff; +use rmcp::transport::common::client_side_sse::SseRetryPolicy; +use rmcp::transport::streamable_http_client::StreamableHttpClient; +use rmcp::transport::streamable_http_client::StreamableHttpError; +use rmcp::transport::streamable_http_client::StreamableHttpPostResponse; + +use crate::http_client_adapter::StreamableHttpClientAdapter; +use crate::http_client_adapter::StreamableHttpClientAdapterError; + +type TransportResult = + std::result::Result>; + +#[derive(Clone)] +pub(crate) struct OAuthTransportClient { + auth_client: AuthClient, + failure_state: OAuthTransportFailureState, +} + +impl OAuthTransportClient { + pub(crate) fn new( + auth_client: AuthClient, + failure_state: OAuthTransportFailureState, + ) -> Self { + Self { + auth_client, + failure_state, + } + } +} + +/// Shared state between RMCP's bearer-only transport and Codex's OAuth session owner. +/// +/// RMCP may issue GET reconnects, DELETE cleanup, and server-response POSTs outside a public +/// `RmcpClient` operation. Those requests may report which access token was rejected, but they +/// must not refresh it: Codex owns the credential transaction and transport rebuild. The state +/// also stops RMCP's unbounded SSE reconnect policy after an auth failure so it cannot repeatedly +/// re-enter this transport with a rejected token while Codex is recovering the session. +#[derive(Clone, Debug, Default)] +pub(crate) struct OAuthTransportFailureState { + inner: Arc, +} + +#[derive(Debug, Default)] +struct OAuthTransportFailureStateInner { + pending_rejected_access_token: Mutex>, +} + +impl OAuthTransportFailureState { + pub(crate) fn record_rejected_access_token(&self, rejected_access_token: AccessToken) { + *self + .inner + .pending_rejected_access_token + .lock() + .unwrap_or_else(PoisonError::into_inner) = Some(rejected_access_token); + } + + pub(crate) fn pending_rejected_access_token(&self) -> Option { + self.inner + .pending_rejected_access_token + .lock() + .unwrap_or_else(PoisonError::into_inner) + .clone() + } + + pub(crate) fn finish_recovery(&self, rejected_access_token: &AccessToken) { + let mut pending = self + .inner + .pending_rejected_access_token + .lock() + .unwrap_or_else(PoisonError::into_inner); + if pending + .as_ref() + .is_some_and(|pending| pending.secret() == rejected_access_token.secret()) + { + *pending = None; + } + } + + pub(crate) fn retry_policy(&self) -> OAuthSseRetryPolicy { + OAuthSseRetryPolicy { + failure_state: self.clone(), + fallback: ExponentialBackoff::default(), + } + } +} + +#[derive(Debug)] +pub(crate) struct OAuthSseRetryPolicy { + failure_state: OAuthTransportFailureState, + fallback: ExponentialBackoff, +} + +impl SseRetryPolicy for OAuthSseRetryPolicy { + fn retry(&self, current_times: usize) -> Option { + if self + .failure_state + .inner + .pending_rejected_access_token + .lock() + .unwrap_or_else(PoisonError::into_inner) + .is_some() + { + None + } else { + self.fallback.retry(current_times) + } + } +} + +impl StreamableHttpClient for OAuthTransportClient { + type Error = StreamableHttpClientAdapterError; + + async fn post_message( + &self, + uri: Arc, + message: ClientJsonRpcMessage, + session_id: Option>, + auth_token: Option, + custom_headers: HashMap, + ) -> TransportResult { + let is_rmcp_owned_response = matches!( + message, + JsonRpcMessage::Response(_) | JsonRpcMessage::Error(_) + ); + let result = self + .auth_client + .post_message(uri, message, session_id, auth_token, custom_headers) + .await; + + // Client-originated requests retain their outer `RmcpClient` recovery boundary, which + // owns caller deadlines and replay decisions. Server responses have no public caller, so + // surface their rejected token to the Codex session owner instead of refreshing here. + if is_rmcp_owned_response + && let Some(rejected_access_token) = + result.as_ref().err().and_then(rejected_access_token) + { + self.failure_state + .record_rejected_access_token(rejected_access_token); + } + result + } + + async fn delete_session( + &self, + uri: Arc, + session_id: Arc, + auth_token: Option, + custom_headers: HashMap, + ) -> TransportResult<()> { + let result = self + .auth_client + .delete_session(uri, session_id, auth_token, custom_headers) + .await; + if let Some(rejected_access_token) = result.as_ref().err().and_then(rejected_access_token) { + self.failure_state + .record_rejected_access_token(rejected_access_token); + } + result + } + + async fn get_stream( + &self, + uri: Arc, + session_id: Arc, + last_event_id: Option, + auth_token: Option, + custom_headers: HashMap, + ) -> TransportResult< + futures::stream::BoxStream<'static, Result>, + > { + let result = self + .auth_client + .get_stream(uri, session_id, last_event_id, auth_token, custom_headers) + .await; + if let Some(rejected_access_token) = result.as_ref().err().and_then(rejected_access_token) { + self.failure_state + .record_rejected_access_token(rejected_access_token); + } + result + } +} + +fn rejected_access_token( + error: &StreamableHttpError, +) -> Option { + match error { + StreamableHttpError::Client(StreamableHttpClientAdapterError::AccessTokenRejected { + rejected_access_token, + }) => Some(rejected_access_token.clone()), + _ => None, + } +} + +#[cfg(test)] +#[path = "oauth_transport_tests.rs"] +mod tests; diff --git a/codex-rs/rmcp-client/src/oauth_transport_tests.rs b/codex-rs/rmcp-client/src/oauth_transport_tests.rs new file mode 100644 index 000000000000..258b0c68f873 --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth_transport_tests.rs @@ -0,0 +1,168 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use oauth2::AccessToken; +use oauth2::RefreshToken; +use oauth2::basic::BasicTokenType; +use reqwest::header::HeaderMap; +use rmcp::model::ClientJsonRpcMessage; +use rmcp::transport::auth::AuthClient; +use rmcp::transport::auth::OAuthState; +use rmcp::transport::auth::OAuthTokenResponse; +use rmcp::transport::auth::VendorExtraTokenFields; +use rmcp::transport::common::client_side_sse::SseRetryPolicy; +use rmcp::transport::streamable_http_client::StreamableHttpClient; +use rmcp::transport::streamable_http_client::StreamableHttpError; +use serde_json::json; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::header; +use wiremock::matchers::method; +use wiremock::matchers::path; + +use super::OAuthTransportClient; +use super::OAuthTransportFailureState; +use crate::http_client_adapter::StreamableHttpClientAdapter; +use crate::http_client_adapter::StreamableHttpClientAdapterError; +use crate::oauth::StoredOAuthTokens; +use crate::oauth::WrappedOAuthTokenResponse; +use crate::oauth::request_oauth_token_response; +use crate::oauth_http_client::OAuthHttpClientAdapter; + +const SERVER_NAME: &str = "oauth-transport-response-test"; +const ACCESS_TOKEN_A: &str = "response-access-a"; +const REFRESH_TOKEN_A: &str = "response-refresh-a"; + +#[tokio::test] +async fn rmcp_owned_response_reports_rejected_token_without_refreshing() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": ["scope-a"], + }))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(500)) + .expect(0) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header("authorization", format!("Bearer {ACCESS_TOKEN_A}"))) + .respond_with(ResponseTemplate::new(401)) + .expect(1) + .mount(&server) + .await; + + let server_url = format!("{}/mcp", server.uri()); + let initial_tokens = initial_tokens(&server_url); + let http_client = codex_exec_server::Environment::default_for_tests().get_http_client(); + let oauth_http_client = Arc::new(OAuthHttpClientAdapter::new( + Arc::clone(&http_client), + HeaderMap::new(), + )); + let mut oauth_state = + OAuthState::new_with_oauth_http_client(server_url.clone(), oauth_http_client).await?; + oauth_state + .set_credentials( + &initial_tokens.client_id, + request_oauth_token_response(&initial_tokens), + ) + .await?; + let manager = match oauth_state { + OAuthState::Authorized(manager) | OAuthState::Unauthorized(manager) => manager, + _ => anyhow::bail!("unexpected OAuth state during response test setup"), + }; + let auth_client = AuthClient::new( + StreamableHttpClientAdapter::new( + Arc::clone(&http_client), + HeaderMap::new(), + /*auth_provider*/ None, + ) + .with_rejected_token_attribution(), + manager, + ); + let failure_state = OAuthTransportFailureState::default(); + let client = OAuthTransportClient::new(auth_client, failure_state.clone()); + let response_message: ClientJsonRpcMessage = serde_json::from_value(json!({ + "jsonrpc": "2.0", + "id": "server-request-1", + "result": { + "action": "accept", + "content": { "confirmed": true } + } + }))?; + + let error = client + .post_message( + Arc::from(server_url), + response_message, + Some(Arc::from("response-session")), + /*auth_token*/ None, + HashMap::new(), + ) + .await + .expect_err("the server should reject the response token"); + + assert!(matches!( + error, + StreamableHttpError::Client(StreamableHttpClientAdapterError::AccessTokenRejected { .. }) + )); + assert_eq!( + failure_state + .pending_rejected_access_token() + .as_ref() + .map(|token| token.secret().as_str()), + Some(ACCESS_TOKEN_A) + ); + assert_eq!( + failure_state.retry_policy().retry(/*current_times*/ 1), + None + ); + server.verify().await; + Ok(()) +} + +#[test] +fn pending_auth_failure_stops_sse_retry_until_recovery_finishes() { + let failure_state = OAuthTransportFailureState::default(); + let rejected_access_token = AccessToken::new(ACCESS_TOKEN_A.to_string()); + + failure_state.record_rejected_access_token(rejected_access_token.clone()); + assert_eq!( + failure_state.retry_policy().retry(/*current_times*/ 1), + None + ); + + failure_state.finish_recovery(&rejected_access_token); + assert!(failure_state.pending_rejected_access_token().is_none()); + assert!( + failure_state + .retry_policy() + .retry(/*current_times*/ 1) + .is_some() + ); +} + +fn initial_tokens(server_url: &str) -> StoredOAuthTokens { + let mut response = OAuthTokenResponse::new( + AccessToken::new(ACCESS_TOKEN_A.to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + response.set_refresh_token(Some(RefreshToken::new(REFRESH_TOKEN_A.to_string()))); + response.set_expires_in(None); + StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: server_url.to_string(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: None, + } +} diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index ed81eefba798..f5d02f72ebb5 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -18,6 +18,7 @@ use codex_exec_server::HttpClient; use codex_keyring_store::DefaultKeyringStore; use futures::FutureExt; use futures::future::BoxFuture; +use oauth2::AccessToken; use oauth2::TokenResponse; use reqwest::header::AUTHORIZATION; use reqwest::header::HeaderMap; @@ -70,8 +71,11 @@ use crate::oauth::OAuthPersistor; use crate::oauth::ResolvedOAuthCredentialStore; use crate::oauth::ResolvedOAuthTokens; use crate::oauth::StoredOAuthTokens; +use crate::oauth::request_oauth_token_response; use crate::oauth::resolve_oauth_tokens_from_store_policy; use crate::oauth_http_client::OAuthHttpClientAdapter; +use crate::oauth_transport::OAuthTransportClient; +use crate::oauth_transport::OAuthTransportFailureState; use crate::stdio_server_launcher::StdioServerCommand; use crate::stdio_server_launcher::StdioServerLauncher; use crate::stdio_server_launcher::StdioServerProcessHandle; @@ -97,18 +101,25 @@ enum PendingTransport { transport: StreamableHttpClientTransport, }, StreamableHttpWithOAuth { - transport: StreamableHttpClientTransport>, - oauth_persistor: OAuthPersistor, + transport: StreamableHttpClientTransport, + oauth: OAuthRuntime, }, } +#[derive(Clone)] +struct OAuthRuntime { + transport_client: OAuthTransportClient, + persistor: OAuthPersistor, + failure_state: OAuthTransportFailureState, +} + enum ClientState { Connecting { transport: Option, }, Ready { service: Arc>, - oauth: Option, + oauth: Option, }, Closed, } @@ -131,6 +142,7 @@ enum TransportRecipe { store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, pinned_credential_store: Arc>, + oauth_runtime: Arc>, http_client: Arc, auth_provider: Option, }, @@ -408,6 +420,7 @@ impl RmcpClient { store_mode, keyring_backend_kind, pinned_credential_store: Arc::new(OnceLock::new()), + oauth_runtime: Arc::new(OnceLock::new()), http_client, auth_provider, }; @@ -450,8 +463,8 @@ impl RmcpClient { } }; - let (service, oauth_persistor) = self - .connect_pending_transport_with_initialize_retries( + let (service, oauth) = self + .connect_pending_transport_with_oauth_recovery( pending_transport, client_service.clone(), timeout, @@ -479,16 +492,10 @@ impl RmcpClient { } *guard = ClientState::Ready { service, - oauth: oauth_persistor.clone(), + oauth: oauth.clone(), }; } - if let Some(runtime) = oauth_persistor - && let Err(error) = runtime.persist_if_needed().await - { - warn!("failed to persist OAuth tokens after initialize: {error}"); - } - Ok(initialize_result) } @@ -504,7 +511,6 @@ impl RmcpClient { async move { service.list_tools(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -539,7 +545,6 @@ impl RmcpClient { }) }) .collect::>>()?; - self.persist_oauth_tokens().await; Ok(ListToolsWithConnectorIdResult { next_cursor: result.next_cursor, tools, @@ -566,7 +571,6 @@ impl RmcpClient { async move { service.list_resources(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -582,7 +586,6 @@ impl RmcpClient { async move { service.list_resource_templates(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -598,7 +601,6 @@ impl RmcpClient { async move { service.read_resource(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -656,7 +658,6 @@ impl RmcpClient { .boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -686,7 +687,6 @@ impl RmcpClient { }, ) .await?; - self.persist_oauth_tokens().await; Ok(()) } @@ -709,20 +709,24 @@ impl RmcpClient { .boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(response) } - async fn service(&self) -> Result>> { + async fn service_and_oauth_runtime( + &self, + ) -> Result<( + Arc>, + Option, + )> { let guard = self.state.lock().await; match &*guard { - ClientState::Ready { service, .. } => Ok(Arc::clone(service)), + ClientState::Ready { service, oauth } => Ok((Arc::clone(service), oauth.clone())), ClientState::Connecting { .. } => Err(anyhow!("MCP client not initialized")), ClientState::Closed => Err(anyhow!("MCP client is shut down")), } } - async fn oauth_persistor(&self) -> Option { + async fn oauth_runtime(&self) -> Option { let guard = self.state.lock().await; match &*guard { ClientState::Ready { @@ -749,22 +753,44 @@ impl RmcpClient { drop(previous_state); } - /// This should be called after every tool call so that if a given tool call triggered - /// a refresh of the OAuth tokens, they are persisted. - async fn persist_oauth_tokens(&self) { - if let Some(runtime) = self.oauth_persistor().await - && let Err(error) = runtime.persist_if_needed().await - { - warn!("failed to persist OAuth tokens: {error}"); + async fn refresh_oauth_if_needed(&self) -> Result<()> { + if self.recover_pending_oauth_failure().await? { + return Ok(()); + } + if let Some(runtime) = self.oauth_runtime().await { + runtime.persistor.refresh_if_needed().await?; } + Ok(()) } - /// OAuth uses independent lock/request bounds and completes before the operation timeout starts. - async fn refresh_oauth_if_needed(&self) -> Result<()> { - if let Some(runtime) = self.oauth_persistor().await { - runtime.refresh_if_needed().await?; + /// Recovers an auth failure emitted by RMCP-owned background traffic before the next public + /// operation runs. The bearer-only transport records the exact rejected token and terminates + /// its SSE reconnect loop; this parent layer owns the serialized refresh and session rebuild. + async fn recover_pending_oauth_failure(&self) -> Result { + let (failed_service, Some(runtime)) = self.service_and_oauth_runtime().await? else { + return Ok(false); + }; + let Some(rejected_access_token) = runtime.failure_state.pending_rejected_access_token() + else { + return Ok(false); + }; + + runtime + .persistor + .refresh_after_unauthorized(rejected_access_token.clone()) + .await?; + match self + .reinitialize_after_session_expiry(&failed_service) + .await + { + Ok(()) => { + runtime + .failure_state + .finish_recovery(&rejected_access_token); + Ok(true) + } + Err(error) => Err(error), } - Ok(()) } async fn create_pending_transport( @@ -790,6 +816,7 @@ impl RmcpClient { store_mode, keyring_backend_kind, pinned_credential_store, + oauth_runtime, http_client, auth_provider, } => { @@ -802,6 +829,20 @@ impl RmcpClient { auth_provider.clone() }; + // Reuse one OAuth manager and persistor across initialize retries and session + // reconstruction. This preserves the lifecycle-pinned store and keeps each failed + // request paired with the manager snapshot that supplied its access token. + if let Some(runtime) = oauth_runtime.get() { + let transport = StreamableHttpClientTransport::with_client( + runtime.transport_client.clone(), + oauth_transport_config(url, &runtime.failure_state), + ); + return Ok(PendingTransport::StreamableHttpWithOAuth { + transport, + oauth: runtime.clone(), + }); + } + let resolved_oauth_tokens = if bearer_token.is_none() && auth_provider.is_none() && !default_headers.contains_key(AUTHORIZATION) @@ -822,8 +863,10 @@ impl RmcpClient { ) { Ok(tokens) => { if let Some(resolved) = tokens.as_ref() { - // Retries and session recovery rebuild this transport. Pin the - // first concrete source so Auto is not reevaluated mid-client. + // Transport retries and session recovery are part of the same + // client lifecycle. Pin the first concrete source in memory so + // rebuilding a transport never re-evaluates Auto and adopts a + // possibly stale credential from another store. pinned_credential_store.set(resolved.store).map_err(|_| { anyhow!( "OAuth credential store pinned concurrently for MCP server `{server_name}`" @@ -847,7 +890,7 @@ impl RmcpClient { store: credential_store, }) = resolved_oauth_tokens { - match create_oauth_transport_and_runtime( + match create_oauth_runtime( server_name, url, initial_tokens.clone(), @@ -857,10 +900,21 @@ impl RmcpClient { ) .await { - Ok((transport, oauth_persistor)) => { + Ok(runtime) => { + oauth_runtime + .set(runtime.clone()) + .map_err(|_| { + anyhow!( + "OAuth runtime resolved concurrently for MCP server `{server_name}`" + ) + })?; + let transport = StreamableHttpClientTransport::with_client( + runtime.transport_client.clone(), + oauth_transport_config(url, &runtime.failure_state), + ); Ok(PendingTransport::StreamableHttpWithOAuth { transport, - oauth_persistor, + oauth: runtime, }) } Err(err) @@ -919,9 +973,9 @@ impl RmcpClient { timeout: Option, ) -> Result<( Arc>, - Option, + Option, )> { - let (transport, oauth_persistor) = match pending_transport { + let (transport, oauth) = match pending_transport { PendingTransport::InProcess { transport } => ( service::serve_client(client_service, transport).boxed(), None, @@ -934,12 +988,9 @@ impl RmcpClient { service::serve_client(client_service, transport).boxed(), None, ), - PendingTransport::StreamableHttpWithOAuth { - transport, - oauth_persistor, - } => ( + PendingTransport::StreamableHttpWithOAuth { transport, oauth } => ( service::serve_client(client_service, transport).boxed(), - Some(oauth_persistor), + Some(oauth), ), }; @@ -956,21 +1007,9 @@ impl RmcpClient { .await .map_err(|source| anyhow::Error::from(HandshakeError { source })), }; - let service = match service_result { - Ok(service) => service, - Err(error) => { - if let Some(runtime) = oauth_persistor.as_ref() - && let Err(persist_error) = runtime.persist_if_needed().await - { - warn!( - "failed to persist OAuth tokens after failed initialize: {persist_error}" - ); - } - return Err(error); - } - }; + let service = service_result?; - Ok((Arc::new(service), oauth_persistor)) + Ok((Arc::new(service), oauth)) } async fn run_service_operation( @@ -983,31 +1022,113 @@ impl RmcpClient { F: Fn(Arc>) -> Fut, Fut: std::future::Future>, { - let service = self.service().await?; - match Self::run_service_operation_with_transient_retries( - 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) => { + let deadline = timeout.map(|duration| Instant::now() + duration); + // Keep the OAuth runtime paired with the service that performs this operation. Session + // recovery can replace both while the request is in flight; rereading only the persistor + // after a 401 could refresh credentials owned by a different transport lifecycle. + let (mut service, mut oauth_runtime) = self.service_and_oauth_runtime().await?; + let mut oauth_recovered = false; + let mut retried_after_newer_credentials = false; + let mut session_recovered = false; + + loop { + let result = Self::run_service_operation_with_transient_retries( + Arc::clone(&service), + label, + timeout, + deadline, + self.elicitation_pause_state.clone(), + &operation, + ) + .await; + + if let Some(rejected_access_token) = result + .as_ref() + .err() + .and_then(Self::rejected_access_token_from_operation_error) + { + let Some(oauth_runtime) = oauth_runtime.as_ref() else { + return result.map_err(Into::into); + }; + if oauth_recovered { + // A delayed 401 can reject B after another operation already committed C. + // Adopt that newer durable credential once and retry without contacting the + // provider. If B is still authoritative, the one provider refresh was + // genuinely rejected and this lifecycle needs reauthentication. + let adopted_newer_credentials = if retried_after_newer_credentials { + false + } else { + let remaining = remaining_operation_timeout(label, timeout, deadline)?; + let adoption = oauth_runtime + .persistor + .adopt_newer_credentials_after_unauthorized(&rejected_access_token); + match remaining { + Some(remaining) => match time::timeout(remaining, adoption).await { + Ok(result) => result?, + Err(_) => { + return Err(ClientOperationError::Timeout { + label: label.to_string(), + duration: timeout.unwrap_or(remaining), + } + .into()); + } + }, + None => adoption.await?, + } + }; + if adopted_newer_credentials { + retried_after_newer_credentials = true; + continue; + } + return Err(AuthError::AuthorizationRequired.into()); + } + + // Public request/notification recovery stays here rather than in the transport + // wrapper because this layer owns the caller deadline. RMCP can continue + // processing a queued transport message after the caller times out; retrying it + // inside the wrapper could replay a timed-out tool call. + let remaining = remaining_operation_timeout(label, timeout, deadline)?; + let refresh = oauth_runtime + .persistor + .refresh_after_unauthorized(rejected_access_token); + let refresh_result = match remaining { + Some(remaining) => match time::timeout(remaining, refresh).await { + Ok(result) => result, + Err(_) => { + // The owned transaction keeps running after this caller stops waiting, + // but the rejected operation is not replayed after its deadline. + return Err(ClientOperationError::Timeout { + label: label.to_string(), + duration: timeout.unwrap_or(remaining), + } + .into()); + } + }, + None => refresh.await, + }; + if let Err(error) = refresh_result { + if let Err(timeout_error) = + remaining_operation_timeout(label, timeout, deadline) + { + return Err(timeout_error.into()); + } + return Err(error); + } + oauth_recovered = true; + continue; + } + + if !session_recovered && result.as_ref().is_err_and(Self::is_session_expired_404) { + // OAuth and session recovery are each one-shot, but either error may arrive first. + // Re-entering this loop lets 404 -> 401 compose just like the existing 401 -> 404 + // path without allowing either recovery to repeat indefinitely. self.reinitialize_after_session_expiry(&service).await?; - let recovered_service = self.service().await?; - Self::run_service_operation_with_transient_retries( - recovered_service, - label, - timeout, - self.elicitation_pause_state.clone(), - &operation, - ) - .await - .map_err(Into::into) + (service, oauth_runtime) = self.service_and_oauth_runtime().await?; + session_recovered = true; + continue; } - Err(error) => Err(error.into()), + + return result.map_err(Into::into); } } @@ -1015,6 +1136,7 @@ impl RmcpClient { service: Arc>, label: &str, timeout: Option, + retry_deadline: Option, pause_state: ElicitationPauseState, operation: &F, ) -> std::result::Result @@ -1022,7 +1144,6 @@ impl RmcpClient { F: Fn(Arc>) -> Fut, Fut: std::future::Future>, { - let retry_deadline = timeout.map(|duration| Instant::now() + duration); for (attempt, retry_delay_ms) in STREAMABLE_HTTP_RETRY_DELAYS_MS .iter() .copied() @@ -1128,6 +1249,34 @@ impl RmcpClient { }) } + fn rejected_access_token_from_operation_error( + error: &ClientOperationError, + ) -> Option { + let ClientOperationError::Service(rmcp::service::ServiceError::TransportSend(error)) = + error + else { + return None; + }; + + error + .error + .downcast_ref::>() + .and_then(Self::rejected_access_token) + } + + pub(super) fn rejected_access_token( + error: &StreamableHttpError, + ) -> Option { + match error { + StreamableHttpError::Client( + StreamableHttpClientAdapterError::AccessTokenRejected { + rejected_access_token, + }, + ) => Some(rejected_access_token.clone()), + _ => None, + } + } + async fn reinitialize_after_session_expiry( &self, failed_service: &Arc>, @@ -1161,8 +1310,8 @@ impl RmcpClient { .clone() .ok_or_else(|| anyhow!("MCP client cannot recover before initialize succeeds"))?; let pending_transport = Self::create_pending_transport(&self.transport_recipe).await?; - let (service, oauth_persistor) = self - .connect_pending_transport_with_initialize_retries( + let (service, oauth) = self + .connect_pending_transport_with_oauth_recovery( pending_transport, initialize_context.client_service, initialize_context.timeout, @@ -1176,31 +1325,22 @@ impl RmcpClient { } *guard = ClientState::Ready { service, - oauth: oauth_persistor.clone(), + oauth: oauth.clone(), }; } - if let Some(runtime) = oauth_persistor - && let Err(error) = runtime.persist_if_needed().await - { - warn!("failed to persist OAuth tokens after session recovery: {error}"); - } - Ok(()) } } -async fn create_oauth_transport_and_runtime( +async fn create_oauth_runtime( server_name: &str, url: &str, initial_tokens: StoredOAuthTokens, credential_store: ResolvedOAuthCredentialStore, default_headers: HeaderMap, http_client: Arc, -) -> Result<( - StreamableHttpClientTransport>, - OAuthPersistor, -)> { +) -> Result { let oauth_http_client = Arc::new(OAuthHttpClientAdapter::new( http_client.clone(), default_headers.clone(), @@ -1211,7 +1351,7 @@ async fn create_oauth_transport_and_runtime( oauth_state .set_credentials( &initial_tokens.client_id, - initial_tokens.token_response.0.clone(), + request_oauth_token_response(&initial_tokens), ) .await?; @@ -1224,25 +1364,36 @@ async fn create_oauth_transport_and_runtime( }; let auth_client = AuthClient::new( - StreamableHttpClientAdapter::new(http_client, default_headers, /*auth_provider*/ None), + StreamableHttpClientAdapter::new(http_client, default_headers, /*auth_provider*/ None) + .with_rejected_token_attribution(), manager, ); let auth_manager = auth_client.auth_manager.clone(); - let transport = StreamableHttpClientTransport::with_client( - auth_client, - StreamableHttpClientTransportConfig::with_uri(url.to_string()), - ); - - let runtime = OAuthPersistor::new( + let persistor = OAuthPersistor::new( server_name.to_string(), url.to_string(), auth_manager, credential_store, Some(initial_tokens), ); + let failure_state = OAuthTransportFailureState::default(); + let transport_client = OAuthTransportClient::new(auth_client, failure_state.clone()); + + Ok(OAuthRuntime { + transport_client, + persistor, + failure_state, + }) +} - Ok((transport, runtime)) +fn oauth_transport_config( + url: &str, + failure_state: &OAuthTransportFailureState, +) -> StreamableHttpClientTransportConfig { + let mut config = StreamableHttpClientTransportConfig::with_uri(url.to_string()); + config.retry_config = Arc::new(failure_state.retry_policy()); + config } #[cfg(test)] diff --git a/codex-rs/rmcp-client/src/startup_error.rs b/codex-rs/rmcp-client/src/startup_error.rs index c74f637e3d26..d10c8cd0ed9c 100644 --- a/codex-rs/rmcp-client/src/startup_error.rs +++ b/codex-rs/rmcp-client/src/startup_error.rs @@ -34,7 +34,7 @@ fn client_initialize_error_requires_authentication(error: &ClientInitializeError error, StreamableHttpError::Auth(auth_error) if auth_error_requires_authentication(auth_error) - ) + ) || matches!(error, StreamableHttpError::AuthRequired(_)) }) } diff --git a/codex-rs/rmcp-client/src/streamable_http_retry.rs b/codex-rs/rmcp-client/src/streamable_http_retry.rs index 5c4d28881533..12778aa34955 100644 --- a/codex-rs/rmcp-client/src/streamable_http_retry.rs +++ b/codex-rs/rmcp-client/src/streamable_http_retry.rs @@ -5,39 +5,175 @@ use std::time::Instant; use anyhow::Result; use anyhow::anyhow; use codex_exec_server::ExecServerError; +use oauth2::AccessToken; use reqwest::StatusCode; use rmcp::service::RoleClient; use rmcp::service::RunningService; +use rmcp::transport::auth::AuthError; use rmcp::transport::streamable_http_client::StreamableHttpError; use tokio::time; use tracing::warn; -use crate::elicitation_client_service::ElicitationClientService; -use crate::http_client_adapter::StreamableHttpClientAdapterError; -use crate::oauth::OAuthPersistor; - +use super::OAuthRuntime; use super::PendingTransport; use super::RmcpClient; +use crate::elicitation_client_service::ElicitationClientService; +use crate::http_client_adapter::StreamableHttpClientAdapterError; const JSON_RPC_INTERNAL_ERROR_CODE: i64 = -32603; pub(super) const STREAMABLE_HTTP_RETRY_DELAYS_MS: [u64; 2] = [250, 1_000]; +#[derive(Default)] +struct InitializeAttemptContext { + oauth: Option, +} + impl RmcpClient { - pub(super) async fn connect_pending_transport_with_initialize_retries( + pub(super) async fn connect_pending_transport_with_oauth_recovery( + &self, + initial_transport: PendingTransport, + client_service: ElicitationClientService, + timeout: Option, + ) -> Result<( + Arc>, + Option, + )> { + let mut initialize_deadline = timeout.map(|duration| Instant::now() + duration); + let mut attempt_context = InitializeAttemptContext::default(); + match self + .connect_pending_transport_with_initialize_retries( + initial_transport, + client_service.clone(), + timeout, + &mut initialize_deadline, + &mut attempt_context, + ) + .await + { + Ok(result) => Ok(result), + Err(error) => { + let Some(rejected_access_token) = + Self::rejected_access_token_from_initialize_error(&error) + else { + return Err(error); + }; + let Some(oauth) = attempt_context.oauth else { + return Err(error); + }; + // Initialization gets one provider refresh and one reconstructed transport. + // Reusing this wrapper for the retry would turn persistent 401s into a refresh + // loop. A later delayed 401 may rebuild once more only when it can adopt an + // already-committed newer token without contacting the provider. The startup + // deadline gates whether recovery starts and bounds transport setup plus retry + // handshakes, but the refresh transaction has its own bounds and is deliberately + // excluded from the startup budget. + remaining_initialize_timeout(timeout, initialize_deadline)?; + let refresh_started_at = Instant::now(); + let refresh_result = oauth + .persistor + .refresh_after_unauthorized(rejected_access_token) + .await; + if let Some(deadline) = initialize_deadline.as_mut() { + *deadline += refresh_started_at.elapsed(); + } + refresh_result?; + let remaining = remaining_initialize_timeout(timeout, initialize_deadline)?; + let transport = match remaining { + Some(remaining) => time::timeout( + remaining, + Self::create_pending_transport(&self.transport_recipe), + ) + .await + .map_err(|_| initialize_timeout_error(timeout, remaining))??, + None => Self::create_pending_transport(&self.transport_recipe).await?, + }; + let mut retry_context = InitializeAttemptContext::default(); + let result = self + .connect_pending_transport_with_initialize_retries( + transport, + client_service.clone(), + timeout, + &mut initialize_deadline, + &mut retry_context, + ) + .await; + if let Some(rejected_access_token) = result + .as_ref() + .err() + .and_then(Self::rejected_access_token_from_initialize_error) + { + let Some(retry_oauth) = retry_context.oauth else { + return Err(AuthError::AuthorizationRequired.into()); + }; + // A delayed B/401 can arrive after another process already committed C. + // Retry initialization once with C if it is now authoritative, but never + // contact the provider again from this one-refresh startup boundary. + let remaining = remaining_initialize_timeout(timeout, initialize_deadline)?; + let adoption = retry_oauth + .persistor + .adopt_newer_credentials_after_unauthorized(&rejected_access_token); + let adopted_newer_credentials = match remaining { + Some(remaining) => time::timeout(remaining, adoption) + .await + .map_err(|_| initialize_timeout_error(timeout, remaining))??, + None => adoption.await?, + }; + if adopted_newer_credentials { + let remaining = remaining_initialize_timeout(timeout, initialize_deadline)?; + let transport = match remaining { + Some(remaining) => time::timeout( + remaining, + Self::create_pending_transport(&self.transport_recipe), + ) + .await + .map_err(|_| initialize_timeout_error(timeout, remaining))??, + None => Self::create_pending_transport(&self.transport_recipe).await?, + }; + let mut adoption_context = InitializeAttemptContext::default(); + let adoption_result = self + .connect_pending_transport_with_initialize_retries( + transport, + client_service, + timeout, + &mut initialize_deadline, + &mut adoption_context, + ) + .await; + if adoption_result + .as_ref() + .err() + .and_then(Self::rejected_access_token_from_initialize_error) + .is_some() + { + return Err(AuthError::AuthorizationRequired.into()); + } + return adoption_result; + } + // The reconstructed transport rejected the still-authoritative refreshed + // token, so preserve Codex's established reauthentication signal. + return Err(AuthError::AuthorizationRequired.into()); + } + result + } + } + } + + async fn connect_pending_transport_with_initialize_retries( &self, initial_transport: PendingTransport, client_service: ElicitationClientService, timeout: Option, + initialize_deadline: &mut Option, + attempt_context: &mut InitializeAttemptContext, ) -> Result<( Arc>, - Option, + Option, )> { let should_retry = match &initial_transport { PendingTransport::InProcess { .. } | PendingTransport::Stdio { .. } => false, PendingTransport::StreamableHttp { .. } | PendingTransport::StreamableHttpWithOAuth { .. } => true, }; - let mut retry_deadline = timeout.map(|duration| Instant::now() + duration); let mut pending_transport = Some(initial_transport); for (attempt, retry_delay_ms) in STREAMABLE_HTTP_RETRY_DELAYS_MS @@ -50,7 +186,7 @@ impl RmcpClient { let transport = match pending_transport.take() { Some(transport) => transport, None => { - let remaining = remaining_initialize_timeout(timeout, retry_deadline)?; + let remaining = remaining_initialize_timeout(timeout, *initialize_deadline)?; match remaining { Some(remaining) => time::timeout( remaining, @@ -62,20 +198,25 @@ impl RmcpClient { } } }; - if let PendingTransport::StreamableHttpWithOAuth { - oauth_persistor, .. - } = &transport - { - // OAuth refresh has its own lock and provider request bounds. Exclude it from the - // MCP handshake budget, and finish persistence before attempting initialize. + if let PendingTransport::StreamableHttpWithOAuth { oauth, .. } = &transport { + // OAuth has independent bounds; pause the MCP handshake budget until refreshed + // credentials are durably committed. let refresh_started_at = Instant::now(); - oauth_persistor.refresh_if_needed().await?; - if let Some(deadline) = retry_deadline.as_mut() { + oauth.persistor.refresh_if_needed().await?; + if let Some(deadline) = initialize_deadline.as_mut() { *deadline += refresh_started_at.elapsed(); } } - let attempt_timeout = remaining_initialize_timeout(timeout, retry_deadline)?; - + // Keep the persistor paired with the transport attempt that returned 401. Rebuilt + // transports reuse the recipe's lifecycle-pinned credential source, and this pairing + // also keeps the authorization manager and snapshot aligned with the failed attempt. + attempt_context.oauth = match &transport { + PendingTransport::StreamableHttpWithOAuth { oauth, .. } => Some(oauth.clone()), + PendingTransport::InProcess { .. } + | PendingTransport::Stdio { .. } + | PendingTransport::StreamableHttp { .. } => None, + }; + let attempt_timeout = remaining_initialize_timeout(timeout, *initialize_deadline)?; match Self::connect_pending_transport( transport, client_service.clone(), @@ -96,7 +237,7 @@ impl RmcpClient { error = %error, "streamable HTTP MCP initialize failed with a retryable error; retrying" ); - if !sleep_with_retry_deadline(delay, retry_deadline).await { + if !sleep_with_retry_deadline(delay, *initialize_deadline).await { let duration = timeout.unwrap_or(delay); return Err(anyhow!( "timed out handshaking with MCP server after {duration:?}" @@ -121,6 +262,33 @@ impl RmcpClient { }) } + fn rejected_access_token_from_initialize_error(error: &anyhow::Error) -> Option { + error.chain().find_map(|source| { + source + .downcast_ref::() + .and_then(|error| { + Self::rejected_access_token_from_client_initialize_error(&error.source) + }) + .or_else(|| { + source + .downcast_ref::() + .and_then(Self::rejected_access_token_from_client_initialize_error) + }) + }) + } + + fn rejected_access_token_from_client_initialize_error( + error: &rmcp::service::ClientInitializeError, + ) -> Option { + match error { + rmcp::service::ClientInitializeError::TransportError { error, .. } => error + .error + .downcast_ref::>() + .and_then(Self::rejected_access_token), + _ => None, + } + } + fn is_retryable_client_initialize_error(error: &rmcp::service::ClientInitializeError) -> bool { match error { rmcp::service::ClientInitializeError::TransportError { error, context } @@ -171,6 +339,9 @@ impl RmcpClient { | StreamableHttpError::ServerDoesNotSupportSse | StreamableHttpError::Deserialize(_) | StreamableHttpError::Client(StreamableHttpClientAdapterError::SessionExpired404) + | StreamableHttpError::Client( + StreamableHttpClientAdapterError::AccessTokenRejected { .. }, + ) | StreamableHttpError::Client(StreamableHttpClientAdapterError::Header(_)) => false, _ => false, } diff --git a/codex-rs/rmcp-client/tests/streamable_http_oauth_internal.rs b/codex-rs/rmcp-client/tests/streamable_http_oauth_internal.rs new file mode 100644 index 000000000000..5574d3717ecb --- /dev/null +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_internal.rs @@ -0,0 +1,254 @@ +mod streamable_http_test_support; + +use std::path::PathBuf; +use std::time::Duration; + +use anyhow::Context; +use codex_config::types::AuthKeyringBackendKind; +use codex_config::types::OAuthCredentialsStoreMode; +use codex_exec_server::Environment; +use codex_rmcp_client::RmcpClient; +use codex_rmcp_client::StoredOAuthTokens; +use codex_rmcp_client::WrappedOAuthTokenResponse; +use codex_rmcp_client::save_oauth_tokens; +use oauth2::AccessToken; +use oauth2::RefreshToken; +use oauth2::basic::BasicTokenType; +use rmcp::transport::auth::OAuthTokenResponse; +use rmcp::transport::auth::VendorExtraTokenFields; +use serde_json::Value; +use serde_json::json; +use tempfile::TempDir; +use tokio::process::Command; +use wiremock::Mock; +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-internal"; +const SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_INTERNAL_SERVER_URL"; +const GET_FAILURE_MARKER_ENV: &str = "MCP_TEST_OAUTH_INTERNAL_GET_FAILURE_MARKER"; +const ACCESS_TOKEN_A: &str = "internal-access-a"; +const REFRESH_TOKEN_A: &str = "internal-refresh-a"; +const ACCESS_TOKEN_B: &str = "internal-access-b"; +const REFRESH_TOKEN_B: &str = "internal-refresh-b"; + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn rmcp_owned_get_reports_auth_failure_for_parent_recovery() -> anyhow::Result<()> { + let codex_home = TempDir::new()?; + let get_failure_marker = codex_home.path().join("get-failure-observed"); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + mount_refresh(&server, REFRESH_TOKEN_A, ACCESS_TOKEN_B, REFRESH_TOKEN_B).await; + + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header("authorization", format!("Bearer {ACCESS_TOKEN_A}"))) + .respond_with(|request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body.get("method").and_then(Value::as_str) { + Some("initialize") => initialize_response(&body), + Some("notifications/initialized") => ResponseTemplate::new(202), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + }) + .expect(2) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/mcp")) + .and(header("authorization", format!("Bearer {ACCESS_TOKEN_A}"))) + .respond_with({ + let get_failure_marker = get_failure_marker.clone(); + move |_request: &Request| { + std::fs::write(&get_failure_marker, b"observed") + .expect("record RMCP-owned GET auth failure"); + ResponseTemplate::new(401) + } + }) + // RMCP's default SSE reconnect policy is unbounded. The Codex failure latch must make + // this logical reconnect terminal instead of repeatedly re-entering get_stream with A. + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header("authorization", format!("Bearer {ACCESS_TOKEN_B}"))) + .respond_with(|request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body.get("method").and_then(Value::as_str) { + Some("initialize") => initialize_response(&body), + Some("notifications/initialized") => ResponseTemplate::new(202), + Some("tools/list") => ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { "tools": [] }, + })), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + }) + .expect(3) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/mcp")) + .and(header("authorization", format!("Bearer {ACCESS_TOKEN_B}"))) + .respond_with(ResponseTemplate::new(405)) + .expect(1) + .mount(&server) + .await; + Mock::given(method("DELETE")) + .and(path("/mcp")) + .and(header("authorization", format!("Bearer {ACCESS_TOKEN_B}"))) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(&server) + .await; + + let status = Command::new(std::env::current_exe()?) + .args([ + "oauth_internal_get_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(SERVER_URL_ENV, format!("{}/mcp", server.uri())) + .env(GET_FAILURE_MARKER_ENV, &get_failure_marker) + .status() + .await?; + anyhow::ensure!(status.success(), "OAuth internal child failed: {status}"); + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by rmcp_owned_get_reports_auth_failure_for_parent_recovery"] +async fn oauth_internal_get_child() -> anyhow::Result<()> { + let client = create_oauth_client().await?; + initialize_client(&client).await?; + wait_for_marker(GET_FAILURE_MARKER_ENV).await?; + + let tools = client + .list_tools(/*params*/ None, Some(Duration::from_secs(/*secs*/ 5))) + .await?; + assert!(tools.tools.is_empty()); + client.shutdown().await; + Ok(()) +} + +async fn wait_for_marker(env_name: &str) -> anyhow::Result<()> { + let path = PathBuf::from( + std::env::var_os(env_name).with_context(|| format!("missing {env_name} environment"))?, + ); + tokio::time::timeout(Duration::from_secs(/*secs*/ 5), async { + while !path.exists() { + tokio::time::sleep(Duration::from_millis(/*millis*/ 10)).await; + } + }) + .await + .with_context(|| format!("timed out waiting for marker {}", path.display())) +} + +async fn create_oauth_client() -> anyhow::Result { + let server_url = std::env::var(SERVER_URL_ENV)?; + save_initial_tokens(&server_url)?; + RmcpClient::new_streamable_http_client( + SERVER_NAME, + &server_url, + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await +} + +fn save_initial_tokens(server_url: &str) -> anyhow::Result<()> { + let mut response = OAuthTokenResponse::new( + AccessToken::new(ACCESS_TOKEN_A.to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + response.set_refresh_token(Some(RefreshToken::new(REFRESH_TOKEN_A.to_string()))); + response.set_expires_in(None); + save_oauth_tokens( + SERVER_NAME, + &StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: server_url.to_string(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: None, + }, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) +} + +async fn mount_oauth_metadata(server: &MockServer) { + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": ["scope-a"], + }))) + .mount(server) + .await; +} + +async fn mount_refresh( + server: &MockServer, + request_refresh_token: &str, + response_access_token: &str, + response_refresh_token: &str, +) { + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={request_refresh_token}" + ))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": response_access_token, + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": response_refresh_token, + "scope": "scope-a", + }))) + .expect(1) + .mount(server) + .await; +} + +fn initialize_response(body: &Value) -> ResponseTemplate { + ResponseTemplate::new(200) + .insert_header("mcp-session-id", "oauth-internal-session") + .set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { + "protocolVersion": body + .pointer("/params/protocolVersion") + .cloned() + .unwrap_or_else(|| json!("2025-06-18")), + "capabilities": {}, + "serverInfo": { + "name": "oauth-internal-test", + "version": "0.0.0-test" + } + } + })) +} 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 32f90bc41e8e..e01c7bb36c1b 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs @@ -1,5 +1,8 @@ mod streamable_http_test_support; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::time::Duration; use std::time::SystemTime; use std::time::UNIX_EPOCH; @@ -40,6 +43,11 @@ const SERVER_NAME: &str = "test-streamable-http-oauth-startup"; const EXPIRED_ACCESS_TOKEN: &str = "expired-access-token"; const REFRESH_TOKEN: &str = "valid-refresh-token"; const REFRESHED_ACCESS_TOKEN: &str = "refreshed-access-token"; +const ROTATED_REFRESH_TOKEN: &str = "rotated-refresh-token"; +const FINAL_ACCESS_TOKEN: &str = "final-access-token"; +const FINAL_REFRESH_TOKEN: &str = "final-refresh-token"; +const REJECTED_RETRY_ACCESS_TOKEN: &str = "rejected-retry-access-token"; +const REJECTED_RETRY_REFRESH_TOKEN: &str = "rejected-retry-refresh-token"; const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_STARTUP_SERVER_URL"; const UNREFRESHABLE_SERVER_URL: &str = "https://unrefreshable.example/mcp"; const UNEXPIRED_SERVER_URL: &str = "https://unexpired.example/mcp"; @@ -123,6 +131,290 @@ async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn recovers_initialization_and_operation_401_once() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": ["scope-a"], + }))) + .expect(1) + .mount(&server) + .await; + mount_refresh( + &server, + REFRESH_TOKEN, + REFRESHED_ACCESS_TOKEN, + ROTATED_REFRESH_TOKEN, + ) + .await; + mount_refresh( + &server, + ROTATED_REFRESH_TOKEN, + FINAL_ACCESS_TOKEN, + FINAL_REFRESH_TOKEN, + ) + .await; + mount_refresh( + &server, + FINAL_REFRESH_TOKEN, + REJECTED_RETRY_ACCESS_TOKEN, + REJECTED_RETRY_REFRESH_TOKEN, + ) + .await; + + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {EXPIRED_ACCESS_TOKEN}"), + )) + .respond_with(ResponseTemplate::new(401)) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {REFRESHED_ACCESS_TOKEN}"), + )) + .respond_with(|request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body.get("method").and_then(Value::as_str) { + Some("initialize") => initialize_response(&body), + Some("notifications/initialized") => ResponseTemplate::new(202), + Some("tools/list") => ResponseTemplate::new(401), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + }) + .expect(3) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {FINAL_ACCESS_TOKEN}"), + )) + .respond_with({ + let resource_attempts = Arc::new(AtomicUsize::new(0)); + move |request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body.get("method").and_then(Value::as_str) { + Some("tools/list") => ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { "tools": [] }, + })), + Some("resources/list") + if resource_attempts.fetch_add(1, Ordering::SeqCst) == 0 => + { + ResponseTemplate::new(404) + } + Some("resources/list") => ResponseTemplate::new(401) + .insert_header("www-authenticate", "Bearer realm=\"mcp\""), + Some("initialize") => initialize_response(&body), + Some("notifications/initialized") => ResponseTemplate::new(202), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + } + }) + .expect(5) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {REJECTED_RETRY_ACCESS_TOKEN}"), + )) + .respond_with( + ResponseTemplate::new(401).insert_header("www-authenticate", "Bearer realm=\"mcp\""), + ) + .expect(1) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + let status = Command::new(std::env::current_exe()?) + .args([ + "oauth_401_recovery_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(CHILD_SERVER_URL_ENV, format!("{}/mcp", server.uri())) + .status() + .await?; + assert!(status.success(), "OAuth recovery child failed: {status}"); + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn rejected_initialize_retry_requires_reauthentication() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": ["scope-a"], + }))) + .expect(1) + .mount(&server) + .await; + mount_refresh( + &server, + REFRESH_TOKEN, + REFRESHED_ACCESS_TOKEN, + ROTATED_REFRESH_TOKEN, + ) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {EXPIRED_ACCESS_TOKEN}"), + )) + .respond_with( + ResponseTemplate::new(401).insert_header("www-authenticate", "Bearer realm=\"mcp\""), + ) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {REFRESHED_ACCESS_TOKEN}"), + )) + .respond_with( + ResponseTemplate::new(401).insert_header("www-authenticate", "Bearer realm=\"mcp\""), + ) + .expect(1) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + let status = Command::new(std::env::current_exe()?) + .args([ + "oauth_rejected_initialize_retry_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(CHILD_SERVER_URL_ENV, format!("{}/mcp", server.uri())) + .status() + .await?; + assert!( + status.success(), + "OAuth rejected-retry child failed: {status}" + ); + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn operation_timeout_bounds_unauthorized_refresh_wait() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": ["scope-a"], + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={REFRESH_TOKEN}" + ))) + .respond_with( + ResponseTemplate::new(200) + .set_delay(Duration::from_millis(/*millis*/ 500)) + .set_body_json(json!({ + "access_token": REFRESHED_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": ROTATED_REFRESH_TOKEN, + "scope": "scope-a", + })), + ) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {EXPIRED_ACCESS_TOKEN}"), + )) + .respond_with(|request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body.get("method").and_then(Value::as_str) { + Some("initialize") => initialize_response(&body), + Some("notifications/initialized") => ResponseTemplate::new(202), + Some("tools/list") => ResponseTemplate::new(401), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + }) + // Three requests always use A: initialize, initialized, and the first tools/list. The next + // operation may wait for the in-flight authorization-manager guard and send B directly, + // or it may send A once, receive 401, and join the same refresh transaction. Both + // interleavings are valid. The exact provider and B expectations below prove that neither + // path performs a second refresh or sends more than one successful retry. + .expect(3..=4) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {REFRESHED_ACCESS_TOKEN}"), + )) + .respond_with(|request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { "tools": [] }, + })) + }) + .expect(1) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + let status = Command::new(std::env::current_exe()?) + .args([ + "oauth_401_timeout_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(CHILD_SERVER_URL_ENV, format!("{}/mcp", server.uri())) + .status() + .await?; + assert!(status.success(), "OAuth timeout child failed: {status}"); + server.verify().await; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn reports_auth_status_for_persisted_credentials() -> anyhow::Result<()> { let codex_home = TempDir::new()?; @@ -236,7 +528,9 @@ async fn persisted_credentials_auth_status_child() -> anyhow::Result<()> { url: UNEXPIRED_SERVER_URL.to_string(), client_id: "test-client-id".to_string(), token_response: WrappedOAuthTokenResponse(response), - expires_at: Some(now.saturating_add(/*rhs*/ 60_000)), + // Keep this outside the 60-second proactive refresh guard band. The test is checking a + // healthy persisted access token, not the boundary where a refresh becomes necessary. + expires_at: Some(now.saturating_add(/*rhs*/ 120_000)), }; save_oauth_tokens( SERVER_NAME, @@ -376,3 +670,142 @@ async fn expired_unrefreshable_startup_child() -> anyhow::Result<()> { assert!(is_authentication_required_error(&error)); Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by recovers_initialization_and_operation_401_once"] +async fn oauth_401_recovery_child() -> anyhow::Result<()> { + let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; + let client = refreshable_oauth_client(&server_url).await?; + initialize_client(&client).await?; + let tools = client + .list_tools(/*params*/ None, Some(Duration::from_secs(/*secs*/ 5))) + .await?; + assert!(tools.tools.is_empty()); + + let error = client + .list_resources(/*params*/ None, Some(Duration::from_secs(/*secs*/ 5))) + .await + .expect_err("a rejected one-shot OAuth retry should require reauthentication"); + assert!(is_authentication_required_error(&error)); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by rejected_initialize_retry_requires_reauthentication"] +async fn oauth_rejected_initialize_retry_child() -> anyhow::Result<()> { + let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; + let client = refreshable_oauth_client(&server_url).await?; + let error = initialize_client(&client) + .await + .expect_err("a rejected initialize retry should require reauthentication"); + assert!(is_authentication_required_error(&error)); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by operation_timeout_bounds_unauthorized_refresh_wait"] +async fn oauth_401_timeout_child() -> anyhow::Result<()> { + let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; + let client = refreshable_oauth_client(&server_url).await?; + initialize_client(&client).await?; + + let error = client + .list_tools( + /*params*/ None, + Some(Duration::from_millis(/*millis*/ 50)), + ) + .await + .expect_err("operation deadline should expire before the delayed refresh"); + assert!( + error.to_string().contains("timed out awaiting tools/list"), + "unexpected operation error: {error:#}" + ); + // The caller stopped waiting, but the owned refresh transaction still holds the credential + // lock. Starting the next operation immediately makes its preflight wait for that transaction, + // then adopt the persisted token without another provider request. + let tools = client + .list_tools(/*params*/ None, Some(Duration::from_secs(/*secs*/ 5))) + .await?; + assert!(tools.tools.is_empty()); + Ok(()) +} + +async fn refreshable_oauth_client(server_url: &str) -> anyhow::Result { + let mut response = OAuthTokenResponse::new( + AccessToken::new(EXPIRED_ACCESS_TOKEN.to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + response.set_refresh_token(Some(RefreshToken::new(REFRESH_TOKEN.to_string()))); + response.set_expires_in(None); + save_oauth_tokens( + SERVER_NAME, + &StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: server_url.to_string(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: None, + }, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + + let client = RmcpClient::new_streamable_http_client( + SERVER_NAME, + server_url, + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await?; + Ok(client) +} + +async fn mount_refresh( + server: &MockServer, + request_refresh_token: &str, + response_access_token: &str, + response_refresh_token: &str, +) { + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={request_refresh_token}" + ))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": response_access_token, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": response_refresh_token, + "scope": "scope-a", + }))) + .expect(1) + .mount(server) + .await; +} + +fn initialize_response(body: &Value) -> ResponseTemplate { + ResponseTemplate::new(200) + .insert_header("mcp-session-id", "oauth-recovery-session") + .set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { + "protocolVersion": body + .pointer("/params/protocolVersion") + .cloned() + .unwrap_or_else(|| json!("2025-06-18")), + "capabilities": {}, + "serverInfo": { + "name": "oauth-401-recovery-test", + "version": "0.0.0-test", + }, + }, + })) +} diff --git a/codex-rs/rmcp-client/tests/streamable_http_recovery.rs b/codex-rs/rmcp-client/tests/streamable_http_recovery.rs index 29d2404ede9e..a53d0cb08109 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_recovery.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_recovery.rs @@ -11,6 +11,8 @@ use codex_exec_server::HttpClient; use codex_exec_server::HttpRequestParams; use codex_exec_server::HttpRequestResponse; use codex_exec_server::HttpResponseBodyStream; +use codex_rmcp_client::RmcpClient; +use codex_rmcp_client::is_authentication_required_error; use futures::FutureExt as _; use futures::future::BoxFuture; use pretty_assertions::assert_eq; @@ -124,7 +126,13 @@ async fn streamable_http_initialize_retries_remote_no_response_error() -> anyhow async fn streamable_http_initialize_retries_transient_http_status() -> anyhow::Result<()> { let (_server, base_url) = spawn_streamable_http_server().await?; - arm_initialize_post_failure(&base_url, /*status*/ 502, /*remaining*/ 1).await?; + arm_initialize_post_failure( + &base_url, + /*status*/ 502, + /*remaining*/ 1, + /*www_authenticate_headers*/ &[], + ) + .await?; let client = create_client(&base_url).await?; let result = call_echo_tool(&client, "after-status-retry").await?; @@ -301,6 +309,37 @@ async fn streamable_http_401_does_not_trigger_recovery() -> anyhow::Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn static_bearer_initialize_401_requires_authentication() -> anyhow::Result<()> { + let (_server, base_url) = spawn_streamable_http_server().await?; + arm_initialize_post_failure( + &base_url, + /*status*/ 401, + /*remaining*/ 1, + /*www_authenticate_headers*/ &[r#"Bearer realm="mcp""#], + ) + .await?; + + let client = RmcpClient::new_streamable_http_client( + "test-static-bearer-401", + &format!("{base_url}/mcp"), + Some("test-bearer".to_string()), + /*http_headers*/ None, + /*env_http_headers*/ None, + codex_config::types::OAuthCredentialsStoreMode::File, + codex_config::types::AuthKeyringBackendKind::default(), + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await?; + let error = streamable_http_test_support::initialize_client(&client) + .await + .expect_err("a challenged static bearer token should require authentication"); + + assert!(is_authentication_required_error(&error)); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn streamable_http_403_scope_challenge_returns_insufficient_scope() -> anyhow::Result<()> { let (_server, base_url) = spawn_streamable_http_server().await?; diff --git a/codex-rs/rmcp-client/tests/streamable_http_test_support.rs b/codex-rs/rmcp-client/tests/streamable_http_test_support.rs index 87a52f55dff6..30d3039f5bb7 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_test_support.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_test_support.rs @@ -258,12 +258,14 @@ pub(crate) async fn arm_initialize_post_failure( base_url: &str, status: u16, remaining: usize, + www_authenticate_headers: &[&str], ) -> anyhow::Result<()> { let response = reqwest::Client::new() .post(format!("{base_url}{INITIALIZE_POST_FAILURE_CONTROL_PATH}")) .json(&json!({ "status": status, "remaining": remaining, + "www_authenticate_headers": www_authenticate_headers, })) .send() .await?;