diff --git a/codex-rs/rmcp-client/src/http_client_adapter.rs b/codex-rs/rmcp-client/src/http_client_adapter.rs index 19befb62355e..a7d934134812 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; @@ -61,6 +62,10 @@ pub(crate) struct StreamableHttpClientAdapter { 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("MCP OAuth operation failed: {0:#}")] + OAuth(#[source] anyhow::Error), #[error(transparent)] HttpRequest(#[from] ExecServerError), #[error("invalid HTTP header: {0}")] @@ -109,7 +114,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 +167,11 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { StreamableHttpClientAdapterError::SessionExpired404, )); } + if 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 +250,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 +284,11 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { if response.status == StatusCode::METHOD_NOT_ALLOWED.as_u16() { return Ok(()); } + if 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 +331,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 +364,11 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { StreamableHttpClientAdapterError::SessionExpired404, )); } + if 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 +391,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 df787945081d..731f7404f5d5 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 d4d4ad13ec0f..a582e39aea23 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -65,6 +65,7 @@ use codex_keyring_store::KeyringStore; use codex_utils_home_dir::find_codex_home; pub(crate) use self::persistor::OAuthPersistor; +pub(crate) use self::persistor::request_oauth_token_response; use self::refresh_lock::RefreshCredentialLock; pub(crate) use self::resolved_store::ResolvedOAuthCredentialStore; pub(crate) use self::resolved_store::ResolvedOAuthTokens; @@ -73,7 +74,8 @@ pub(crate) use self::resolved_store::resolve_oauth_tokens; const KEYRING_SERVICE: &str = "Codex MCP Credentials"; const MCP_OAUTH_SECRET_PREFIX: &str = "MCP_OAUTH"; -const REFRESH_SKEW_MILLIS: u64 = 30_000; +// Refresh proactively so ordinary requests do not race token expiry. +const REFRESH_SKEW_MILLIS: u64 = 60_000; #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct StoredOAuthTokens { diff --git a/codex-rs/rmcp-client/src/oauth/persistor.rs b/codex-rs/rmcp-client/src/oauth/persistor.rs index 7fc54198e794..700433b2d54a 100644 --- a/codex-rs/rmcp-client/src/oauth/persistor.rs +++ b/codex-rs/rmcp-client/src/oauth/persistor.rs @@ -71,7 +71,6 @@ impl OAuthPersistor { .await } - #[expect(dead_code, reason = "wired by part 4 of this stack")] pub(crate) async fn refresh_after_unauthorized( &self, rejected_access_token: AccessToken, 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..fbdda2641647 --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth_transport.rs @@ -0,0 +1,217 @@ +//! 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. +//! +//! 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. + +use std::collections::HashMap; +use std::sync::Arc; + +use reqwest::header::HeaderName; +use reqwest::header::HeaderValue; +use rmcp::model::ClientJsonRpcMessage; +use rmcp::model::JsonRpcMessage; +use rmcp::transport::auth::AuthClient; +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>; + +#[derive(Clone)] +pub(crate) struct OAuthTransportClient { + auth_client: AuthClient, + persistor: OAuthPersistor, +} + +impl OAuthTransportClient { + pub(crate) fn new( + auth_client: AuthClient, + persistor: OAuthPersistor, + ) -> Self { + Self { + auth_client, + persistor, + } + } + + pub(crate) fn persistor(&self) -> OAuthPersistor { + self.persistor.clone() + } + + 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) + } + + 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 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(_) + ); + 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(), + ) + .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? + { + self.auth_client + .post_message(uri, message, session_id, auth_token, custom_headers) + .await + } else { + result + } + } + + async fn delete_session( + &self, + uri: Arc, + session_id: Arc, + 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(), + ) + .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? + { + self.auth_client + .delete_session(uri, session_id, auth_token, custom_headers) + .await + } else { + 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>, + > { + 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(), + ) + .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? + { + self.auth_client + .get_stream(uri, session_id, last_event_id, auth_token, custom_headers) + .await + } else { + result + } + } +} + +fn rejected_access_token( + error: &StreamableHttpError, +) -> Option { + match error { + StreamableHttpError::Client(StreamableHttpClientAdapterError::AccessTokenRejected { + rejected_access_token, + }) => Some(rejected_access_token.clone()), + _ => None, + } +} + +fn oauth_transport_error( + error: anyhow::Error, +) -> StreamableHttpError { + StreamableHttpError::Client(StreamableHttpClientAdapterError::OAuth(error)) +} diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 55feb4a8ffd6..0c8c61240e65 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; @@ -71,8 +72,10 @@ use crate::oauth::ResolvedOAuthCredentialStore; use crate::oauth::ResolvedOAuthTokens; use crate::oauth::StoredOAuthTokens; use crate::oauth::load_oauth_tokens_from_store; +use crate::oauth::request_oauth_token_response; use crate::oauth::resolve_oauth_tokens; use crate::oauth_http_client::OAuthHttpClientAdapter; +use crate::oauth_transport::OAuthTransportClient; use crate::stdio_server_launcher::StdioServerCommand; use crate::stdio_server_launcher::StdioServerLauncher; use crate::stdio_server_launcher::StdioServerProcessHandle; @@ -99,7 +102,7 @@ enum PendingTransport { transport: StreamableHttpClientTransport, }, StreamableHttpWithOAuth { - transport: StreamableHttpClientTransport>, + transport: StreamableHttpClientTransport, oauth_persistor: OAuthPersistor, }, } @@ -133,6 +136,7 @@ enum TransportRecipe { store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, resolved_store: Arc>, + oauth_client: Arc>, http_client: Arc, auth_provider: Option, }, @@ -410,6 +414,7 @@ impl RmcpClient { store_mode, keyring_backend_kind, resolved_store: Arc::new(OnceLock::new()), + oauth_client: Arc::new(OnceLock::new()), http_client, auth_provider, }; @@ -454,7 +459,7 @@ impl RmcpClient { let mut initialize_deadline = timeout.map(|duration| Instant::now() + duration); let (service, oauth_persistor) = self - .connect_pending_transport_with_initialize_retries( + .connect_pending_transport_with_oauth_recovery( pending_transport, client_service.clone(), timeout, @@ -704,9 +709,20 @@ impl RmcpClient { } async fn service(&self) -> Result>> { + self.service_and_oauth_persistor() + .await + .map(|(service, _oauth_persistor)| service) + } + + async fn service_and_oauth_persistor( + &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")), } @@ -767,6 +783,7 @@ impl RmcpClient { store_mode, keyring_backend_kind, resolved_store, + oauth_client, http_client, auth_provider, } => { @@ -779,6 +796,23 @@ impl RmcpClient { auth_provider.clone() }; + // Reuse the OAuth manager and persistor across initialize retries and session + // reconstruction. Besides pinning Auto's resolved source for this client + // lifecycle, this keeps a successfully refreshed in-memory token authoritative if + // writing it to the resolved store fails. Rebuilding from durable state in that + // condition could reinstall the already-consumed refresh token. + if let Some(oauth_client) = oauth_client.get() { + let runtime = oauth_client.persistor(); + let transport = StreamableHttpClientTransport::with_client( + oauth_client.clone(), + StreamableHttpClientTransportConfig::with_uri(url.clone()), + ); + return Ok(PendingTransport::StreamableHttpWithOAuth { + transport, + oauth_persistor: runtime, + }); + } + let resolved_oauth_tokens = if bearer_token.is_none() && auth_provider.is_none() && !default_headers.contains_key(AUTHORIZATION) @@ -823,7 +857,7 @@ impl RmcpClient { store: credential_store, }) = resolved_oauth_tokens { - match create_oauth_transport_and_runtime( + match create_oauth_transport_client( server_name, url, initial_tokens.clone(), @@ -833,7 +867,19 @@ impl RmcpClient { ) .await { - Ok((transport, oauth_persistor)) => { + Ok(resolved_oauth_client) => { + oauth_client + .set(resolved_oauth_client.clone()) + .map_err(|_| { + anyhow!( + "OAuth client 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()), + ); Ok(PendingTransport::StreamableHttpWithOAuth { transport, oauth_persistor, @@ -961,38 +1007,78 @@ impl RmcpClient { F: Fn(Arc>) -> Fut, Fut: std::future::Future>, { - let service = self.service().await?; - match Self::run_service_operation_with_transient_retries( + let deadline = timeout.map(|duration| Instant::now() + duration); + // Keep the OAuth persistor 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 (service, oauth_persistor) = self.service_and_oauth_persistor().await?; + let mut result = Self::run_service_operation_with_transient_retries( Arc::clone(&service), label, timeout, + deadline, self.elicitation_pause_state.clone(), &operation, ) - .await + .await; + + if let Some(rejected_access_token) = result + .as_ref() + .err() + .and_then(Self::rejected_access_token_from_operation_error) + && let Some(oauth_persistor) = oauth_persistor { - Ok(result) => Ok(result), - Err(error) if Self::is_session_expired_404(&error) => { - self.reinitialize_after_session_expiry(&service).await?; - let recovered_service = self.service().await?; - Self::run_service_operation_with_transient_retries( - recovered_service, - label, - timeout, - self.elicitation_pause_state.clone(), - &operation, - ) - .await - .map_err(Into::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 therefore replay a timed-out tool call. The refresh transaction itself is + // independently owned and completes to its bounded provider timeout if this caller is + // canceled. + remaining_operation_timeout(label, timeout, deadline)?; + let refresh_result = oauth_persistor + .refresh_after_unauthorized(rejected_access_token) + .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); } - Err(error) => Err(error.into()), + result = Self::run_service_operation_with_transient_retries( + Arc::clone(&service), + label, + timeout, + deadline, + self.elicitation_pause_state.clone(), + &operation, + ) + .await; } + + if result.as_ref().is_err_and(Self::is_session_expired_404) { + // Session recovery remains one-shot and runs after the optional OAuth retry, so a 401 + // followed by the old session's 404 still reconstructs the transport before retrying. + self.reinitialize_after_session_expiry(&service).await?; + let recovered_service = self.service().await?; + result = Self::run_service_operation_with_transient_retries( + recovered_service, + label, + timeout, + deadline, + self.elicitation_pause_state.clone(), + &operation, + ) + .await; + } + + result.map_err(Into::into) } async fn run_service_operation_with_transient_retries( service: Arc>, label: &str, timeout: Option, + retry_deadline: Option, pause_state: ElicitationPauseState, operation: &F, ) -> std::result::Result @@ -1000,7 +1086,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() @@ -1106,6 +1191,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>, @@ -1143,7 +1256,7 @@ impl RmcpClient { .timeout .map(|duration| Instant::now() + duration); let (service, oauth_persistor) = self - .connect_pending_transport_with_initialize_retries( + .connect_pending_transport_with_oauth_recovery( pending_transport, initialize_context.client_service, initialize_context.timeout, @@ -1166,17 +1279,14 @@ impl RmcpClient { } } -async fn create_oauth_transport_and_runtime( +async fn create_oauth_transport_client( 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(), @@ -1187,7 +1297,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?; @@ -1205,11 +1315,6 @@ async fn create_oauth_transport_and_runtime( ); 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( server_name.to_string(), url.to_string(), @@ -1218,7 +1323,7 @@ async fn create_oauth_transport_and_runtime( Some(initial_tokens), ); - Ok((transport, runtime)) + Ok(OAuthTransportClient::new(auth_client, runtime)) } #[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 2794489b2a76..540209e55de0 100644 --- a/codex-rs/rmcp-client/src/streamable_http_retry.rs +++ b/codex-rs/rmcp-client/src/streamable_http_retry.rs @@ -5,6 +5,7 @@ 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; @@ -22,8 +23,13 @@ use super::RmcpClient; 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, +} + 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, @@ -32,6 +38,75 @@ impl RmcpClient { ) -> Result<( Arc>, Option, + )> { + let mut attempt_context = InitializeAttemptContext::default(); + match self + .connect_pending_transport_with_initialize_retries( + initial_transport, + client_service.clone(), + timeout, + 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_persistor) = attempt_context.oauth_persistor 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. + 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(); + self.connect_pending_transport_with_initialize_retries( + transport, + client_service, + timeout, + initialize_deadline, + &mut retry_context, + ) + .await + } + } + } + + 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, )> { let should_retry = match &initial_transport { PendingTransport::InProcess { .. } | PendingTransport::Stdio { .. } => false, @@ -62,6 +137,17 @@ 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()), + PendingTransport::InProcess { .. } + | PendingTransport::Stdio { .. } + | PendingTransport::StreamableHttp { .. } => None, + }; match Self::connect_pending_transport( transport, client_service.clone(), @@ -108,6 +194,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 } @@ -158,6 +271,10 @@ impl RmcpClient { | StreamableHttpError::ServerDoesNotSupportSse | StreamableHttpError::Deserialize(_) | StreamableHttpError::Client(StreamableHttpClientAdapterError::SessionExpired404) + | 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_startup.rs b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs index 866daa67a4c0..ab193e7358d2 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs @@ -182,7 +182,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,