From aaebdab06778eeefe3de5a1eac63fc24f74b1eca Mon Sep 17 00:00:00 2001 From: Steven Lee Date: Fri, 10 Jul 2026 02:35:30 +0000 Subject: [PATCH] Keep MCP OAuth refresh ownership in Codex --- .../rmcp-client/src/http_client_adapter.rs | 2 - .../src/oauth/refresh_transaction.rs | 41 +++ .../src/oauth/tests/persistor_tests.rs | 76 ++++++ codex-rs/rmcp-client/src/oauth_transport.rs | 256 ++++++++---------- .../rmcp-client/src/oauth_transport_tests.rs | 150 ++++------ codex-rs/rmcp-client/src/rmcp_client.rs | 181 +++++++++---- .../rmcp-client/src/streamable_http_retry.rs | 101 +++++-- .../tests/streamable_http_oauth_internal.rs | 86 +++--- 8 files changed, 524 insertions(+), 369 deletions(-) diff --git a/codex-rs/rmcp-client/src/http_client_adapter.rs b/codex-rs/rmcp-client/src/http_client_adapter.rs index eccce8c03d..75bb42a5e3 100644 --- a/codex-rs/rmcp-client/src/http_client_adapter.rs +++ b/codex-rs/rmcp-client/src/http_client_adapter.rs @@ -65,8 +65,6 @@ pub(crate) enum StreamableHttpClientAdapterError { SessionExpired404, #[error("MCP server rejected the access token with HTTP 401 Unauthorized")] AccessTokenRejected { rejected_access_token: AccessToken }, - #[error("MCP OAuth operation failed: {0:#}")] - OAuth(#[source] anyhow::Error), #[error(transparent)] HttpRequest(#[from] ExecServerError), #[error("invalid HTTP header: {0}")] diff --git a/codex-rs/rmcp-client/src/oauth/refresh_transaction.rs b/codex-rs/rmcp-client/src/oauth/refresh_transaction.rs index 9bd03cc401..8420b63ad1 100644 --- a/codex-rs/rmcp-client/src/oauth/refresh_transaction.rs +++ b/codex-rs/rmcp-client/src/oauth/refresh_transaction.rs @@ -54,6 +54,47 @@ impl OAuthPersistor { .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); + } + + 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, 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 c07be4af65..ff9121c628 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; @@ -209,6 +210,81 @@ async fn delayed_unauthorized_retries_adopt_the_winning_token() -> Result<()> { 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" diff --git a/codex-rs/rmcp-client/src/oauth_transport.rs b/codex-rs/rmcp-client/src/oauth_transport.rs index e5ea3fc2f1..ad4a9fadb9 100644 --- a/codex-rs/rmcp-client/src/oauth_transport.rs +++ b/codex-rs/rmcp-client/src/oauth_transport.rs @@ -1,34 +1,35 @@ //! Codex-owned OAuth policy for RMCP Streamable HTTP traffic. //! -//! RMCP remains responsible for transport mechanics and bearer-token injection. Codex owns the -//! credential lifecycle: every POST, SSE GET/reconnect, and session DELETE receives proactive -//! refresh from its owning Codex layer, and each path has at most one 401 recovery. The -//! authorization manager only receives request-safe credentials, so it cannot independently -//! refresh outside Codex's serialized transaction. +//! 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. //! -//! POST recovery is split at an intentional ownership boundary. Client-originated requests and -//! notifications retain their outer `RmcpClient` recovery, which knows the startup/tool deadline -//! and can avoid replaying a request after its caller timed out. RMCP-owned responses to -//! server-initiated requests have no such outer operation, so they recover here. GET/reconnect and -//! DELETE are always RMCP-owned and also recover here. +//! 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::auth::AuthError; +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 tracing::debug; use crate::http_client_adapter::StreamableHttpClientAdapter; use crate::http_client_adapter::StreamableHttpClientAdapterError; -use crate::oauth::OAuthPersistor; type TransportResult = std::result::Result>; @@ -36,53 +37,97 @@ type TransportResult = #[derive(Clone)] pub(crate) struct OAuthTransportClient { auth_client: AuthClient, - persistor: OAuthPersistor, + failure_state: OAuthTransportFailureState, } impl OAuthTransportClient { pub(crate) fn new( auth_client: AuthClient, - persistor: OAuthPersistor, + failure_state: OAuthTransportFailureState, ) -> Self { Self { auth_client, - persistor, + 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 persistor(&self) -> OAuthPersistor { - self.persistor.clone() + pub(crate) fn retry_policy(&self) -> OAuthSseRetryPolicy { + OAuthSseRetryPolicy { + failure_state: self.clone(), + fallback: ExponentialBackoff::default(), + } } +} - async fn preflight(&self, operation: &'static str) -> TransportResult<()> { - debug!( - operation, - "checking MCP OAuth credentials before transport request" - ); - self.persistor - .refresh_if_needed() - .await - .map_err(oauth_transport_error) - } +#[derive(Debug)] +pub(crate) struct OAuthSseRetryPolicy { + failure_state: OAuthTransportFailureState, + fallback: ExponentialBackoff, +} - async fn recover_after_unauthorized( - &self, - operation: &'static str, - rejected_access_token: Option, - ) -> TransportResult { - let Some(rejected_access_token) = rejected_access_token else { - return Ok(false); - }; - - debug!( - operation, - "recovering once after MCP transport rejected an OAuth access token" - ); - self.persistor - .refresh_after_unauthorized(rejected_access_token) - .await - .map_err(oauth_transport_error)?; - Ok(true) +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) + } } } @@ -101,40 +146,22 @@ impl StreamableHttpClient for OAuthTransportClient { message, JsonRpcMessage::Response(_) | JsonRpcMessage::Error(_) ); - if is_rmcp_owned_response { - self.preflight("post_message").await?; - } let result = self .auth_client - .post_message( - Arc::clone(&uri), - message.clone(), - session_id.clone(), - auth_token.clone(), - custom_headers.clone(), - ) + .post_message(uri, message, session_id, auth_token, custom_headers) .await; - // RMCP queues client-originated requests independently of the caller waiting on them. If - // recovery happened here, a timed-out public tool call could still be replayed after its - // refresh finished. The outer RmcpClient path owns those deadlines. Responses to - // server-initiated requests have no outer operation and therefore recover here. - if !is_rmcp_owned_response { - return result; - } - let rejected_access_token = result.as_ref().err().and_then(rejected_access_token); - if self - .recover_after_unauthorized("post_message", rejected_access_token) - .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) { - authorization_required_after_retry( - self.auth_client - .post_message(uri, message, session_id, auth_token, custom_headers) - .await, - ) - } else { - result + self.failure_state + .record_rejected_access_token(rejected_access_token); } + result } async fn delete_session( @@ -144,29 +171,15 @@ impl StreamableHttpClient for OAuthTransportClient { auth_token: Option, custom_headers: HashMap, ) -> TransportResult<()> { - self.preflight("delete_session").await?; let result = self .auth_client - .delete_session( - Arc::clone(&uri), - Arc::clone(&session_id), - auth_token.clone(), - custom_headers.clone(), - ) + .delete_session(uri, session_id, auth_token, custom_headers) .await; - let rejected_access_token = result.as_ref().err().and_then(rejected_access_token); - if self - .recover_after_unauthorized("delete_session", rejected_access_token) - .await? - { - authorization_required_after_retry( - self.auth_client - .delete_session(uri, session_id, auth_token, custom_headers) - .await, - ) - } else { - result + 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( @@ -179,43 +192,15 @@ impl StreamableHttpClient for OAuthTransportClient { ) -> TransportResult< futures::stream::BoxStream<'static, Result>, > { - self.preflight("get_stream").await?; let result = self .auth_client - .get_stream( - Arc::clone(&uri), - Arc::clone(&session_id), - last_event_id.clone(), - auth_token.clone(), - custom_headers.clone(), - ) + .get_stream(uri, session_id, last_event_id, auth_token, custom_headers) .await; - let rejected_access_token = result.as_ref().err().and_then(rejected_access_token); - if self - .recover_after_unauthorized("get_stream", rejected_access_token) - .await? - { - authorization_required_after_retry( - self.auth_client - .get_stream(uri, session_id, last_event_id, auth_token, custom_headers) - .await, - ) - } else { - result + 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); } - } -} - -fn authorization_required_after_retry(result: TransportResult) -> TransportResult { - match result { - // The first 401 carries the token that was actually rejected so concurrent recovery can - // distinguish A from a newer B. Once the single retry also rejects B, attribution is no - // longer useful: surface the existing reauthentication marker instead of leaking the - // adapter-only error past the Codex-owned recovery boundary. - Err(StreamableHttpError::Client( - StreamableHttpClientAdapterError::AccessTokenRejected { .. }, - )) => Err(StreamableHttpError::Auth(AuthError::AuthorizationRequired)), - result => result, + result } } @@ -230,25 +215,6 @@ fn rejected_access_token( } } -fn oauth_transport_error( - error: anyhow::Error, -) -> StreamableHttpError { - if let Some(auth_error) = - error - .chain() - .find_map(|source| match source.downcast_ref::() { - Some(AuthError::AuthorizationRequired) => Some(AuthError::AuthorizationRequired), - Some(AuthError::TokenExpired) => Some(AuthError::TokenExpired), - _ => None, - }) - { - // Preserve RMCP's established reauthentication variants across Codex's transport policy - // boundary. Other OAuth failures retain their context-rich adapter error. - return StreamableHttpError::Auth(auth_error); - } - StreamableHttpError::Client(StreamableHttpClientAdapterError::OAuth(error)) -} - #[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 index acbca36c10..258b0c68f8 100644 --- a/codex-rs/rmcp-client/src/oauth_transport_tests.rs +++ b/codex-rs/rmcp-client/src/oauth_transport_tests.rs @@ -1,80 +1,41 @@ use std::collections::HashMap; use std::sync::Arc; -use codex_config::types::AuthKeyringBackendKind; -use codex_config::types::OAuthCredentialsStoreMode; -use codex_exec_server::Environment; 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::AuthError; 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 rmcp::transport::streamable_http_client::StreamableHttpPostResponse; use serde_json::json; -use tempfile::TempDir; -use tokio::process::Command; use wiremock::Mock; use wiremock::MockServer; use wiremock::ResponseTemplate; -use wiremock::matchers::body_string_contains; use wiremock::matchers::header; use wiremock::matchers::method; use wiremock::matchers::path; use super::OAuthTransportClient; -use super::authorization_required_after_retry; -use super::oauth_transport_error; +use super::OAuthTransportFailureState; use crate::http_client_adapter::StreamableHttpClientAdapter; use crate::http_client_adapter::StreamableHttpClientAdapterError; -use crate::oauth::OAuthPersistor; -use crate::oauth::ResolvedOAuthCredentialStore; use crate::oauth::StoredOAuthTokens; use crate::oauth::WrappedOAuthTokenResponse; use crate::oauth::request_oauth_token_response; -use crate::oauth::save_oauth_tokens; use crate::oauth_http_client::OAuthHttpClientAdapter; const SERVER_NAME: &str = "oauth-transport-response-test"; -const SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_RESPONSE_SERVER_URL"; const ACCESS_TOKEN_A: &str = "response-access-a"; const REFRESH_TOKEN_A: &str = "response-refresh-a"; -const ACCESS_TOKEN_B: &str = "response-access-b"; -const REFRESH_TOKEN_B: &str = "response-refresh-b"; - -#[test] -fn exhausted_transport_retry_requires_reauthentication() { - let result = authorization_required_after_retry::<()>(Err(StreamableHttpError::Client( - StreamableHttpClientAdapterError::AccessTokenRejected { - rejected_access_token: AccessToken::new(ACCESS_TOKEN_B.to_string()), - }, - ))); - - assert!(matches!( - result, - Err(StreamableHttpError::Auth(AuthError::AuthorizationRequired)) - )); -} - -#[test] -fn oauth_transport_preserves_reauthentication_errors() { - let error = anyhow::Error::new(AuthError::AuthorizationRequired) - .context("refreshing rejected MCP access token"); - - assert!(matches!( - oauth_transport_error(error), - StreamableHttpError::Auth(AuthError::AuthorizationRequired) - )); -} #[tokio::test] -async fn server_response_post_receives_one_shot_oauth_recovery() -> anyhow::Result<()> { +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")) @@ -87,18 +48,8 @@ async fn server_response_post_receives_one_shot_oauth_recovery() -> anyhow::Resu .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_A}" - ))) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "access_token": ACCESS_TOKEN_B, - "token_type": "Bearer", - "expires_in": 3600, - "refresh_token": REFRESH_TOKEN_B, - "scope": "scope-a", - }))) - .expect(1) + .respond_with(ResponseTemplate::new(500)) + .expect(0) .mount(&server) .await; Mock::given(method("POST")) @@ -108,44 +59,10 @@ async fn server_response_post_receives_one_shot_oauth_recovery() -> anyhow::Resu .expect(1) .mount(&server) .await; - Mock::given(method("POST")) - .and(path("/mcp")) - .and(header("authorization", format!("Bearer {ACCESS_TOKEN_B}"))) - .respond_with(ResponseTemplate::new(202)) - .expect(1) - .mount(&server) - .await; - let codex_home = TempDir::new()?; - let status = Command::new(std::env::current_exe()?) - .args([ - "oauth_transport::tests::server_response_post_child", - "--exact", - "--ignored", - "--nocapture", - ]) - .env("CODEX_HOME", codex_home.path()) - .env(SERVER_URL_ENV, format!("{}/mcp", server.uri())) - .status() - .await?; - anyhow::ensure!(status.success(), "OAuth response child failed: {status}"); - server.verify().await; - Ok(()) -} - -#[tokio::test] -#[ignore = "spawned by server_response_post_receives_one_shot_oauth_recovery"] -async fn server_response_post_child() -> anyhow::Result<()> { - let server_url = std::env::var(SERVER_URL_ENV)?; + let server_url = format!("{}/mcp", server.uri()); let initial_tokens = initial_tokens(&server_url); - save_oauth_tokens( - SERVER_NAME, - &initial_tokens, - OAuthCredentialsStoreMode::File, - AuthKeyringBackendKind::default(), - )?; - - let http_client = Environment::default_for_tests().get_http_client(); + 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(), @@ -171,14 +88,8 @@ async fn server_response_post_child() -> anyhow::Result<()> { .with_rejected_token_attribution(), manager, ); - let persistor = OAuthPersistor::new( - SERVER_NAME.to_string(), - server_url.clone(), - Arc::clone(&auth_client.auth_manager), - ResolvedOAuthCredentialStore::File, - Some(initial_tokens), - ); - let client = OAuthTransportClient::new(auth_client, persistor); + 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", @@ -188,7 +99,7 @@ async fn server_response_post_child() -> anyhow::Result<()> { } }))?; - let response = client + let error = client .post_message( Arc::from(server_url), response_message, @@ -196,12 +107,49 @@ async fn server_response_post_child() -> anyhow::Result<()> { /*auth_token*/ None, HashMap::new(), ) - .await?; + .await + .expect_err("the server should reject the response token"); - assert!(matches!(response, StreamableHttpPostResponse::Accepted)); + 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()), diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 99f95f486a..f5d02f72eb 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -75,6 +75,7 @@ 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; @@ -101,17 +102,24 @@ enum PendingTransport { }, StreamableHttpWithOAuth { transport: StreamableHttpClientTransport, - oauth_persistor: OAuthPersistor, + 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, } @@ -134,7 +142,7 @@ enum TransportRecipe { store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, pinned_credential_store: Arc>, - oauth_client: Arc>, + oauth_runtime: Arc>, http_client: Arc, auth_provider: Option, }, @@ -412,7 +420,7 @@ impl RmcpClient { store_mode, keyring_backend_kind, pinned_credential_store: Arc::new(OnceLock::new()), - oauth_client: Arc::new(OnceLock::new()), + oauth_runtime: Arc::new(OnceLock::new()), http_client, auth_provider, }; @@ -455,7 +463,7 @@ impl RmcpClient { } }; - let (service, oauth_persistor) = self + let (service, oauth) = self .connect_pending_transport_with_oauth_recovery( pending_transport, client_service.clone(), @@ -484,7 +492,7 @@ impl RmcpClient { } *guard = ClientState::Ready { service, - oauth: oauth_persistor.clone(), + oauth: oauth.clone(), }; } @@ -704,11 +712,11 @@ impl RmcpClient { Ok(response) } - async fn service_and_oauth_persistor( + async fn service_and_oauth_runtime( &self, ) -> Result<( Arc>, - Option, + Option, )> { let guard = self.state.lock().await; match &*guard { @@ -718,7 +726,7 @@ impl RmcpClient { } } - async fn oauth_persistor(&self) -> Option { + async fn oauth_runtime(&self) -> Option { let guard = self.state.lock().await; match &*guard { ClientState::Ready { @@ -746,12 +754,45 @@ impl RmcpClient { } async fn refresh_oauth_if_needed(&self) -> Result<()> { - if let Some(runtime) = self.oauth_persistor().await { - runtime.refresh_if_needed().await?; + if self.recover_pending_oauth_failure().await? { + return Ok(()); + } + if let Some(runtime) = self.oauth_runtime().await { + runtime.persistor.refresh_if_needed().await?; } Ok(()) } + /// 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), + } + } + async fn create_pending_transport( transport_recipe: &TransportRecipe, ) -> Result { @@ -775,7 +816,7 @@ impl RmcpClient { store_mode, keyring_backend_kind, pinned_credential_store, - oauth_client, + oauth_runtime, http_client, auth_provider, } => { @@ -791,15 +832,14 @@ impl RmcpClient { // 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(oauth_client) = oauth_client.get() { - let runtime = oauth_client.persistor(); + if let Some(runtime) = oauth_runtime.get() { let transport = StreamableHttpClientTransport::with_client( - oauth_client.clone(), - StreamableHttpClientTransportConfig::with_uri(url.clone()), + runtime.transport_client.clone(), + oauth_transport_config(url, &runtime.failure_state), ); return Ok(PendingTransport::StreamableHttpWithOAuth { transport, - oauth_persistor: runtime, + oauth: runtime.clone(), }); } @@ -850,7 +890,7 @@ impl RmcpClient { store: credential_store, }) = resolved_oauth_tokens { - match create_oauth_transport_client( + match create_oauth_runtime( server_name, url, initial_tokens.clone(), @@ -860,22 +900,21 @@ impl RmcpClient { ) .await { - Ok(resolved_oauth_client) => { - oauth_client - .set(resolved_oauth_client.clone()) + Ok(runtime) => { + oauth_runtime + .set(runtime.clone()) .map_err(|_| { anyhow!( - "OAuth client resolved concurrently for MCP server `{server_name}`" + "OAuth runtime resolved concurrently for MCP server `{server_name}`" ) })?; - let oauth_persistor = resolved_oauth_client.persistor(); let transport = StreamableHttpClientTransport::with_client( - resolved_oauth_client, - StreamableHttpClientTransportConfig::with_uri(url.clone()), + runtime.transport_client.clone(), + oauth_transport_config(url, &runtime.failure_state), ); Ok(PendingTransport::StreamableHttpWithOAuth { transport, - oauth_persistor, + oauth: runtime, }) } Err(err) @@ -934,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, @@ -949,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), ), }; @@ -973,7 +1009,7 @@ impl RmcpClient { }; let service = service_result?; - Ok((Arc::new(service), oauth_persistor)) + Ok((Arc::new(service), oauth)) } async fn run_service_operation( @@ -987,11 +1023,12 @@ impl RmcpClient { Fut: std::future::Future>, { let deadline = timeout.map(|duration| Instant::now() + duration); - // Keep the OAuth persistor paired with the service that performs this operation. Session + // 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_persistor) = self.service_and_oauth_persistor().await?; + 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 { @@ -1010,21 +1047,50 @@ impl RmcpClient { .err() .and_then(Self::rejected_access_token_from_operation_error) { - if oauth_recovered { - // The rejected token is needed only to attribute the first 401. A second 401 - // after the one allowed refresh means this lifecycle needs reauthentication. - return Err(AuthError::AuthorizationRequired.into()); - } - let Some(oauth_persistor) = oauth_persistor.as_ref() else { + 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_persistor.refresh_after_unauthorized(rejected_access_token); + 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, @@ -1057,7 +1123,7 @@ impl RmcpClient { // 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?; - (service, oauth_persistor) = self.service_and_oauth_persistor().await?; + (service, oauth_runtime) = self.service_and_oauth_runtime().await?; session_recovered = true; continue; } @@ -1244,7 +1310,7 @@ 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 + let (service, oauth) = self .connect_pending_transport_with_oauth_recovery( pending_transport, initialize_context.client_service, @@ -1259,7 +1325,7 @@ impl RmcpClient { } *guard = ClientState::Ready { service, - oauth: oauth_persistor.clone(), + oauth: oauth.clone(), }; } @@ -1267,14 +1333,14 @@ impl RmcpClient { } } -async fn create_oauth_transport_client( +async fn create_oauth_runtime( server_name: &str, url: &str, initial_tokens: StoredOAuthTokens, credential_store: ResolvedOAuthCredentialStore, default_headers: HeaderMap, http_client: Arc, -) -> Result { +) -> Result { let oauth_http_client = Arc::new(OAuthHttpClientAdapter::new( http_client.clone(), default_headers.clone(), @@ -1304,15 +1370,30 @@ async fn create_oauth_transport_client( ); let auth_manager = auth_client.auth_manager.clone(); - 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(OAuthTransportClient::new(auth_client, runtime)) + Ok(OAuthRuntime { + transport_client, + persistor, + failure_state, + }) +} + +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/streamable_http_retry.rs b/codex-rs/rmcp-client/src/streamable_http_retry.rs index 9bff1d548a..12778aa349 100644 --- a/codex-rs/rmcp-client/src/streamable_http_retry.rs +++ b/codex-rs/rmcp-client/src/streamable_http_retry.rs @@ -14,19 +14,18 @@ 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_persistor: Option, + oauth: Option, } impl RmcpClient { @@ -37,7 +36,7 @@ impl RmcpClient { timeout: Option, ) -> Result<( Arc>, - Option, + Option, )> { let mut initialize_deadline = timeout.map(|duration| Instant::now() + duration); let mut attempt_context = InitializeAttemptContext::default(); @@ -58,17 +57,20 @@ impl RmcpClient { else { return Err(error); }; - let Some(oauth_persistor) = attempt_context.oauth_persistor else { + let Some(oauth) = attempt_context.oauth else { return Err(error); }; - // Initialization gets one OAuth refresh and one reconstructed transport. Reusing - // this wrapper for the retry would turn persistent 401s into a refresh loop. The - // startup deadline gates whether recovery starts and bounds transport setup plus - // the retry handshake, but the refresh transaction has its own bounds and is - // deliberately excluded from the startup budget. + // 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 + let refresh_result = oauth + .persistor .refresh_after_unauthorized(rejected_access_token) .await; if let Some(deadline) = initialize_deadline.as_mut() { @@ -89,21 +91,66 @@ impl RmcpClient { let result = self .connect_pending_transport_with_initialize_retries( transport, - client_service, + client_service.clone(), timeout, &mut initialize_deadline, &mut retry_context, ) .await; - if result + if let Some(rejected_access_token) = result .as_ref() .err() .and_then(Self::rejected_access_token_from_initialize_error) - .is_some() { - // The first 401 identifies which access token failed. If the reconstructed - // transport still rejects the refreshed token, preserve Codex's established - // signal that the user must authenticate again. + 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 @@ -120,7 +167,7 @@ impl RmcpClient { attempt_context: &mut InitializeAttemptContext, ) -> Result<( Arc>, - Option, + Option, )> { let should_retry = match &initial_transport { PendingTransport::InProcess { .. } | PendingTransport::Stdio { .. } => false, @@ -151,14 +198,11 @@ impl RmcpClient { } } }; - if let PendingTransport::StreamableHttpWithOAuth { - oauth_persistor, .. - } = &transport - { + 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?; + oauth.persistor.refresh_if_needed().await?; if let Some(deadline) = initialize_deadline.as_mut() { *deadline += refresh_started_at.elapsed(); } @@ -166,10 +210,8 @@ impl RmcpClient { // 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_persistor = match &transport { - PendingTransport::StreamableHttpWithOAuth { - oauth_persistor, .. - } => Some(oauth_persistor.clone()), + attempt_context.oauth = match &transport { + PendingTransport::StreamableHttpWithOAuth { oauth, .. } => Some(oauth.clone()), PendingTransport::InProcess { .. } | PendingTransport::Stdio { .. } | PendingTransport::StreamableHttp { .. } => None, @@ -300,7 +342,6 @@ impl RmcpClient { | StreamableHttpError::Client( StreamableHttpClientAdapterError::AccessTokenRejected { .. }, ) - | StreamableHttpError::Client(StreamableHttpClientAdapterError::OAuth(_)) | 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 index fc40872a15..5574d3717e 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_internal.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_internal.rs @@ -33,24 +33,19 @@ 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_RETRY_MARKER_ENV: &str = "MCP_TEST_OAUTH_INTERNAL_GET_RETRY_MARKER"; -const DELETE_RETRY_MARKER_ENV: &str = "MCP_TEST_OAUTH_INTERNAL_DELETE_RETRY_MARKER"; +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"; -const ACCESS_TOKEN_C: &str = "internal-access-c"; -const REFRESH_TOKEN_C: &str = "internal-refresh-c"; #[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn rmcp_owned_get_and_delete_receive_oauth_recovery() -> anyhow::Result<()> { +async fn rmcp_owned_get_reports_auth_failure_for_parent_recovery() -> anyhow::Result<()> { let codex_home = TempDir::new()?; - let get_retry_marker = codex_home.path().join("get-retry-observed"); - let delete_retry_marker = codex_home.path().join("delete-retry-observed"); + 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; - mount_refresh(&server, REFRESH_TOKEN_B, ACCESS_TOKEN_C, REFRESH_TOKEN_C).await; Mock::given(method("POST")) .and(path("/mcp")) @@ -70,59 +65,64 @@ async fn rmcp_owned_get_and_delete_receive_oauth_recovery() -> anyhow::Result<() Mock::given(method("GET")) .and(path("/mcp")) .and(header("authorization", format!("Bearer {ACCESS_TOKEN_A}"))) - .respond_with(ResponseTemplate::new(401)) + .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}"))) - // A 405 tells RMCP that the optional common SSE stream is unsupported. Reaching this - // response proves that the wrapper retried the RMCP-owned GET with B after refreshing A. - .respond_with({ - let get_retry_marker = get_retry_marker.clone(); - move |_request: &Request| { - std::fs::write(&get_retry_marker, b"observed") - .expect("record retried RMCP-owned GET"); - ResponseTemplate::new(405) - } - }) + .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(401)) - .expect(1) - .mount(&server) - .await; - Mock::given(method("DELETE")) - .and(path("/mcp")) - .and(header("authorization", format!("Bearer {ACCESS_TOKEN_C}"))) - .respond_with({ - let delete_retry_marker = delete_retry_marker.clone(); - move |_request: &Request| { - std::fs::write(&delete_retry_marker, b"observed") - .expect("record retried RMCP-owned DELETE"); - ResponseTemplate::new(204) - } - }) + .respond_with(ResponseTemplate::new(204)) .expect(1) .mount(&server) .await; let status = Command::new(std::env::current_exe()?) .args([ - "oauth_internal_get_delete_child", + "oauth_internal_get_child", "--exact", "--ignored", "--nocapture", ]) .env("CODEX_HOME", codex_home.path()) .env(SERVER_URL_ENV, format!("{}/mcp", server.uri())) - .env(GET_RETRY_MARKER_ENV, &get_retry_marker) - .env(DELETE_RETRY_MARKER_ENV, &delete_retry_marker) + .env(GET_FAILURE_MARKER_ENV, &get_failure_marker) .status() .await?; anyhow::ensure!(status.success(), "OAuth internal child failed: {status}"); @@ -131,13 +131,17 @@ async fn rmcp_owned_get_and_delete_receive_oauth_recovery() -> anyhow::Result<() } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] -#[ignore = "spawned by rmcp_owned_get_and_delete_receive_oauth_recovery"] -async fn oauth_internal_get_delete_child() -> anyhow::Result<()> { +#[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_RETRY_MARKER_ENV).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; - wait_for_marker(DELETE_RETRY_MARKER_ENV).await?; Ok(()) }