openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-22 17:58:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Enable coordinated MCP OAuth refresh (#42413)

Browse Source 
## What changed

- Honor `McpOAuthRefreshMode::Coordinated` for streamable HTTP MCP connections by letting RMCP refresh and persist credentials through the pinned credential store.
- Reread credentials while holding the refresh lock, persist refreshed tokens before retrying requests, and serialize login and logout writes with in-flight refreshes.
- Keep OAuth preparation outside MCP operation and handshake deadlines, and allow refresh persistence to finish if the initiating caller is cancelled.

## Testing

- Cover coordinated startup refresh, concurrent refreshes, `401` retries, credential mutation races, provider timeouts, and caller cancellation.

GitOrigin-RevId: 100ae788048087be23e26bb43f868b1fcebb007a
This commit is contained in:
Matthew Zeng
2026-09-03 02:00:12 +00:00
committed by copyberry
parent fdf23b4097
commit 88912c04cd
17 changed files with 683 additions and 128 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/cli/src/mcp_cmd.rs
+3 -1
View File
@@ -660,7 +660,9 @@ async fn run_logout(config: &Config, logout_args: LogoutArgs) -> Result<()> {
&url,
config.mcp_oauth_credentials_store_mode,
config.auth_keyring_backend_kind(),
) {
)
.await
{
Ok(true) => println!("Removed OAuth credentials for '{name}'."),
Ok(false) => println!("No OAuth credentials stored for '{name}'."),
Err(err) => return Err(anyhow!("failed to delete OAuth credentials: {err}")),
codex-rs/codex-mcp/src/rmcp_client.rs
+1 -5
View File
@@ -1088,11 +1088,6 @@ pub(crate) async fn make_rmcp_client(
runtime_auth_provider: Option<SharedAuthProvider>,
protocol_mode: McpProtocolMode,
) -> Result<RmcpClient, StartupOutcomeError> {
if oauth_refresh_mode == McpOAuthRefreshMode::Coordinated {
warn!(
"MCP OAuth refresh coordination is not available in this build; using legacy refresh"
);
}
let config = server.config().clone();
if matches!(config.auth, McpServerAuth::ChatGpt)
&& !config.is_local_environment()
@@ -1215,6 +1210,7 @@ pub(crate) async fn make_rmcp_client(
runtime_auth_provider,
protocol_mode,
redirect_mode,
oauth_refresh_mode,
)
.await
.map_err(StartupOutcomeError::from)
codex-rs/core/tests/suite/mcp_oauth_refresh_tests.rs
+2 -1
View File
@@ -66,7 +66,8 @@ async fn oauth_mode_refresh_replaces_the_live_connection(
"synthetic-access-token",
"synthetic-refresh-token",
OAuthCredentialExpiry::Valid,
)?;
)
.await?;
let fixture = test_codex()
.with_home(home)
.with_config(move |config| {
codex-rs/core/tests/suite/rmcp_client.rs
+21 -10
View File
@@ -3774,7 +3774,8 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> {
"expired-access-token",
refresh_token,
OAuthCredentialExpiry::Expired,
)?;
)
.await?;
let discovered_credential_name =
credential_config.oauth_credential_name(discovered_server_name);
write_fallback_oauth_tokens(
@@ -3784,7 +3785,8 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> {
expected_token,
refresh_token,
OAuthCredentialExpiry::Valid,
)?;
)
.await?;
// Phase 4: configure Codex with the OAuth-backed Streamable HTTP MCP
// server and build the fixture in the active local or remote-aware mode.
@@ -3792,6 +3794,10 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> {
.with_model_info_override("gpt-5.4", |model| model.supports_search_tool = false)
.with_home(temp_home.clone())
.with_config(move |config| {
config
.features
.enable(Feature::McpOAuthRefreshCoordination)
.expect("test config should allow coordinated MCP OAuth refresh");
config.mcp_oauth_credentials_store_mode = OAuthCredentialsStoreMode::Auto;
insert_mcp_server(
config,
@@ -3957,12 +3963,15 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> {
.await
.context("the newly discovered OAuth server did not recover after its store was unlocked")??;
assert!(codex_rmcp_client::delete_oauth_tokens(
discovered_credential_name.as_ref(),
http_server.url(),
OAuthCredentialsStoreMode::File,
codex_config::types::AuthKeyringBackendKind::default(),
)?);
assert!(
codex_rmcp_client::delete_oauth_tokens(
discovered_credential_name.as_ref(),
http_server.url(),
OAuthCredentialsStoreMode::File,
codex_config::types::AuthKeyringBackendKind::default(),
)
.await?
);
fixture.codex.refresh_runtime_config(refreshed_config).await;
let logged_out_startup = tokio::time::timeout(
Duration::from_secs(5),
@@ -3990,7 +3999,8 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> {
expected_token,
refresh_token,
OAuthCredentialExpiry::Valid,
)?;
)
.await?;
// Phase 6: submit the user turn that should invoke the OAuth-backed tool.
fixture
@@ -4431,7 +4441,7 @@ enum OAuthCredentialExpiry {
Expired,
}
fn write_fallback_oauth_tokens(
async fn write_fallback_oauth_tokens(
server_name: &str,
server_url: &str,
client_id: &str,
@@ -4468,6 +4478,7 @@ fn write_fallback_oauth_tokens(
OAuthCredentialsStoreMode::File,
codex_config::types::AuthKeyringBackendKind::default(),
)
.await
}
struct EnvVarGuard {
codex-rs/rmcp-client/src/oauth.rs
+6 -4
View File
@@ -16,13 +16,13 @@
//!
//! If the keyring is not available or fails, we fall back to CODEX_HOME/.credentials.json which is consistent with other coding CLI agents.
#[cfg(test)]
mod credential_store;
mod ema_identity;
mod issuer_binding;
mod refresh_lock;
mod refresh_transaction;
mod resolved_store;
mod runtime;
mod store_lock;
#[cfg(test)]
@@ -74,7 +74,6 @@ use tokio::sync::Mutex;
use codex_utils_home_dir::find_codex_home;
#[cfg(test)]
pub(crate) use self::credential_store::OAuthCredentialStore;
pub(crate) use self::ema_identity::stored_oidc_identity;
pub(crate) use self::issuer_binding::validate_authorization_server_endpoints;
@@ -85,6 +84,7 @@ 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;
use self::resolved_store::try_resolve_oauth_tokens_from_store_policy;
pub(crate) use self::runtime::OAuthRuntime;
const KEYRING_SERVICE: &str = "Codex MCP Credentials";
const MCP_OAUTH_SECRET_PREFIX: &str = "MCP_OAUTH";
@@ -440,12 +440,13 @@ enum OAuthKeyringLoadError {
Backend(#[from] anyhow::Error),
}
pub fn save_oauth_tokens(
pub async fn save_oauth_tokens(
server_name: &str,
tokens: &StoredOAuthTokens,
store_mode: OAuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
) -> Result<()> {
let _lock = RefreshCredentialLock::acquire_for_server(server_name, &tokens.url).await?;
let keyring_store = DefaultKeyringStore;
match store_mode {
OAuthCredentialsStoreMode::Auto => save_oauth_tokens_with_keyring_with_fallback_to_file(
@@ -585,12 +586,13 @@ fn save_oauth_tokens_with_keyring_with_fallback_to_file<K: KeyringStore + Clone
}
}
pub fn delete_oauth_tokens(
pub async fn delete_oauth_tokens(
server_name: &str,
url: &str,
store_mode: OAuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
) -> Result<bool> {
let _lock = RefreshCredentialLock::acquire_for_server(server_name, url).await?;
let keyring_store = DefaultKeyringStore;
delete_oauth_tokens_from_keyring_and_file(
&keyring_store,
codex-rs/rmcp-client/src/oauth/credential_store.rs
+87 -24
View File
@@ -4,8 +4,10 @@
//! against an in-memory store. Initialization may save tokenless client credentials, which
//! this refresh adapter does not support.
//!
//! RMCP must hold the credential lock across its authoritative reread, token exchange, and
//! save. Saves and clears require an active guard; no store operation reacquires the lock.
//! Ordinary token reads use the cached credentials. Refresh-guard acquisition rereads
//! the pinned store before RMCP exchanges the token and saves the result. Codex
//! preparation rechecks freshness under that guard before asking RMCP to refresh.
//! Saves and clears require an active guard; no store operation reacquires the lock.
//! The runtime snapshot advances only for credentials compatible with this connection,
//! so replacement logins still cause a rebuild.
//! Synchronous store operations run off the Tokio workers so caller deadlines remain pollable.
@@ -22,21 +24,29 @@ use futures::future::BoxFuture;
use oauth2::Scope;
use oauth2::TokenResponse;
use rmcp::transport::auth::AuthError;
use rmcp::transport::auth::AuthorizationManager;
use rmcp::transport::auth::CredentialRefreshGuard;
use rmcp::transport::auth::CredentialStore;
use rmcp::transport::auth::StoredCredentials;
use tokio::sync::Mutex;
use tracing::warn;
use crate::oauth_http_client::PROACTIVE_REFRESH_TIMEOUT;
use super::RefreshCredentialLock;
use super::ResolvedOAuthCredentialStore;
use super::StoredOAuthTokens;
use super::WrappedOAuthTokenResponse;
use super::normalized_oauth_credentials;
use super::refresh_expires_in_from_timestamp;
use super::refresh_transaction::REFRESH_REQUEST_TIMEOUT;
use super::token_needs_refresh;
use super::validate_refresh_token_issuer;
#[derive(Clone)]
pub(crate) struct OAuthCredentialStore<K = DefaultKeyringStore> {
inner: Arc<OAuthCredentialStoreInner<K>>,
held_refresh_guard: Option<Arc<RefreshCredentialLock>>,
}
struct OAuthCredentialStoreInner<K> {
@@ -67,6 +77,39 @@ impl<K: KeyringStore + Clone + 'static> OAuthCredentialStore<K> {
last_credentials: Mutex::new(Some(tokens)),
refresh_guard: Mutex::new(Weak::new()),
}),
held_refresh_guard: None,
}
}
pub(crate) async fn refresh_if_needed(&self, manager: &mut AuthorizationManager) -> Result<()> {
let guard = self.acquire_transaction_guard().await?;
let tokens = self
.inner
.last_credentials
.lock()
.await
.clone()
.ok_or(AuthError::AuthorizationRequired)?;
if !token_needs_refresh(tokens.expires_at) {
return Ok(());
}
let metadata = manager.resolve_metadata().await?.metadata;
validate_refresh_token_issuer(&metadata, &tokens)?;
manager.set_metadata(metadata);
manager.configure_client_id(&tokens.client_id)?;
// Reuse the guard for RMCP's exchange and save after the locked freshness check.
manager.set_credential_store(Self {
inner: Arc::clone(&self.inner),
held_refresh_guard: Some(guard),
});
let result = PROACTIVE_REFRESH_TIMEOUT
.scope(REFRESH_REQUEST_TIMEOUT, manager.refresh_token())
.await;
manager.set_credential_store(self.clone());
match result {
Ok(_) => Ok(()),
Err(AuthError::TokenRefreshRejected(_)) => Err(AuthError::AuthorizationRequired.into()),
Err(error) => Err(error.into()),
}
}
@@ -74,11 +117,27 @@ impl<K: KeyringStore + Clone + 'static> OAuthCredentialStore<K> {
normalized_oauth_credentials(self.inner.last_credentials.lock().await.as_ref())
}
pub(crate) async fn acquire_transaction_guard(&self) -> Result<Arc<RefreshCredentialLock>> {
pub(crate) async fn acquire_transaction_guard(
&self,
) -> Result<Arc<RefreshCredentialLock>, AuthError> {
let guard = Arc::new(
RefreshCredentialLock::acquire_for_server(&self.inner.server_name, &self.inner.url)
.await?,
.await
.map_err(credential_store_error)?,
);
let inner = Arc::clone(&self.inner);
let tokens = tokio::task::spawn_blocking(move || {
inner
.store
.load(&inner.keyring, &inner.server_name, &inner.url)
})
.await
.context("OAuth credential load task failed")
.map_err(credential_store_error)?
.map_err(credential_store_error)?
.ok_or(AuthError::AuthorizationRequired)?;
self.validate_connection(&tokens)?;
*self.inner.last_credentials.lock().await = Some(tokens);
*self.inner.refresh_guard.lock().await = Arc::downgrade(&guard);
Ok(guard)
}
@@ -111,25 +170,13 @@ impl<K: KeyringStore + Clone + 'static> CredentialStore for OAuthCredentialStore
Self: 'async_trait,
{
Box::pin(async move {
let inner = Arc::clone(&self.inner);
let Some(tokens) = tokio::task::spawn_blocking(move || {
inner
.store
.load(&inner.keyring, &inner.server_name, &inner.url)
})
.await
.context("OAuth credential load task failed")
.map_err(credential_store_error)?
.map_err(credential_store_error)?
else {
// Do not acknowledge removal as a live connection update: reconciliation
// must see that this connection still belongs to the previous credentials.
return Ok(None);
};
self.validate_connection(&tokens)?;
let credentials = rmcp_credentials(&tokens);
*self.inner.last_credentials.lock().await = Some(tokens);
Ok(Some(credentials))
Ok(self
.inner
.last_credentials
.lock()
.await
.as_ref()
.map(rmcp_credentials))
})
}
@@ -228,10 +275,26 @@ impl<K: KeyringStore + Clone + 'static> CredentialStore for OAuthCredentialStore
.map_err(credential_store_error)
})
}
fn acquire_refresh_guard<'life0, 'async_trait>(
&'life0 self,
) -> BoxFuture<'async_trait, Result<Option<CredentialRefreshGuard>, AuthError>>
where
'life0: 'async_trait,
Self: 'async_trait,
{
Box::pin(async move {
let guard = match &self.held_refresh_guard {
Some(guard) => Arc::clone(guard),
None => self.acquire_transaction_guard().await?,
};
Ok(Some(CredentialRefreshGuard::new(guard)))
})
}
}
fn credential_store_error(error: anyhow::Error) -> AuthError {
AuthError::InternalError(format!("{error:#}"))
AuthError::CredentialStoreError(format!("{error:#}"))
}
fn rmcp_credentials(tokens: &StoredOAuthTokens) -> StoredCredentials {
codex-rs/rmcp-client/src/oauth/refresh_transaction.rs
+1 -1
View File
@@ -28,7 +28,7 @@ use super::refresh_lock::RefreshCredentialLock;
use super::token_needs_refresh;
use super::validate_refresh_token_issuer;
const REFRESH_REQUEST_TIMEOUT: Duration = Duration::from_secs(45);
pub(super) const REFRESH_REQUEST_TIMEOUT: Duration = Duration::from_secs(45);
impl OAuthPersistor {
pub(crate) async fn refresh_if_needed(&self) -> Result<()> {
codex-rs/rmcp-client/src/oauth/runtime.rs
+64
View File
@@ -0,0 +1,64 @@
//! Keeps one OAuth refresh and persistence owner for the lifetime of a connection.
//!
//! Both modes use cached expiry to prepare credentials before the MCP request budget.
//! Coordinated preparation owns a task so caller cancellation cannot interrupt persistence.
//! It locks the manager before the credential store, matching RMCP's lock order.
//! RMCP commits through the pinned store while Legacy retains Codex's persistor.
use std::sync::Arc;
use anyhow::Context;
use anyhow::Result;
use rmcp::transport::auth::AuthorizationManager;
use tokio::sync::Mutex;
use tracing::warn;
use super::OAuthPersistor;
use super::credential_store::OAuthCredentialStore;
use super::token_needs_refresh;
#[derive(Clone)]
pub(crate) enum OAuthRuntime {
Legacy(OAuthPersistor),
Coordinated {
auth_manager: Arc<Mutex<AuthorizationManager>>,
store: OAuthCredentialStore,
},
}
impl OAuthRuntime {
#[expect(
clippy::await_holding_invalid_type,
reason = "AuthorizationManager access follows RMCP's manager-before-credential lock order"
)]
pub(crate) async fn refresh_if_needed(&self) -> Result<()> {
match self {
Self::Legacy(persistor) => persistor.refresh_if_needed().await,
Self::Coordinated {
auth_manager,
store,
} => {
let expires_at = store
.stored_credentials()
.await
.and_then(|tokens| tokens.expires_at);
if !token_needs_refresh(expires_at) {
return Ok(());
}
let auth_manager = Arc::clone(auth_manager);
let store = store.clone();
tokio::spawn(async move {
let mut manager = auth_manager.lock().await;
let result = store.refresh_if_needed(&mut manager).await;
if result.is_err() {
// Keep the summary in the owned task without logging credential/provider data.
warn!("MCP OAuth preparation failed");
}
result
})
.await
.context("OAuth refresh task failed")?
}
}
}
}
codex-rs/rmcp-client/src/oauth/tests/credential_store_tests.rs
+31 -9
View File
@@ -51,14 +51,22 @@ async fn mutations_require_and_retain_the_transaction_guard() -> Result<()> {
store.save(credentials.clone()).await
};
assert!(
matches!(result, Err(AuthError::InternalError(error)) if error.contains("active refresh guard"))
matches!(result, Err(AuthError::CredentialStoreError(error)) if error.contains("active refresh guard"))
);
assert_eq!(
load_oauth_tokens_from_file(&initial.server_name, &initial.url)?.is_none(),
!clear,
);
// Guard acquisition now rereads storage. Restore the deleted credential first,
// then delete it again under the guard so the queued save must recreate it.
save_oauth_tokens_to_file(&initial)?;
let guard = CredentialStore::acquire_refresh_guard(&store)
.await?
.expect("coordinated refresh guard");
if !clear {
store.clear().await?;
}
let aggregate_lock = OAuthStoreLock::acquire_for_write(OAuthStore::File)?;
let guard = store.acquire_transaction_guard().await?;
let mut mutation = if clear {
store.clear()
} else {
@@ -117,7 +125,7 @@ async fn save_publishes_only_persisted_credentials() -> Result<()> {
.expect("durable keyring credentials");
let mut expected = initial.clone();
if fail_save {
assert!(matches!(result, Err(AuthError::InternalError(_))));
assert!(matches!(result, Err(AuthError::CredentialStoreError(_))));
} else {
result?;
expected
@@ -156,9 +164,14 @@ async fn pinned_read_failure_does_not_adopt_fallback_credentials() -> Result<()>
ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct),
keyring,
);
let cached = store.load().await?.expect("cached credentials");
assert_eq!(
cached.token_response.unwrap().access_token().secret(),
initial.token_response.0.access_token().secret()
);
assert!(matches!(
store.load().await,
Err(AuthError::InternalError(_))
store.acquire_transaction_guard().await,
Err(AuthError::CredentialStoreError(_))
));
assert_eq!(
store.stored_credentials().await,
@@ -191,14 +204,21 @@ async fn replacement_or_removal_does_not_acknowledge_a_new_runtime_snapshot() ->
replacement.issuer = issuer;
save_oauth_tokens_to_file(&replacement)?;
assert!(matches!(
store.load().await,
store.acquire_transaction_guard().await,
Err(AuthError::AuthorizationRequired)
));
assert_eq!(store.stored_credentials().await, original_snapshot);
}
let _guard = store.acquire_transaction_guard().await?;
save_oauth_tokens_to_file(&initial)?;
let guard = store.acquire_transaction_guard().await?;
store.clear().await?;
assert!(store.load().await?.is_none());
assert!(load_oauth_tokens_from_file(&initial.server_name, &initial.url)?.is_none());
drop(guard);
assert!(matches!(
store.acquire_transaction_guard().await,
Err(AuthError::AuthorizationRequired)
));
assert!(store.load().await?.is_some());
assert_eq!(store.stored_credentials().await, original_snapshot);
Ok(())
}
@@ -219,7 +239,7 @@ async fn storage_roundtrip_preserves_absolute_and_unknown_expiry() -> Result<()>
(Some(future_received_at), Some(future_deadline)),
(None, None),
] {
let _guard = store.acquire_transaction_guard().await?;
let guard = store.acquire_transaction_guard().await?;
let mut credentials = store.load().await?.unwrap();
credentials.token_received_at = received_at;
credentials
@@ -236,6 +256,8 @@ async fn storage_roundtrip_preserves_absolute_and_unknown_expiry() -> Result<()>
.0
.set_expires_in(expected_deadline.map(|_| Duration::ZERO).as_ref());
assert_tokens_match_without_expiry(&durable, &expected);
drop(guard);
let _guard = store.acquire_transaction_guard().await?;
let reloaded = store.load().await?.unwrap();
let expires_in = reloaded.token_response.unwrap().expires_in();
if let Some(expected_deadline) = expected_deadline {
codex-rs/rmcp-client/src/oauth/tests/persistor_tests.rs
+302 -17
View File
@@ -9,14 +9,18 @@ use codex_config::types::OAuthCredentialsStoreMode;
use codex_exec_server::RouteAwareHttpClient;
use codex_http_client::HttpClientFactory;
use codex_http_client::OutboundProxyPolicy;
use codex_keyring_store::DefaultKeyringStore;
use http::HeaderMap;
use keyring::Error as KeyringError;
use oauth2::AccessToken;
use oauth2::RefreshToken;
use oauth2::TokenResponse;
use pretty_assertions::assert_eq;
use rmcp::transport::auth::AuthClient;
use rmcp::transport::auth::AuthError;
use rmcp::transport::auth::AuthorizationManager;
use rmcp::transport::auth::OAuthState;
use rmcp::transport::streamable_http_client::StreamableHttpClient;
use tokio::sync::Mutex as TokioMutex;
use tracing::Event;
use tracing::Id;
@@ -29,6 +33,7 @@ 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;
@@ -36,21 +41,221 @@ use super::MockKeyringStore;
use super::TempCodexHome;
use super::assert_tokens_match_without_expiry;
use super::sample_tokens;
use crate::http_client_adapter::StreamableHttpClientAdapter;
use crate::http_client_adapter::StreamableHttpRedirectMode;
use crate::oauth::OAuthCredentialStore;
use crate::oauth::OAuthPersistor;
use crate::oauth::OAuthRuntime;
use crate::oauth::ResolvedOAuthCredentialStore;
use crate::oauth::StoredOAuthTokens;
use crate::oauth::WrappedOAuthTokenResponse;
use crate::oauth::compute_store_key;
use crate::oauth::delete_oauth_tokens;
use crate::oauth::load_oauth_tokens_from_file;
use crate::oauth::refresh_lock::RefreshCredentialLock;
use crate::oauth::save_oauth_tokens;
use crate::oauth::save_oauth_tokens_to_file;
use crate::oauth::stored_oauth_credentials;
use crate::oauth_http_client::OAuthHttpClientAdapter;
use crate::oauth_http_client::PROACTIVE_REFRESH_TIMEOUT;
use crate::startup_error::is_authentication_required_error;
const REFRESH_LOCK_CONTENTION_EVENT_TARGET: &str =
"codex_rmcp_client::oauth::refresh_lock::contention";
#[tokio::test(flavor = "current_thread")]
async fn login_and_logout_follow_the_completed_refresh() -> Result<()> {
let (_env, _server, initial) = test_context().await?;
let mut replacement = initial.clone();
replacement
.token_response
.0
.set_access_token(AccessToken::new("replacement-access-token".into()));
replacement
.token_response
.0
.set_refresh_token(Some(RefreshToken::new("replacement-refresh-token".into())));
let mut refreshed = initial.clone();
refreshed
.token_response
.0
.set_refresh_token(Some(RefreshToken::new("rotated-refresh-token".into())));
for expected in [Some(replacement), None] {
save_oauth_tokens_to_file(&initial)?;
let held_lock =
RefreshCredentialLock::acquire_for_server(&initial.server_name, &initial.url).await?;
let (contended_tx, contended_rx) = mpsc::channel();
let _subscriber_guard =
tracing::subscriber::set_default(LockContentionSubscriber { contended_tx });
let mutation = tokio::spawn({
let initial = initial.clone();
let expected = expected.clone();
async move {
match expected {
Some(tokens) => {
save_oauth_tokens(
&tokens.server_name,
&tokens,
OAuthCredentialsStoreMode::File,
AuthKeyringBackendKind::Direct,
)
.await
}
None => {
assert!(
delete_oauth_tokens(
&initial.server_name,
&initial.url,
OAuthCredentialsStoreMode::File,
AuthKeyringBackendKind::Direct,
)
.await?
);
Ok(())
}
}
}
});
wait_for_lock_contention(contended_rx, /*expected_count*/ 1).await?;
assert!(!mutation.is_finished());
assert_eq!(
load_oauth_tokens_from_file(&initial.server_name, &initial.url)?,
Some(initial.clone())
);
// Complete the in-flight refresh before allowing the login or logout to write.
save_oauth_tokens_to_file(&refreshed)?;
drop(held_lock);
mutation.await??;
assert_eq!(
load_oauth_tokens_from_file(&initial.server_name, &initial.url)?,
expected
);
}
Ok(())
}
#[tokio::test(flavor = "current_thread")]
async fn fresh_preparation_skips_storage_and_refresh_lock() -> Result<()> {
let (env, server, mut initial) = test_context().await?;
let fresh = sample_tokens();
initial.token_response = fresh.token_response;
initial.expires_at = fresh.expires_at;
std::fs::create_dir(env.path().join(".credentials.json"))?;
Mock::given(method("POST"))
.and(path("/oauth/token"))
.respond_with(ResponseTemplate::new(500))
.expect(0)
.mount(&server)
.await;
for mode in [
crate::McpOAuthRefreshMode::Legacy,
crate::McpOAuthRefreshMode::Coordinated,
] {
let runtime = runtime_for(&initial, mode).await?;
let held_lock =
RefreshCredentialLock::acquire_for_server(&initial.server_name, &initial.url).await?;
tokio::time::timeout(
Duration::from_millis(/*millis*/ 100),
runtime.refresh_if_needed(),
)
.await
.context("fresh preparation must not wait for the refresh lock")??;
drop(held_lock);
}
server.verify().await;
Ok(())
}
#[tokio::test(flavor = "current_thread")]
async fn coordinated_401_refresh_rereads_and_persists_before_retry() -> Result<()> {
let (_env, server, mut initial) = test_context().await?;
let fresh = sample_tokens();
initial.token_response = fresh.token_response;
initial.expires_at = fresh.expires_at;
save_oauth_tokens_to_file(&initial)?;
let mut manager = authorization_manager_for(&initial).await?;
manager.set_credential_store(OAuthCredentialStore::new(
initial.clone(),
ResolvedOAuthCredentialStore::File,
DefaultKeyringStore,
));
let client = AuthClient::new(
StreamableHttpClientAdapter::new(
Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new(
OutboundProxyPolicy::ReqwestDefault,
))),
HeaderMap::new(),
/*auth_provider*/ None,
/*has_configured_headers*/ false,
StreamableHttpRedirectMode::Legacy,
Arc::default(),
),
manager,
);
let mut durable = initial.clone();
durable
.token_response
.0
.set_access_token(AccessToken::new("durable-access-token".into()));
durable
.token_response
.0
.set_refresh_token(Some(RefreshToken::new("durable-refresh-token".into())));
save_oauth_tokens_to_file(&durable)?;
Mock::given(method("POST"))
.and(path("/mcp"))
.and(header("authorization", "Bearer access-token"))
.respond_with(ResponseTemplate::new(401).insert_header("www-authenticate", "Bearer"))
.expect(1)
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/oauth/token"))
.and(body_string_contains("refresh_token=durable-refresh-token"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"access_token": "refreshed-access-token",
"refresh_token": "rotated-refresh-token",
"token_type": "Bearer",
"expires_in": 3600,
})))
.expect(1)
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/mcp"))
.and(header("authorization", "Bearer refreshed-access-token"))
.respond_with(move |_: &wiremock::Request| {
let saved = load_oauth_tokens_from_file(&durable.server_name, &durable.url)
.expect("read persisted refresh before retry")
.expect("refresh must be persisted before retry");
let mut expected = durable.token_response.0.clone();
expected.set_access_token(AccessToken::new("refreshed-access-token".into()));
expected.set_refresh_token(Some(RefreshToken::new("rotated-refresh-token".into())));
expected.set_expires_in(saved.token_response.0.expires_in().as_ref());
assert_eq!(saved.token_response, WrappedOAuthTokenResponse(expected));
ResponseTemplate::new(202)
})
.expect(1)
.mount(&server)
.await;
client
.post_message(
initial.url.into(),
serde_json::from_value(
serde_json::json!({"jsonrpc": "2.0", "id": 1, "method": "ping"}),
)?,
/*session_id*/ None,
/*auth_token*/ None,
Default::default(),
)
.await?;
server.verify().await;
Ok(())
}
struct LockContentionSubscriber {
contended_tx: mpsc::Sender<()>,
}
@@ -95,6 +300,11 @@ impl Subscriber for LockContentionSubscriber {
#[tokio::test(flavor = "current_thread")]
async fn concurrent_refreshes_call_provider_once_and_carry_omitted_fields() -> Result<()> {
assert_concurrent_refreshes(crate::McpOAuthRefreshMode::Legacy).await?;
assert_concurrent_refreshes(crate::McpOAuthRefreshMode::Coordinated).await
}
async fn assert_concurrent_refreshes(mode: crate::McpOAuthRefreshMode) -> Result<()> {
let (_env, server, initial) = test_context().await?;
Mock::given(method("POST"))
.and(path("/oauth/token"))
@@ -119,9 +329,8 @@ async fn concurrent_refreshes_call_provider_once_and_carry_omitted_fields() -> R
let _subscriber_guard =
tracing::subscriber::set_default(LockContentionSubscriber { contended_tx });
let first = persistor_for(&initial).await?;
let second = persistor_for(&initial).await?;
let initial_credentials = first.stored_credentials().await;
let first = runtime_for(&initial, mode).await?;
let second = runtime_for(&initial, mode).await?;
let first_task = tokio::spawn({
let first = first.clone();
async move { first.refresh_if_needed().await }
@@ -140,11 +349,16 @@ async fn concurrent_refreshes_call_provider_once_and_carry_omitted_fields() -> R
// 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?;
if let OAuthRuntime::Legacy(persistor) = &first {
persistor.persist_if_needed().await?;
}
let stored = load_oauth_tokens_from_file(&initial.server_name, &initial.url)?
.expect("refreshed credentials should be stored");
assert_eq!(stored.issuer, initial.issuer);
let live_credentials = first.stored_credentials().await;
let live_credentials = match &first {
OAuthRuntime::Legacy(persistor) => persistor.stored_credentials().await,
OAuthRuntime::Coordinated { store, .. } => store.stored_credentials().await,
};
let disk_credentials = stored_oauth_credentials(
&initial.server_name,
&initial.url,
@@ -152,7 +366,6 @@ async fn concurrent_refreshes_call_provider_once_and_carry_omitted_fields() -> R
AuthKeyringBackendKind::Direct,
)?;
assert_eq!(live_credentials, disk_credentials);
assert_ne!(live_credentials, initial_credentials);
let mut expected_response = initial.token_response.0.clone();
expected_response.set_access_token(AccessToken::new("refreshed-access-token".to_string()));
// File loads derive `expires_in` from stable `expires_at`, so it may tick down before this
@@ -176,7 +389,7 @@ async fn resolved_keyring_read_error_preserves_in_memory_credentials() -> Result
let keyring_store = MockKeyringStore::default();
let key = compute_store_key(&initial.server_name, &initial.url)?;
keyring_store.set_error(&key, KeyringError::Invalid("error".into(), "load".into()));
let manager = authorization_manager_for(&initial).await?;
let manager = Arc::new(TokioMutex::new(authorization_manager_for(&initial).await?));
let persistor = OAuthPersistor::new(
initial.server_name.clone(),
initial.url.clone(),
@@ -354,6 +567,11 @@ async fn transient_refresh_failure_does_not_require_reauthorization() -> Result<
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn caller_cancellation_does_not_cancel_refresh_persistence() -> Result<()> {
assert_caller_cancellation(crate::McpOAuthRefreshMode::Legacy).await?;
assert_caller_cancellation(crate::McpOAuthRefreshMode::Coordinated).await
}
async fn assert_caller_cancellation(mode: crate::McpOAuthRefreshMode) -> Result<()> {
let (_env, server, initial) = test_context().await?;
let (request_received_tx, request_received_rx) = mpsc::channel();
let (release_response_tx, release_response_rx) = mpsc::channel();
@@ -381,11 +599,8 @@ async fn caller_cancellation_does_not_cancel_refresh_persistence() -> Result<()>
.mount(&server)
.await;
save_oauth_tokens_to_file(&initial)?;
let persistor = persistor_for(&initial).await?;
let caller = tokio::spawn({
let persistor = persistor.clone();
async move { persistor.refresh_if_needed().await }
});
let runtime = runtime_for(&initial, mode).await?;
let caller = tokio::spawn(async move { runtime.refresh_if_needed().await });
tokio::task::spawn_blocking(move || {
request_received_rx
@@ -452,11 +667,83 @@ async fn provider_timeout_releases_lock_and_preserves_durable_credentials() -> 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 coordinated_provider_timeout_excludes_lock_wait() -> Result<()> {
let (_env, server, initial) = test_context().await?;
mount_delayed_refresh(&server, "late-access-token").await;
save_oauth_tokens_to_file(&initial)?;
let (auth_manager, _) = coordinated_manager_for(&initial).await?;
let held_lock =
RefreshCredentialLock::acquire_for_server(&initial.server_name, &initial.url).await?;
let (contended_tx, contended_rx) = mpsc::channel();
let _subscriber_guard =
tracing::subscriber::set_default(LockContentionSubscriber { contended_tx });
let mut refresh = tokio::spawn(
PROACTIVE_REFRESH_TIMEOUT.scope(Duration::from_millis(/*millis*/ 50), async move {
auth_manager.lock().await.get_access_token().await
}),
);
wait_for_lock_contention(contended_rx, /*expected_count*/ 1).await?;
assert!(
tokio::time::timeout(Duration::from_millis(/*millis*/ 100), &mut refresh)
.await
.is_err(),
"lock wait must not consume the provider timeout"
);
drop(held_lock);
let error = refresh
.await?
.expect_err("the scoped provider request must time out");
assert!(matches!(error, AuthError::TokenRefreshFailed(_)));
let _lock = tokio::time::timeout(
Duration::from_millis(/*millis*/ 100),
RefreshCredentialLock::acquire_for_server(&initial.server_name, &initial.url),
)
.await??;
let stored = load_oauth_tokens_from_file(&initial.server_name, &initial.url)?.unwrap();
assert_tokens_match_without_expiry(&stored, &initial);
server.verify().await;
Ok(())
}
async fn runtime_for(
tokens: &StoredOAuthTokens,
mode: crate::McpOAuthRefreshMode,
) -> Result<OAuthRuntime> {
Ok(match mode {
crate::McpOAuthRefreshMode::Legacy => OAuthRuntime::Legacy(persistor_for(tokens).await?),
crate::McpOAuthRefreshMode::Coordinated => {
let (auth_manager, store) = coordinated_manager_for(tokens).await?;
OAuthRuntime::Coordinated {
auth_manager,
store,
}
}
})
}
async fn coordinated_manager_for(
tokens: &StoredOAuthTokens,
) -> Result<(Arc<TokioMutex<AuthorizationManager>>, OAuthCredentialStore)> {
let mut manager = authorization_manager_for(tokens).await?;
let store = OAuthCredentialStore::new(
tokens.clone(),
ResolvedOAuthCredentialStore::File,
DefaultKeyringStore,
);
manager.set_credential_store(store.clone());
Ok((Arc::new(TokioMutex::new(manager)), store))
}
async fn persistor_for(tokens: &StoredOAuthTokens) -> Result<OAuthPersistor> {
Ok(OAuthPersistor::new(
tokens.server_name.clone(),
tokens.url.clone(),
authorization_manager_for(tokens).await?,
Arc::new(TokioMutex::new(authorization_manager_for(tokens).await?)),
ResolvedOAuthCredentialStore::File,
Some(tokens.clone()),
))
@@ -470,9 +757,7 @@ async fn test_context() -> Result<(TempCodexHome, MockServer, StoredOAuthTokens)
Ok((env, server, tokens))
}
async fn authorization_manager_for(
tokens: &StoredOAuthTokens,
) -> Result<Arc<TokioMutex<AuthorizationManager>>> {
async fn authorization_manager_for(tokens: &StoredOAuthTokens) -> Result<AuthorizationManager> {
let oauth_http_client = Arc::new(OAuthHttpClientAdapter::new(
Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new(
OutboundProxyPolicy::ReqwestDefault,
@@ -492,7 +777,7 @@ async fn authorization_manager_for(
}
_ => anyhow::bail!("unexpected OAuth state"),
};
Ok(Arc::new(TokioMutex::new(manager)))
Ok(manager)
}
async fn mount_oauth_metadata(server: &MockServer) {
codex-rs/rmcp-client/src/oauth_http_client.rs
+15 -1
View File
@@ -38,6 +38,11 @@ const MAX_OAUTH_HTTP_RESPONSE_BODY_BYTES: usize = 1024 * 1024;
const MAX_OAUTH_HTTP_REDIRECTS: usize = 10;
static NEXT_OAUTH_REQUEST_ID: AtomicU64 = AtomicU64::new(0);
tokio::task_local! {
/// Bounds provider HTTP work during preparation, excluding credential lock waits and saves.
pub(crate) static PROACTIVE_REFRESH_TIMEOUT: Duration;
}
#[derive(Debug, thiserror::Error)]
enum OAuthHttpClientAdapterError {
#[error("unsupported OAuth HTTP redirect policy")]
@@ -324,7 +329,16 @@ fn oauth_redirect_policy(
impl OAuthHttpClient for OAuthHttpClientAdapter {
fn execute(&self, request: OAuthHttpRequest) -> OAuthHttpClientFuture<'_> {
Box::pin(self.execute_request(request.request, request.redirect_policy, request.timeout))
Box::pin(async move {
let operation =
self.execute_request(request.request, request.redirect_policy, request.timeout);
match PROACTIVE_REFRESH_TIMEOUT.try_with(|duration| *duration) {
Ok(duration) => tokio::time::timeout(duration, operation)
.await
.map_err(|_| oauth_http_client_error(OAuthHttpClientAdapterError::TimedOut))?,
Err(_) => operation.await,
}
})
}
}
codex-rs/rmcp-client/src/perform_oauth_login.rs
+2 -1
View File
@@ -721,7 +721,8 @@ impl OauthLoginFlow {
&stored,
self.store_mode,
self.keyring_backend_kind,
)?;
)
.await?;
Ok(())
}
codex-rs/rmcp-client/src/rmcp_client.rs
+59 -28
View File
@@ -77,7 +77,9 @@ use crate::http_client_adapter::StreamableHttpClientAdapter;
use crate::http_client_adapter::StreamableHttpClientAdapterError;
use crate::http_client_adapter::StreamableHttpRedirectMode;
use crate::in_process_transport::InProcessTransportFactory;
use crate::oauth::OAuthCredentialStore;
use crate::oauth::OAuthPersistor;
use crate::oauth::OAuthRuntime;
use crate::oauth::ResolvedOAuthCredentialStore;
use crate::oauth::ResolvedOAuthTokens;
use crate::oauth::StoredOAuthTokens;
@@ -85,6 +87,7 @@ use crate::oauth::install_tokens_in_manager;
use crate::oauth::resolve_oauth_tokens_from_store_policy;
use crate::oauth::validate_refresh_token_issuer;
use crate::oauth_http_client::OAuthHttpClientAdapter;
use crate::oauth_refresh_mode::McpOAuthRefreshMode;
use crate::protocol_mode::McpProtocolMode;
use crate::stdio_server_launcher::StdioServerCommand;
use crate::stdio_server_launcher::StdioServerLauncher;
@@ -112,7 +115,7 @@ enum PendingTransport {
},
StreamableHttpWithOAuth {
transport: StreamableHttpClientTransport<AuthClient<StreamableHttpClientAdapter>>,
oauth_persistor: OAuthPersistor,
oauth_runtime: OAuthRuntime,
},
StreamableHttpWithAccessTokenOnly {
transport: StreamableHttpClientTransport<AuthClient<StreamableHttpClientAdapter>>,
@@ -125,7 +128,7 @@ enum ClientState {
},
Ready {
service: Arc<RunningService<RoleClient, ElicitationClientService>>,
oauth: Option<OAuthPersistor>,
oauth: Option<OAuthRuntime>,
},
Closed,
}
@@ -160,6 +163,7 @@ enum TransportRecipe {
http_client: Arc<dyn HttpClient>,
auth_provider: Option<SharedAuthProvider>,
redirect_mode: StreamableHttpRedirectMode,
oauth_refresh_mode: McpOAuthRefreshMode,
initialize_deadline: Arc<StdMutex<Option<Instant>>>,
},
}
@@ -539,6 +543,7 @@ impl RmcpClient {
auth_provider,
protocol_mode,
StreamableHttpRedirectMode::Legacy,
McpOAuthRefreshMode::Legacy,
)
.await
}
@@ -556,6 +561,7 @@ impl RmcpClient {
auth_provider: Option<SharedAuthProvider>,
protocol_mode: McpProtocolMode,
redirect_mode: StreamableHttpRedirectMode,
oauth_refresh_mode: McpOAuthRefreshMode,
) -> Result<Self> {
let transport_recipe = TransportRecipe::StreamableHttp {
server_name: server_name.to_string(),
@@ -569,6 +575,7 @@ impl RmcpClient {
http_client,
auth_provider,
redirect_mode,
oauth_refresh_mode,
initialize_deadline: Arc::new(StdMutex::new(None)),
};
let transport = Self::create_pending_transport(&transport_recipe).await?;
@@ -611,7 +618,7 @@ impl RmcpClient {
}
};
let (service, oauth_persistor) = self
let (service, oauth_runtime) = self
.connect_pending_transport_with_initialize_retries(
pending_transport,
client_service.clone(),
@@ -640,11 +647,11 @@ impl RmcpClient {
}
*guard = ClientState::Ready {
service,
oauth: oauth_persistor.clone(),
oauth: oauth_runtime.clone(),
};
}
if let Some(runtime) = oauth_persistor
if let Some(OAuthRuntime::Legacy(runtime)) = oauth_runtime
&& let Err(error) = runtime.persist_if_needed().await
{
warn!("failed to persist OAuth tokens after initialize: {error}");
@@ -937,7 +944,7 @@ impl RmcpClient {
}
}
async fn oauth_persistor(&self) -> Option<OAuthPersistor> {
async fn oauth_runtime(&self) -> Option<OAuthRuntime> {
let guard = self.state.lock().await;
match &*guard {
ClientState::Ready {
@@ -950,8 +957,11 @@ impl RmcpClient {
/// Returns `None` when this client does not manage stored OAuth credentials.
pub async fn managed_oauth_credentials(&self) -> Option<Option<StoredOAuthTokens>> {
let persistor = self.oauth_persistor().await?;
Some(persistor.stored_credentials().await)
let runtime = self.oauth_runtime().await?;
Some(match runtime {
OAuthRuntime::Legacy(persistor) => persistor.stored_credentials().await,
OAuthRuntime::Coordinated { store, .. } => store.stored_credentials().await,
})
}
/// Returns whether an initialized transport or its underlying service has stopped.
@@ -985,16 +995,16 @@ impl RmcpClient {
/// 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
if let Some(OAuthRuntime::Legacy(runtime)) = self.oauth_runtime().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.
/// Prepares managed OAuth before the MCP operation timeout starts.
async fn refresh_oauth_if_needed(&self) -> Result<()> {
if let Some(runtime) = self.oauth_persistor().await {
if let Some(runtime) = self.oauth_runtime().await {
runtime.refresh_if_needed().await?;
}
Ok(())
@@ -1026,6 +1036,7 @@ impl RmcpClient {
http_client,
auth_provider,
redirect_mode,
oauth_refresh_mode,
initialize_deadline,
} => {
let has_configured_headers = matches!(
@@ -1112,6 +1123,7 @@ impl RmcpClient {
Arc::clone(http_client),
has_configured_headers,
*redirect_mode,
*oauth_refresh_mode,
Arc::clone(initialize_deadline),
)
.await
@@ -1180,7 +1192,7 @@ impl RmcpClient {
timeout: Option<Duration>,
) -> Result<(
Arc<RunningService<RoleClient, ElicitationClientService>>,
Option<OAuthPersistor>,
Option<OAuthRuntime>,
)> {
let _initialize_deadline = match &self.transport_recipe {
TransportRecipe::StreamableHttp {
@@ -1198,7 +1210,7 @@ impl RmcpClient {
TransportRecipe::InProcess { .. } | TransportRecipe::Stdio { .. } => None,
};
let lifecycle = self.protocol_mode.client_lifecycle();
let (transport, oauth_persistor) = match pending_transport {
let (transport, oauth_runtime) = match pending_transport {
PendingTransport::InProcess { transport } => (
client_service
.serve_with_lifecycle(transport, lifecycle)
@@ -1219,12 +1231,12 @@ impl RmcpClient {
),
PendingTransport::StreamableHttpWithOAuth {
transport,
oauth_persistor,
oauth_runtime,
} => (
client_service
.serve_with_lifecycle(transport, lifecycle)
.boxed(),
Some(oauth_persistor),
Some(oauth_runtime),
),
PendingTransport::StreamableHttpWithAccessTokenOnly { transport } => (
client_service
@@ -1250,7 +1262,7 @@ impl RmcpClient {
let service = match service_result {
Ok(service) => service,
Err(error) => {
if let Some(runtime) = oauth_persistor.as_ref()
if let Some(OAuthRuntime::Legacy(runtime)) = oauth_runtime.as_ref()
&& let Err(persist_error) = runtime.persist_if_needed().await
{
warn!(
@@ -1268,7 +1280,7 @@ impl RmcpClient {
.set_response_cache_config(ClientCacheConfig::disabled())
.await;
Ok((Arc::new(service), oauth_persistor))
Ok((Arc::new(service), oauth_runtime))
}
async fn run_service_operation<T, F, Fut>(
@@ -1459,7 +1471,7 @@ impl RmcpClient {
.clone()
.ok_or_else(|| anyhow!("MCP client cannot recover before initialize succeeds"))?;
let pending_transport = Self::create_pending_transport(&self.transport_recipe).await?;
let (service, oauth_persistor) = self
let (service, oauth_runtime) = self
.connect_pending_transport_with_initialize_retries(
pending_transport,
initialize_context.client_service,
@@ -1478,11 +1490,11 @@ impl RmcpClient {
}
*guard = ClientState::Ready {
service,
oauth: oauth_persistor.clone(),
oauth: oauth_runtime.clone(),
};
}
if let Some(runtime) = oauth_persistor
if let Some(OAuthRuntime::Legacy(runtime)) = oauth_runtime
&& let Err(error) = runtime.persist_if_needed().await
{
warn!("failed to persist OAuth tokens after session recovery: {error}");
@@ -1502,6 +1514,7 @@ async fn create_oauth_transport_and_runtime(
http_client: Arc<dyn HttpClient>,
has_configured_headers: bool,
redirect_mode: StreamableHttpRedirectMode,
oauth_refresh_mode: McpOAuthRefreshMode,
initialize_deadline: Arc<StdMutex<Option<Instant>>>,
) -> Result<PendingTransport> {
let oauth_http_client = Arc::new(OAuthHttpClientAdapter::new_with_redirect_mode(
@@ -1533,6 +1546,18 @@ async fn create_oauth_transport_and_runtime(
runtime_tokens.issuer = None;
}
install_tokens_in_manager(&mut manager, &runtime_tokens).await?;
let coordinated_store = match oauth_refresh_mode {
McpOAuthRefreshMode::Coordinated if !use_stored_access_token_only => {
let store = OAuthCredentialStore::new(
initial_tokens.clone(),
credential_store,
DefaultKeyringStore,
);
manager.set_credential_store(store.clone());
Some(store)
}
McpOAuthRefreshMode::Legacy | McpOAuthRefreshMode::Coordinated => None,
};
let auth_client = AuthClient::new(
StreamableHttpClientAdapter::new(
@@ -1559,17 +1584,23 @@ async fn create_oauth_transport_and_runtime(
return Ok(PendingTransport::StreamableHttpWithAccessTokenOnly { transport });
}
let runtime = OAuthPersistor::new(
server_name.to_string(),
url.to_string(),
auth_manager,
credential_store,
Some(initial_tokens),
);
let runtime = match coordinated_store {
Some(store) => OAuthRuntime::Coordinated {
auth_manager,
store,
},
None => OAuthRuntime::Legacy(OAuthPersistor::new(
server_name.to_string(),
url.to_string(),
auth_manager,
credential_store,
Some(initial_tokens),
)),
};
Ok(PendingTransport::StreamableHttpWithOAuth {
transport,
oauth_persistor: runtime,
oauth_runtime: runtime,
})
}
codex-rs/rmcp-client/src/streamable_http_retry.rs
+4 -7
View File
@@ -14,7 +14,7 @@ use tracing::warn;
use crate::elicitation_client_service::ElicitationClientService;
use crate::http_client_adapter::StreamableHttpClientAdapterError;
use crate::oauth::OAuthPersistor;
use crate::oauth::OAuthRuntime;
use super::PendingTransport;
use super::RmcpClient;
@@ -30,7 +30,7 @@ impl RmcpClient {
timeout: Option<Duration>,
) -> Result<(
Arc<RunningService<RoleClient, ElicitationClientService>>,
Option<OAuthPersistor>,
Option<OAuthRuntime>,
)> {
let should_retry = match &initial_transport {
PendingTransport::InProcess { .. } | PendingTransport::Stdio { .. } => false,
@@ -63,14 +63,11 @@ impl RmcpClient {
}
}
};
if let PendingTransport::StreamableHttpWithOAuth {
oauth_persistor, ..
} = &transport
{
if let PendingTransport::StreamableHttpWithOAuth { oauth_runtime, .. } = &transport {
// OAuth refresh has its own lock and provider request bounds. Exclude it from the
// MCP handshake budget, and finish persistence before attempting initialize.
let refresh_started_at = Instant::now();
oauth_persistor.refresh_if_needed().await?;
oauth_runtime.refresh_if_needed().await?;
if let Some(deadline) = retry_deadline.as_mut() {
*deadline += refresh_started_at.elapsed();
}
codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs
+70 -15
View File
@@ -13,6 +13,8 @@ use codex_config::types::OAuthCredentialsStoreMode;
use codex_exec_server::Environment;
use codex_rmcp_client::McpAuthState;
use codex_rmcp_client::McpLoginRequirement;
use codex_rmcp_client::McpOAuthRefreshMode;
use codex_rmcp_client::McpProtocolMode;
use codex_rmcp_client::OAuthDiscoveryTimeout;
use codex_rmcp_client::RmcpClient;
use codex_rmcp_client::StoredOAuthTokens;
@@ -43,6 +45,7 @@ use wiremock::matchers::method;
use wiremock::matchers::path;
use streamable_http_test_support::initialize_client;
use streamable_http_test_support::initialize_client_with_timeout;
const SERVER_NAME: &str = "test-streamable-http-oauth-startup";
const EXPIRED_ACCESS_TOKEN: &str = "expired-access-token";
@@ -52,6 +55,7 @@ const RESOURCE_API_KEY: &str = "resource-api-key-secret";
const RESOURCE_USER_AGENT: &str = "resource-only-user-agent";
const MCP_USER_AGENT: &str = concat!("codex-mcp-client/", env!("CARGO_PKG_VERSION"));
const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_STARTUP_SERVER_URL";
const CHILD_REFRESH_MODE_ENV: &str = "MCP_TEST_OAUTH_REFRESH_MODE";
const CHILD_HELPER_COMMAND_ENV: &str = "MCP_TEST_OAUTH_STARTUP_HELPER_COMMAND";
const CHILD_RESOURCE_API_KEY_ENV: &str = "MCP_TEST_OAUTH_STARTUP_RESOURCE_API_KEY";
const CHILD_STORED_ISSUER_ENV: &str = "MCP_TEST_OAUTH_STARTUP_STORED_ISSUER";
@@ -76,25 +80,44 @@ enum IssuerMismatchAccessToken {
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result<()> {
assert_expired_token_refresh(OAuthStartupScenario::DirectAuthorizationMetadata).await
assert_expired_token_refresh(
OAuthStartupScenario::DirectAuthorizationMetadata,
McpOAuthRefreshMode::Coordinated,
)
.await
}
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn refreshes_oauth_with_gateway_headers_helper() -> anyhow::Result<()> {
assert_expired_token_refresh(OAuthStartupScenario::GatewayHeadersHelper).await
assert_expired_token_refresh(
OAuthStartupScenario::GatewayHeadersHelper,
McpOAuthRefreshMode::Legacy,
)
.await
}
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn refreshes_oauth_after_gateway_rejects_token_request() -> anyhow::Result<()> {
assert_expired_token_refresh(OAuthStartupScenario::SameOriginGatewayHeadersHelper).await
assert_expired_token_refresh(
OAuthStartupScenario::SameOriginGatewayHeadersHelper,
McpOAuthRefreshMode::Legacy,
)
.await
}
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn refresh_uses_discovered_protected_resource_audience() -> anyhow::Result<()> {
assert_expired_token_refresh(OAuthStartupScenario::ProtectedResourceMetadata).await
assert_expired_token_refresh(
OAuthStartupScenario::ProtectedResourceMetadata,
McpOAuthRefreshMode::Legacy,
)
.await
}
async fn assert_expired_token_refresh(scenario: OAuthStartupScenario) -> anyhow::Result<()> {
async fn assert_expired_token_refresh(
scenario: OAuthStartupScenario,
refresh_mode: McpOAuthRefreshMode,
) -> anyhow::Result<()> {
let server = MockServer::start().await;
let authorization_server = MockServer::start().await;
let same_origin_gateway = matches!(
@@ -198,12 +221,18 @@ async fn assert_expired_token_refresh(scenario: OAuthStartupScenario) -> anyhow:
{
ResponseTemplate::new(/*s*/ 401)
} else {
ResponseTemplate::new(/*s*/ 200).set_body_json(json!({
let response = ResponseTemplate::new(/*s*/ 200).set_body_json(json!({
"access_token": REFRESHED_ACCESS_TOKEN,
"token_type": "Bearer",
"expires_in": 7200,
"refresh_token": REFRESH_TOKEN,
}))
}));
if refresh_mode == McpOAuthRefreshMode::Coordinated {
// Longer than the child's handshake timeout: refresh must finish first.
response.set_delay(Duration::from_secs(/*secs*/ 2))
} else {
response
}
}
})
.expect(if same_origin_gateway { 2 } else { 1 })
@@ -266,6 +295,14 @@ async fn assert_expired_token_refresh(scenario: OAuthStartupScenario) -> anyhow:
.env(CHILD_SERVER_URL_ENV, server_url)
.env(CHILD_STORED_ISSUER_ENV, authorization_server_issuer)
.env(CHILD_RESOURCE_API_KEY_ENV, RESOURCE_API_KEY)
.env(
CHILD_REFRESH_MODE_ENV,
if refresh_mode == McpOAuthRefreshMode::Coordinated {
"coordinated"
} else {
"legacy"
},
)
.env("MCP_TEST_AMBIENT_SECRET", "must-not-reach-helper");
if with_headers_helper {
command.env(
@@ -632,7 +669,8 @@ async fn persisted_credentials_auth_status_child() -> anyhow::Result<()> {
&tokens,
OAuthCredentialsStoreMode::File,
AuthKeyringBackendKind::default(),
)?;
)
.await?;
let status = auth_status(LEGACY_REFRESHABLE_SERVER_URL).await?;
assert_eq!(
@@ -662,7 +700,8 @@ async fn persisted_credentials_auth_status_child() -> anyhow::Result<()> {
&tokens,
OAuthCredentialsStoreMode::File,
AuthKeyringBackendKind::default(),
)?;
)
.await?;
let status = auth_status(UNEXPIRED_SERVER_URL).await?;
assert_eq!(status, McpAuthState::OAuth);
@@ -686,7 +725,8 @@ async fn persisted_credentials_auth_status_child() -> anyhow::Result<()> {
&tokens,
OAuthCredentialsStoreMode::File,
AuthKeyringBackendKind::default(),
)?;
)
.await?;
let status = auth_status(REFRESHABLE_SERVER_URL).await?;
assert_eq!(status, McpAuthState::OAuth);
@@ -713,6 +753,11 @@ async fn auth_status(server_url: &str) -> anyhow::Result<McpAuthState> {
#[ignore = "spawned by refreshes_expired_persisted_token_before_initialize"]
async fn oauth_startup_child() -> anyhow::Result<()> {
let server_url = std::env::var(CHILD_SERVER_URL_ENV)?;
let refresh_mode = if std::env::var(CHILD_REFRESH_MODE_ENV).as_deref() == Ok("coordinated") {
McpOAuthRefreshMode::Coordinated
} else {
McpOAuthRefreshMode::Legacy
};
// Save an expired access token with a valid refresh token so startup must
// refresh before sending the initialize request.
@@ -736,7 +781,8 @@ async fn oauth_startup_child() -> anyhow::Result<()> {
&tokens,
OAuthCredentialsStoreMode::File,
AuthKeyringBackendKind::default(),
)?;
)
.await?;
// This mirrors create_client's transport and initialization setup, except
// it omits the direct bearer token. Supplying that token would bypass the
@@ -750,7 +796,7 @@ async fn oauth_startup_child() -> anyhow::Result<()> {
std::env::current_dir()?,
)?;
}
let client = RmcpClient::new_streamable_http_client(
let client = RmcpClient::new_streamable_http_client_with_protocol_mode_and_redirect_mode(
SERVER_NAME,
&server_url,
/*bearer_token*/ None,
@@ -766,10 +812,17 @@ async fn oauth_startup_child() -> anyhow::Result<()> {
AuthKeyringBackendKind::default(),
http_client,
/*auth_provider*/ None,
McpProtocolMode::Legacy,
StreamableHttpRedirectMode::Legacy,
refresh_mode,
)
.await?;
initialize_client(&client).await?;
if refresh_mode == McpOAuthRefreshMode::Coordinated {
initialize_client_with_timeout(&client, Duration::from_secs(/*secs*/ 1)).await?;
} else {
initialize_client(&client).await?;
}
Ok(())
}
@@ -795,7 +848,8 @@ async fn expired_unrefreshable_startup_child() -> anyhow::Result<()> {
&tokens,
OAuthCredentialsStoreMode::File,
AuthKeyringBackendKind::default(),
)?;
)
.await?;
let client = RmcpClient::new_streamable_http_client(
SERVER_NAME,
@@ -852,7 +906,8 @@ async fn issuer_startup_child() -> anyhow::Result<()> {
&tokens,
OAuthCredentialsStoreMode::File,
AuthKeyringBackendKind::default(),
)?;
)
.await?;
let client = RmcpClient::new_streamable_http_client(
SERVER_NAME,
codex-rs/rmcp-client/tests/streamable_http_oauth_store_pinning.rs
+7 -3
View File
@@ -203,6 +203,7 @@ async fn transport_provided_bearer_token_avoids_placeholder_headers_and_redirect
/*auth_provider*/ None,
McpProtocolMode::Legacy,
StreamableHttpRedirectMode::AgentPluginV1,
codex_rmcp_client::McpOAuthRefreshMode::Legacy,
)
.await?;
@@ -267,7 +268,8 @@ async fn auto_store_remains_pinned_across_session_recovery_child() -> anyhow::Re
&file_tokens,
OAuthCredentialsStoreMode::File,
AuthKeyringBackendKind::Direct,
)?;
)
.await?;
let file_snapshot = stored_oauth_credential_snapshot(
SERVER_NAME,
&server_url,
@@ -281,13 +283,15 @@ async fn auto_store_remains_pinned_across_session_recovery_child() -> anyhow::Re
&keyring_tokens,
OAuthCredentialsStoreMode::Auto,
AuthKeyringBackendKind::Direct,
)?;
)
.await?;
save_oauth_tokens(
SERVER_NAME,
&file_tokens,
OAuthCredentialsStoreMode::File,
AuthKeyringBackendKind::Direct,
)?;
)
.await?;
assert_eq!(
file_snapshot.reload(
SERVER_NAME,
codex-rs/rmcp-client/tests/streamable_http_test_support.rs
+8 -1
View File
@@ -105,10 +105,17 @@ pub(crate) async fn create_client_with_http_client(
}
pub(crate) async fn initialize_client(client: &RmcpClient) -> anyhow::Result<()> {
initialize_client_with_timeout(client, Duration::from_secs(/*secs*/ 5)).await
}
pub(crate) async fn initialize_client_with_timeout(
client: &RmcpClient,
timeout: Duration,
) -> anyhow::Result<()> {
client
.initialize(
init_params(),
Some(Duration::from_secs(5)),
Some(timeout),
Box::new(|_, _| {
async {
Ok(ElicitationResponse {