github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-21 14:23:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

Restack Codex-owned MCP OAuth recovery

Browse Source
This commit is contained in:
Steven Lee
2026-07-10 01:48:51 +00:00
parent f5a04628b6
commit 2697a2933f
14 changed files with 1739 additions and 215 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/rmcp-client/src/http_client_adapter.rs
+48 -3
View File
@@ -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,17 @@ pub(crate) struct StreamableHttpClientAdapter {
http_client: Arc<dyn HttpClient>,
default_headers: HeaderMap,
auth_provider: Option<SharedAuthProvider>,
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("MCP OAuth operation failed: {0:#}")]
OAuth(#[source] anyhow::Error),
#[error(transparent)]
HttpRequest(#[from] ExecServerError),
#[error("invalid HTTP header: {0}")]
@@ -77,8 +83,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 +122,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 +175,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 +259,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 +293,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 +341,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 +374,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 +402,20 @@ impl StreamableHttpClient for StreamableHttpClientAdapter {
}
}
fn access_token_rejected(
auth_token: Option<&str>,
) -> Option<StreamableHttpError<StreamableHttpClientAdapterError>> {
// 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 {
codex-rs/rmcp-client/src/lib.rs
+1
View File
@@ -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;
codex-rs/rmcp-client/src/oauth.rs
+1 -65
View File
@@ -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";
codex-rs/rmcp-client/src/oauth/refresh_transaction.rs
+116 -37
View File
@@ -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,62 @@ 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
}
pub(crate) async fn refresh_after_unauthorized(
&self,
rejected_access_token: AccessToken,
) -> Result<()> {
self.refresh_in(
DefaultKeyringStore,
RefreshReason::Unauthorized {
rejected_access_token,
},
REFRESH_REQUEST_TIMEOUT,
)
.await
}
/// Injects the credential backend and provider timeout for deterministic failure-path tests.
pub(super) async fn refresh_if_needed_in<K: KeyringStore + Clone + 'static>(
pub(super) async fn refresh_in<K: KeyringStore + Clone + 'static>(
&self,
keyring_store: &K,
keyring_store: K,
reason: RefreshReason,
refresh_request_timeout: Duration,
) -> Result<()> {
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(());
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 +112,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<K: KeyringStore + Clone + 'static>(
&self,
keyring_store: &K,
reason: RefreshReason,
refresh_request_timeout: Duration,
) -> Result<()> {
debug!("waiting for the MCP OAuth credential transaction lock");
@@ -132,11 +154,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 +194,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 +237,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 +278,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 +300,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 +364,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,
codex-rs/rmcp-client/src/oauth/tests/persistor_tests.rs
+73 -7
View File
@@ -38,6 +38,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 +130,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 +145,70 @@ 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"
@@ -168,7 +229,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 +375,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
codex-rs/rmcp-client/src/oauth_transport.rs
+254
View File
@@ -0,0 +1,254 @@
//! 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::auth::AuthError;
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<T> =
std::result::Result<T, StreamableHttpError<StreamableHttpClientAdapterError>>;
#[derive(Clone)]
pub(crate) struct OAuthTransportClient {
auth_client: AuthClient<StreamableHttpClientAdapter>,
persistor: OAuthPersistor,
}
impl OAuthTransportClient {
pub(crate) fn new(
auth_client: AuthClient<StreamableHttpClientAdapter>,
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<oauth2::AccessToken>,
) -> TransportResult<bool> {
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<str>,
message: ClientJsonRpcMessage,
session_id: Option<Arc<str>>,
auth_token: Option<String>,
custom_headers: HashMap<HeaderName, HeaderValue>,
) -> TransportResult<StreamableHttpPostResponse> {
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?
{
authorization_required_after_retry(
self.auth_client
.post_message(uri, message, session_id, auth_token, custom_headers)
.await,
)
} else {
result
}
}
async fn delete_session(
&self,
uri: Arc<str>,
session_id: Arc<str>,
auth_token: Option<String>,
custom_headers: HashMap<HeaderName, HeaderValue>,
) -> 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?
{
authorization_required_after_retry(
self.auth_client
.delete_session(uri, session_id, auth_token, custom_headers)
.await,
)
} else {
result
}
}
async fn get_stream(
&self,
uri: Arc<str>,
session_id: Arc<str>,
last_event_id: Option<String>,
auth_token: Option<String>,
custom_headers: HashMap<HeaderName, HeaderValue>,
) -> TransportResult<
futures::stream::BoxStream<'static, Result<sse_stream::Sse, sse_stream::Error>>,
> {
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?
{
authorization_required_after_retry(
self.auth_client
.get_stream(uri, session_id, last_event_id, auth_token, custom_headers)
.await,
)
} else {
result
}
}
}
fn authorization_required_after_retry<T>(result: TransportResult<T>) -> TransportResult<T> {
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,
}
}
fn rejected_access_token(
error: &StreamableHttpError<StreamableHttpClientAdapterError>,
) -> Option<oauth2::AccessToken> {
match error {
StreamableHttpError::Client(StreamableHttpClientAdapterError::AccessTokenRejected {
rejected_access_token,
}) => Some(rejected_access_token.clone()),
_ => None,
}
}
fn oauth_transport_error(
error: anyhow::Error,
) -> StreamableHttpError<StreamableHttpClientAdapterError> {
if let Some(auth_error) =
error
.chain()
.find_map(|source| match source.downcast_ref::<AuthError>() {
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;
codex-rs/rmcp-client/src/oauth_transport_tests.rs
+220
View File
@@ -0,0 +1,220 @@
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::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 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<()> {
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"))
.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)
.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;
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 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 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 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 response_message: ClientJsonRpcMessage = serde_json::from_value(json!({
"jsonrpc": "2.0",
"id": "server-request-1",
"result": {
"action": "accept",
"content": { "confirmed": true }
}
}))?;
let response = client
.post_message(
Arc::from(server_url),
response_message,
Some(Arc::from("response-session")),
/*auth_token*/ None,
HashMap::new(),
)
.await?;
assert!(matches!(response, StreamableHttpPostResponse::Accepted));
Ok(())
}
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,
}
}
codex-rs/rmcp-client/src/rmcp_client.rs
+161 -91
View File
@@ -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,10 @@ 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::stdio_server_launcher::StdioServerCommand;
use crate::stdio_server_launcher::StdioServerLauncher;
use crate::stdio_server_launcher::StdioServerProcessHandle;
@@ -97,7 +100,7 @@ enum PendingTransport {
transport: StreamableHttpClientTransport<StreamableHttpClientAdapter>,
},
StreamableHttpWithOAuth {
transport: StreamableHttpClientTransport<AuthClient<StreamableHttpClientAdapter>>,
transport: StreamableHttpClientTransport<OAuthTransportClient>,
oauth_persistor: OAuthPersistor,
},
}
@@ -131,6 +134,7 @@ enum TransportRecipe {
store_mode: OAuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
pinned_credential_store: Arc<OnceLock<ResolvedOAuthCredentialStore>>,
oauth_client: Arc<OnceLock<OAuthTransportClient>>,
http_client: Arc<dyn HttpClient>,
auth_provider: Option<SharedAuthProvider>,
},
@@ -408,6 +412,7 @@ impl RmcpClient {
store_mode,
keyring_backend_kind,
pinned_credential_store: Arc::new(OnceLock::new()),
oauth_client: Arc::new(OnceLock::new()),
http_client,
auth_provider,
};
@@ -451,7 +456,7 @@ impl RmcpClient {
};
let (service, oauth_persistor) = self
.connect_pending_transport_with_initialize_retries(
.connect_pending_transport_with_oauth_recovery(
pending_transport,
client_service.clone(),
timeout,
@@ -483,12 +488,6 @@ impl RmcpClient {
};
}
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 +503,6 @@ impl RmcpClient {
async move { service.list_tools(params).await }.boxed()
})
.await?;
self.persist_oauth_tokens().await;
Ok(result)
}
@@ -539,7 +537,6 @@ impl RmcpClient {
})
})
.collect::<Result<Vec<_>>>()?;
self.persist_oauth_tokens().await;
Ok(ListToolsWithConnectorIdResult {
next_cursor: result.next_cursor,
tools,
@@ -566,7 +563,6 @@ impl RmcpClient {
async move { service.list_resources(params).await }.boxed()
})
.await?;
self.persist_oauth_tokens().await;
Ok(result)
}
@@ -582,7 +578,6 @@ impl RmcpClient {
async move { service.list_resource_templates(params).await }.boxed()
})
.await?;
self.persist_oauth_tokens().await;
Ok(result)
}
@@ -598,7 +593,6 @@ impl RmcpClient {
async move { service.read_resource(params).await }.boxed()
})
.await?;
self.persist_oauth_tokens().await;
Ok(result)
}
@@ -656,7 +650,6 @@ impl RmcpClient {
.boxed()
})
.await?;
self.persist_oauth_tokens().await;
Ok(result)
}
@@ -686,7 +679,6 @@ impl RmcpClient {
},
)
.await?;
self.persist_oauth_tokens().await;
Ok(())
}
@@ -709,14 +701,18 @@ impl RmcpClient {
.boxed()
})
.await?;
self.persist_oauth_tokens().await;
Ok(response)
}
async fn service(&self) -> Result<Arc<RunningService<RoleClient, ElicitationClientService>>> {
async fn service_and_oauth_persistor(
&self,
) -> Result<(
Arc<RunningService<RoleClient, ElicitationClientService>>,
Option<OAuthPersistor>,
)> {
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")),
}
@@ -749,17 +745,6 @@ 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}");
}
}
/// 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?;
@@ -790,6 +775,7 @@ impl RmcpClient {
store_mode,
keyring_backend_kind,
pinned_credential_store,
oauth_client,
http_client,
auth_provider,
} => {
@@ -802,6 +788,21 @@ 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(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)
@@ -822,8 +823,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 +850,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(),
@@ -857,7 +860,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,
@@ -956,19 +971,7 @@ 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))
}
@@ -983,31 +986,83 @@ impl RmcpClient {
F: Fn(Arc<RunningService<RoleClient, ElicitationClientService>>) -> Fut,
Fut: std::future::Future<Output = std::result::Result<T, rmcp::service::ServiceError>>,
{
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) => {
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)
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 (mut service, mut oauth_persistor) = self.service_and_oauth_persistor().await?;
let mut oauth_recovered = 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)
{
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 {
return result.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 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_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;
}
Err(error) => Err(error.into()),
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?;
(service, oauth_persistor) = self.service_and_oauth_persistor().await?;
session_recovered = true;
continue;
}
return result.map_err(Into::into);
}
}
@@ -1015,6 +1070,7 @@ impl RmcpClient {
service: Arc<RunningService<RoleClient, ElicitationClientService>>,
label: &str,
timeout: Option<Duration>,
retry_deadline: Option<Instant>,
pause_state: ElicitationPauseState,
operation: &F,
) -> std::result::Result<T, ClientOperationError>
@@ -1022,7 +1078,6 @@ impl RmcpClient {
F: Fn(Arc<RunningService<RoleClient, ElicitationClientService>>) -> Fut,
Fut: std::future::Future<Output = std::result::Result<T, rmcp::service::ServiceError>>,
{
let retry_deadline = timeout.map(|duration| Instant::now() + duration);
for (attempt, retry_delay_ms) in STREAMABLE_HTTP_RETRY_DELAYS_MS
.iter()
.copied()
@@ -1128,6 +1183,34 @@ impl RmcpClient {
})
}
fn rejected_access_token_from_operation_error(
error: &ClientOperationError,
) -> Option<AccessToken> {
let ClientOperationError::Service(rmcp::service::ServiceError::TransportSend(error)) =
error
else {
return None;
};
error
.error
.downcast_ref::<StreamableHttpError<StreamableHttpClientAdapterError>>()
.and_then(Self::rejected_access_token)
}
pub(super) fn rejected_access_token(
error: &StreamableHttpError<StreamableHttpClientAdapterError>,
) -> Option<AccessToken> {
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<RunningService<RoleClient, ElicitationClientService>>,
@@ -1162,7 +1245,7 @@ impl RmcpClient {
.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(
.connect_pending_transport_with_oauth_recovery(
pending_transport,
initialize_context.client_service,
initialize_context.timeout,
@@ -1180,27 +1263,18 @@ impl RmcpClient {
};
}
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_transport_client(
server_name: &str,
url: &str,
initial_tokens: StoredOAuthTokens,
credential_store: ResolvedOAuthCredentialStore,
default_headers: HeaderMap,
http_client: Arc<dyn HttpClient>,
) -> Result<(
StreamableHttpClientTransport<AuthClient<StreamableHttpClientAdapter>>,
OAuthPersistor,
)> {
) -> Result<OAuthTransportClient> {
let oauth_http_client = Arc::new(OAuthHttpClientAdapter::new(
http_client.clone(),
default_headers.clone(),
@@ -1211,7 +1285,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,16 +1298,12 @@ 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(
server_name.to_string(),
url.to_string(),
@@ -1242,7 +1312,7 @@ async fn create_oauth_transport_and_runtime(
Some(initial_tokens),
);
Ok((transport, runtime))
Ok(OAuthTransportClient::new(auth_client, runtime))
}
#[cfg(test)]
codex-rs/rmcp-client/src/startup_error.rs
+1 -1
View File
@@ -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(_))
})
}
codex-rs/rmcp-client/src/streamable_http_retry.rs
+139 -9
View File
@@ -5,9 +5,11 @@ 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;
@@ -22,8 +24,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<OAuthPersistor>,
}
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,
@@ -31,13 +38,95 @@ impl RmcpClient {
) -> Result<(
Arc<RunningService<RoleClient, ElicitationClientService>>,
Option<OAuthPersistor>,
)> {
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_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();
let result = self
.connect_pending_transport_with_initialize_retries(
transport,
client_service,
timeout,
&mut initialize_deadline,
&mut retry_context,
)
.await;
if 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.
return Err(AuthError::AuthorizationRequired.into());
}
result
}
}
}
async fn connect_pending_transport_with_initialize_retries(
&self,
initial_transport: PendingTransport,
client_service: ElicitationClientService,
timeout: Option<Duration>,
initialize_deadline: &mut Option<Instant>,
attempt_context: &mut InitializeAttemptContext,
) -> Result<(
Arc<RunningService<RoleClient, ElicitationClientService>>,
Option<OAuthPersistor>,
)> {
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 +139,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,
@@ -66,16 +155,26 @@ impl RmcpClient {
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.
// 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() {
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_persistor = match &transport {
PendingTransport::StreamableHttpWithOAuth {
oauth_persistor, ..
} => Some(oauth_persistor.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 +195,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 +220,33 @@ impl RmcpClient {
})
}
fn rejected_access_token_from_initialize_error(error: &anyhow::Error) -> Option<AccessToken> {
error.chain().find_map(|source| {
source
.downcast_ref::<HandshakeError>()
.and_then(|error| {
Self::rejected_access_token_from_client_initialize_error(&error.source)
})
.or_else(|| {
source
.downcast_ref::<rmcp::service::ClientInitializeError>()
.and_then(Self::rejected_access_token_from_client_initialize_error)
})
})
}
fn rejected_access_token_from_client_initialize_error(
error: &rmcp::service::ClientInitializeError,
) -> Option<AccessToken> {
match error {
rmcp::service::ClientInitializeError::TransportError { error, .. } => error
.error
.downcast_ref::<StreamableHttpError<StreamableHttpClientAdapterError>>()
.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 +297,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,
}
codex-rs/rmcp-client/tests/streamable_http_oauth_internal.rs
+250
View File
@@ -0,0 +1,250 @@
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_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 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<()> {
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 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"))
.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(ResponseTemplate::new(401))
.expect(1)
.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)
}
})
.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)
}
})
.expect(1)
.mount(&server)
.await;
let status = Command::new(std::env::current_exe()?)
.args([
"oauth_internal_get_delete_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)
.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_and_delete_receive_oauth_recovery"]
async fn oauth_internal_get_delete_child() -> anyhow::Result<()> {
let client = create_oauth_client().await?;
initialize_client(&client).await?;
wait_for_marker(GET_RETRY_MARKER_ENV).await?;
client.shutdown().await;
wait_for_marker(DELETE_RETRY_MARKER_ENV).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<RmcpClient> {
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"
}
}
}))
}
codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs
+433 -1
View File
@@ -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,289 @@ 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:?}")),
}
})
// The later operation starts immediately, so it can send A once before its own 401 joins
// the in-flight recovery. Exactly four requests use A: initialize, initialized, and one
// rejected tools/list per operation. The single provider request below proves both 401s
// converge on the same refresh transaction.
.expect(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 +527,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 +669,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<RmcpClient> {
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",
},
},
}))
}
codex-rs/rmcp-client/tests/streamable_http_recovery.rs
+40 -1
View File
@@ -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?;
codex-rs/rmcp-client/tests/streamable_http_test_support.rs
+2
View File
@@ -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?;