openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-05 20:10:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add Codex-owned MCP OAuth recovery

Browse Source
This commit is contained in:
Steven Lee
2026-06-25 18:31:50 +00:00
parent f800bd9e4a
commit 448dcbd8c1
5 changed files with 702 additions and 741 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/rmcp-client/src/oauth.rs
+326 -417
View File
@@ -61,8 +61,6 @@ use codex_keyring_store::KeyringStore;
use codex_utils_home_dir::find_codex_home;
pub(crate) use self::persistor::OAuthPersistor;
#[cfg(test)]
use self::persistor::install_tokens_in_manager;
use self::refresh_lock::RefreshCredentialLock;
pub(crate) use self::resolved_store::LoadedOAuthTokens;
pub(crate) use self::resolved_store::ResolvedOAuthCredentialStore;
@@ -77,7 +75,8 @@ use rmcp::transport::auth::AuthorizationManager;
const KEYRING_SERVICE: &str = "Codex MCP Credentials";
const MCP_OAUTH_SECRET_PREFIX: &str = "MCP_OAUTH";
const REFRESH_SKEW_MILLIS: u64 = 30_000;
// Refresh proactively so ordinary requests do not race token expiry.
const REFRESH_SKEW_MILLIS: u64 = 60_000;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct StoredOAuthTokens {
@@ -1200,15 +1199,7 @@ mod tests {
-> Result<()> {
let _env = TempCodexHome::new();
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", "scope-b"],
})))
.mount(&server)
.await;
mount_oauth_metadata(&server).await;
let store = MockKeyringStore::default();
let initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri()));
let key = super::compute_store_key(&initial_tokens.server_name, &initial_tokens.url)?;
@@ -1244,29 +1235,43 @@ mod tests {
let _env = TempCodexHome::new();
let server = MockServer::start().await;
mount_oauth_metadata(&server).await;
let refresh_started = mount_refresh_response_with_signal(
&server,
"refresh-token",
"updated-access-token",
"updated-refresh-token",
Duration::from_millis(200),
)
.await;
let store = MockKeyringStore::default();
let mut initial_tokens = sample_tokens();
initial_tokens.url = format!("{}/mcp", server.uri());
let mut updated_tokens = initial_tokens.clone();
updated_tokens
.token_response
.0
.set_access_token(AccessToken::new("updated-access-token".to_string()));
let initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri()));
super::save_oauth_tokens_with_keyring_store(
&store,
&initial_tokens.server_name,
&initial_tokens,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)?;
let manager = authorization_manager_for(&updated_tokens).await?;
let manager = authorization_manager_for(&initial_tokens).await?;
let persistor = OAuthPersistor::new(
initial_tokens.server_name.clone(),
initial_tokens.url.clone(),
manager.clone(),
ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct),
Some(initial_tokens),
Some(initial_tokens.clone()),
);
let key = super::compute_store_key(&updated_tokens.server_name, &updated_tokens.url)?;
store.set_error(&key, KeyringError::Invalid("error".into(), "save".into()));
let refresh_task = tokio::spawn({
let persistor = persistor.clone();
let store = store.clone();
async move { persistor.refresh_if_needed_with_keyring_store(&store).await }
});
let error = persistor
.persist_if_needed_with_keyring_store(&store)
.await
wait_for_signal(refresh_started).await?;
let key = super::compute_store_key(&initial_tokens.server_name, &initial_tokens.url)?;
store.set_error(&key, KeyringError::Invalid("error".into(), "save".into()));
let error = refresh_task
.await?
.expect_err("resolved keyring write should fail instead of falling back");
assert!(
@@ -1276,139 +1281,107 @@ mod tests {
"unexpected error: {error:#}"
);
assert!(!super::fallback_file_path()?.exists());
let manager_tokens = tokens_from_manager(&manager).await?;
assert_eq!(access_token(&manager_tokens), "updated-access-token");
assert_eq!(
refresh_token(&manager_tokens),
Some("updated-refresh-token".to_string())
);
// B remains authoritative in memory without retrying the failed durable write.
persistor
.persist_if_needed_with_keyring_store(&store)
.await?;
// B is now the in-process authority even though its save failed. A later operation may use
// B while it remains healthy, but we deliberately do not retry persistence here. If B
// later needs refresh, the persistor refuses to reread and replay durable A.
persistor
.refresh_if_needed_with_keyring_store(&store)
.await?;
let manager_tokens = tokens_from_manager(&manager).await?;
assert_eq!(access_token(&manager_tokens), "updated-access-token");
Ok(())
}
#[tokio::test]
async fn unpersisted_refresh_does_not_reinstall_unchanged_durable_credentials() -> Result<()> {
let _env = TempCodexHome::new();
let server = MockServer::start().await;
mount_oauth_metadata(&server).await;
Mock::given(method("POST"))
.and(path("/oauth/token"))
.respond_with(ResponseTemplate::new(500))
.expect(0)
.mount(&server)
.await;
let failing_store = MockKeyringStore::default();
let mut initial_tokens = sample_tokens();
initial_tokens.url = format!("{}/mcp", server.uri());
super::save_oauth_tokens_with_keyring_store(
&failing_store,
&initial_tokens.server_name,
&initial_tokens,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)?;
let mut unpersisted_tokens = initial_tokens.clone();
unpersisted_tokens
.token_response
.0
.set_access_token(AccessToken::new("unpersisted-access-token".to_string()));
let manager = authorization_manager_for(&unpersisted_tokens).await?;
let persistor = OAuthPersistor::new(
initial_tokens.server_name.clone(),
initial_tokens.url.clone(),
manager.clone(),
ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct),
Some(initial_tokens.clone()),
);
let key = super::compute_store_key(&initial_tokens.server_name, &initial_tokens.url)?;
failing_store.set_error(&key, KeyringError::Invalid("error".into(), "save".into()));
persistor
.persist_if_needed_with_keyring_store(&failing_store)
.await
.expect_err("the refreshed credential should remain unpersisted");
let unchanged_store = MockKeyringStore::default();
super::save_oauth_tokens_with_keyring_store(
&unchanged_store,
&initial_tokens.server_name,
&initial_tokens,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)?;
let mut subsequent_tokens = unpersisted_tokens.clone();
subsequent_tokens
.token_response
.0
.set_access_token(AccessToken::new("subsequent-access-token".to_string()));
install_tokens_in_manager(&manager, &subsequent_tokens).await?;
persistor
.persist_if_needed_locked_with_keyring_store(&unchanged_store)
.refresh_after_unauthorized_with_keyring_store(
&store,
initial_tokens.token_response.0.access_token().clone(),
)
.await?;
let manager_tokens = tokens_from_manager(&manager).await?;
assert_eq!(
manager_tokens.token_response,
subsequent_tokens.token_response
let error = persistor
.refresh_after_unauthorized_with_keyring_store(
&store,
AccessToken::new("updated-access-token".to_string()),
)
.await
.expect_err("B cannot be refreshed by rereading durable A");
assert!(
error
.to_string()
.contains("previous refresh succeeded but its credentials were not persisted"),
"unexpected error: {error:#}"
);
let stored = super::load_oauth_tokens_from_keyring(
&unchanged_store,
AuthKeyringBackendKind::Direct,
&subsequent_tokens.server_name,
&subsequent_tokens.url,
)?
.expect("the subsequent credentials should replace the unchanged stale snapshot");
assert_eq!(stored.token_response, subsequent_tokens.token_response);
server.verify().await;
Ok(())
}
#[tokio::test]
async fn resolved_secrets_write_does_not_cleanup_fallback_file() -> Result<()> {
async fn unauthorized_transaction_without_expiry_refreshes_once_across_clients() -> Result<()> {
let _env = TempCodexHome::new();
let server = MockServer::start().await;
mount_oauth_metadata(&server).await;
mount_refresh_response(
&server,
"refresh-token",
"refreshed-access-token",
"rotated-refresh-token",
)
.await;
let store = MockKeyringStore::default();
let mut initial_tokens = sample_tokens();
initial_tokens.url = format!("{}/mcp", server.uri());
super::save_oauth_tokens_to_file(&initial_tokens)?;
let fallback_path = super::fallback_file_path()?;
let fallback_before = fs::read(&fallback_path)?;
initial_tokens.expires_at = None;
initial_tokens.token_response.0.set_expires_in(None);
super::save_oauth_tokens_with_keyring_store(
&store,
&initial_tokens.server_name,
&initial_tokens,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)?;
let mut updated_tokens = initial_tokens.clone();
updated_tokens
.token_response
.0
.set_access_token(AccessToken::new("updated-access-token".to_string()));
let manager = authorization_manager_for(&updated_tokens).await?;
let persistor = OAuthPersistor::new(
let first_manager = authorization_manager_for(&initial_tokens).await?;
let second_manager = authorization_manager_for(&initial_tokens).await?;
let first = OAuthPersistor::new(
initial_tokens.server_name.clone(),
initial_tokens.url.clone(),
manager,
ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Secrets),
first_manager,
ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct),
Some(initial_tokens.clone()),
);
let second = OAuthPersistor::new(
initial_tokens.server_name.clone(),
initial_tokens.url.clone(),
second_manager.clone(),
ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct),
Some(initial_tokens.clone()),
);
persistor
.persist_if_needed_with_keyring_store(&store)
first
.refresh_after_unauthorized_with_keyring_store(&store)
.await?;
second
.refresh_after_unauthorized_with_keyring_store(&store)
.await?;
assert_eq!(fs::read(&fallback_path)?, fallback_before);
let secrets_tokens = super::load_oauth_tokens_from_keyring(
server.verify().await;
let stored = super::load_oauth_tokens_from_keyring(
&store,
AuthKeyringBackendKind::Secrets,
&updated_tokens.server_name,
&updated_tokens.url,
AuthKeyringBackendKind::Direct,
&initial_tokens.server_name,
&initial_tokens.url,
)?
.expect("updated tokens should be persisted to resolved Secrets storage");
let mut expected_tokens = updated_tokens;
expected_tokens.expires_at = secrets_tokens.expires_at;
assert_tokens_match_without_expiry(&secrets_tokens, &expected_tokens);
.expect("refreshed tokens should be persisted");
assert_eq!(access_token(&stored), "refreshed-access-token");
assert_eq!(
refresh_token(&stored),
Some("rotated-refresh-token".to_string())
);
let second_tokens = tokens_from_manager(&second_manager).await?;
assert_eq!(second_tokens.token_response, stored.token_response);
Ok(())
}
@@ -1461,19 +1434,15 @@ mod tests {
.await?;
let manager_tokens = tokens_from_manager(&manager).await?;
assert_eq!(
access_token(&manager_tokens),
"already-refreshed-access-token"
);
assert_eq!(
refresh_token(&manager_tokens),
Some("already-rotated-refresh-token".to_string())
assert_token_response_match_without_expiry(
&manager_tokens.token_response,
&latest_tokens.token_response,
);
Ok(())
}
#[tokio::test]
async fn refresh_transaction_refreshes_when_only_derived_expires_in_drifted() -> Result<()> {
async fn refresh_transaction_refreshes_in_guard_band_despite_expires_in_drift() -> Result<()> {
let _env = TempCodexHome::new();
let server = MockServer::start().await;
mount_oauth_metadata(&server).await;
@@ -1488,7 +1457,7 @@ mod tests {
let store = MockKeyringStore::default();
let mut initial_tokens = sample_tokens();
initial_tokens.url = format!("{}/mcp", server.uri());
initial_tokens.expires_at = Some(now_millis().saturating_add(5_000));
initial_tokens.expires_at = Some(now_millis().saturating_add(45_000));
initial_tokens
.token_response
.0
@@ -1598,232 +1567,6 @@ mod tests {
Ok(())
}
#[tokio::test]
async fn completed_login_prevents_delayed_post_operation_overwrite() -> Result<()> {
let _env = TempCodexHome::new();
let server = MockServer::start().await;
let store = MockKeyringStore::default();
let (initial_tokens, manager, persistor) =
persistor_after_rmcp_refresh(&store, &server).await?;
let login_tokens = tokens_with_credentials(
initial_tokens.clone(),
"login-access-token",
"login-refresh-token",
);
super::save_oauth_tokens_locked_with_keyring_store(
&store,
&login_tokens.server_name,
&login_tokens,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)
.await?;
persistor
.persist_if_needed_locked_with_keyring_store(&store)
.await?;
let stored = super::load_oauth_tokens_with_source_and_keyring_store(
&store,
&initial_tokens.server_name,
&initial_tokens.url,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)?
.expect("login tokens should remain persisted");
assert_eq!(access_token(&stored.tokens), "login-access-token");
assert_eq!(
refresh_token(&stored.tokens),
Some("login-refresh-token".to_string())
);
let manager_tokens = tokens_from_manager(&manager).await?;
assert_eq!(access_token(&manager_tokens), "login-access-token");
Ok(())
}
#[tokio::test]
async fn locked_login_save_after_refresh_still_wins() -> Result<()> {
let _env = TempCodexHome::new();
let server = MockServer::start().await;
mount_oauth_metadata(&server).await;
let refresh_started = mount_refresh_response_with_signal(
&server,
"refresh-token",
"refreshed-before-login",
"rotated-before-login",
Duration::from_millis(200),
)
.await;
let store = MockKeyringStore::default();
let initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri()));
super::save_oauth_tokens_with_keyring_store(
&store,
&initial_tokens.server_name,
&initial_tokens,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)?;
let manager = authorization_manager_for(&initial_tokens).await?;
let persistor = OAuthPersistor::new(
initial_tokens.server_name.clone(),
initial_tokens.url.clone(),
manager,
ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct),
Some(initial_tokens.clone()),
);
let refresh_task = tokio::spawn({
let persistor = persistor.clone();
let store = store.clone();
async move { persistor.refresh_if_needed_with_keyring_store(&store).await }
});
wait_for_signal(refresh_started).await?;
let mut login_tokens = sample_tokens();
login_tokens.url.clone_from(&initial_tokens.url);
login_tokens
.token_response
.0
.set_access_token(AccessToken::new("login-after-refresh-access".to_string()));
login_tokens
.token_response
.0
.set_refresh_token(Some(RefreshToken::new(
"login-after-refresh-token".to_string(),
)));
super::save_oauth_tokens_locked_with_keyring_store(
&store,
&login_tokens.server_name,
&login_tokens,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)
.await?;
refresh_task.await??;
server.verify().await;
let stored = super::load_oauth_tokens_with_source_and_keyring_store(
&store,
&initial_tokens.server_name,
&initial_tokens.url,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)?
.expect("login tokens should remain persisted");
assert_eq!(access_token(&stored.tokens), "login-after-refresh-access");
assert_eq!(
refresh_token(&stored.tokens),
Some("login-after-refresh-token".to_string())
);
Ok(())
}
#[tokio::test]
#[expect(
clippy::await_holding_invalid_type,
reason = "AuthorizationManager async access must be serialized through its mutex"
)]
async fn completed_logout_prevents_delayed_post_operation_resurrection() -> Result<()> {
let _env = TempCodexHome::new();
let server = MockServer::start().await;
let store = MockKeyringStore::default();
let (initial_tokens, manager, persistor) =
persistor_after_rmcp_refresh(&store, &server).await?;
let removed = super::delete_oauth_tokens_locked_with_keyring_store(
&store,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
&initial_tokens.server_name,
&initial_tokens.url,
)
.await?;
assert!(removed);
persistor
.persist_if_needed_locked_with_keyring_store(&store)
.await?;
let stored = super::load_oauth_tokens_with_source_and_keyring_store(
&store,
&initial_tokens.server_name,
&initial_tokens.url,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)?;
assert!(stored.is_none());
let guard = manager.lock().await;
let (_, manager_tokens) = guard.get_credentials().await?;
assert!(manager_tokens.is_none());
Ok(())
}
#[tokio::test]
async fn locked_logout_after_refresh_still_deletes() -> Result<()> {
let _env = TempCodexHome::new();
let server = MockServer::start().await;
mount_oauth_metadata(&server).await;
let refresh_started = mount_refresh_response_with_signal(
&server,
"refresh-token",
"refreshed-before-logout",
"rotated-before-logout",
Duration::from_millis(200),
)
.await;
let store = MockKeyringStore::default();
let initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri()));
super::save_oauth_tokens_with_keyring_store(
&store,
&initial_tokens.server_name,
&initial_tokens,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)?;
let manager = authorization_manager_for(&initial_tokens).await?;
let persistor = OAuthPersistor::new(
initial_tokens.server_name.clone(),
initial_tokens.url.clone(),
manager,
ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct),
Some(initial_tokens.clone()),
);
let refresh_task = tokio::spawn({
let persistor = persistor.clone();
let store = store.clone();
async move { persistor.refresh_if_needed_with_keyring_store(&store).await }
});
wait_for_signal(refresh_started).await?;
let removed = super::delete_oauth_tokens_locked_with_keyring_store(
&store,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
&initial_tokens.server_name,
&initial_tokens.url,
)
.await?;
refresh_task.await??;
server.verify().await;
assert!(removed);
let stored = super::load_oauth_tokens_with_source_and_keyring_store(
&store,
&initial_tokens.server_name,
&initial_tokens.url,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)?;
assert!(stored.is_none());
Ok(())
}
#[tokio::test]
async fn provider_refresh_timeout_permits_a_later_serialized_retry() -> Result<()> {
let _env = TempCodexHome::new();
@@ -1970,6 +1713,223 @@ mod tests {
Ok(())
}
#[tokio::test]
async fn provider_refresh_timeout_releases_lock_without_persisting() -> Result<()> {
let _env = TempCodexHome::new();
let server = MockServer::start().await;
mount_oauth_metadata(&server).await;
let refresh_started = mount_refresh_response_with_signal(
&server,
"refresh-token",
"late-access-token",
"late-refresh-token",
Duration::from_secs(1),
)
.await;
let store = MockKeyringStore::default();
let initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri()));
super::save_oauth_tokens_with_keyring_store(
&store,
&initial_tokens.server_name,
&initial_tokens,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)?;
let manager = authorization_manager_for(&initial_tokens).await?;
let persistor = OAuthPersistor::new(
initial_tokens.server_name.clone(),
initial_tokens.url.clone(),
manager,
ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct),
Some(initial_tokens.clone()),
);
let refresh_task = tokio::spawn({
let persistor = persistor.clone();
let store = store.clone();
async move {
persistor
.refresh_if_needed_with_keyring_store_and_timeout(
&store,
Duration::from_millis(200),
)
.await
}
});
wait_for_signal(refresh_started).await?;
let error = refresh_task
.await?
.expect_err("delayed provider response should time out");
assert_eq!(
error.to_string(),
"timed out after 200ms refreshing OAuth tokens for server test-server"
);
let store_key = super::compute_store_key(&initial_tokens.server_name, &initial_tokens.url)?;
let _lock =
RefreshCredentialLock::acquire_with_timeout(&store_key, Duration::from_millis(100))
.await?;
let stored = super::load_oauth_tokens_with_source_and_keyring_store(
&store,
&initial_tokens.server_name,
&initial_tokens.url,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)?
.expect("original tokens should remain persisted");
assert_eq!(access_token(&stored.tokens), "access-token");
assert_eq!(
refresh_token(&stored.tokens),
Some("refresh-token".to_string())
);
server.verify().await;
Ok(())
}
#[tokio::test]
async fn locked_login_save_after_refresh_still_wins() -> Result<()> {
let _env = TempCodexHome::new();
let server = MockServer::start().await;
mount_oauth_metadata(&server).await;
let refresh_started = mount_refresh_response_with_signal(
&server,
"refresh-token",
"refreshed-before-login",
"rotated-before-login",
Duration::from_millis(200),
)
.await;
let store = MockKeyringStore::default();
let initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri()));
super::save_oauth_tokens_with_keyring_store(
&store,
&initial_tokens.server_name,
&initial_tokens,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)?;
let manager = authorization_manager_for(&initial_tokens).await?;
let persistor = OAuthPersistor::new(
initial_tokens.server_name.clone(),
initial_tokens.url.clone(),
manager,
ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct),
Some(initial_tokens.clone()),
);
let refresh_task = tokio::spawn({
let persistor = persistor.clone();
let store = store.clone();
async move { persistor.refresh_if_needed_with_keyring_store(&store).await }
});
wait_for_signal(refresh_started).await?;
let mut login_tokens = sample_tokens();
login_tokens.url.clone_from(&initial_tokens.url);
login_tokens
.token_response
.0
.set_access_token(AccessToken::new("login-after-refresh-access".to_string()));
login_tokens
.token_response
.0
.set_refresh_token(Some(RefreshToken::new(
"login-after-refresh-token".to_string(),
)));
super::save_oauth_tokens_locked_with_keyring_store(
&store,
&login_tokens.server_name,
&login_tokens,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)
.await?;
refresh_task.await??;
server.verify().await;
let stored = super::load_oauth_tokens_with_source_and_keyring_store(
&store,
&initial_tokens.server_name,
&initial_tokens.url,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)?
.expect("login tokens should remain persisted");
assert_eq!(access_token(&stored.tokens), "login-after-refresh-access");
assert_eq!(
refresh_token(&stored.tokens),
Some("login-after-refresh-token".to_string())
);
Ok(())
}
#[tokio::test]
async fn locked_logout_after_refresh_still_deletes() -> Result<()> {
let _env = TempCodexHome::new();
let server = MockServer::start().await;
mount_oauth_metadata(&server).await;
let refresh_started = mount_refresh_response_with_signal(
&server,
"refresh-token",
"refreshed-before-logout",
"rotated-before-logout",
Duration::from_millis(200),
)
.await;
let store = MockKeyringStore::default();
let initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri()));
super::save_oauth_tokens_with_keyring_store(
&store,
&initial_tokens.server_name,
&initial_tokens,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)?;
let manager = authorization_manager_for(&initial_tokens).await?;
let persistor = OAuthPersistor::new(
initial_tokens.server_name.clone(),
initial_tokens.url.clone(),
manager,
ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct),
Some(initial_tokens.clone()),
);
let refresh_task = tokio::spawn({
let persistor = persistor.clone();
let store = store.clone();
async move { persistor.refresh_if_needed_with_keyring_store(&store).await }
});
wait_for_signal(refresh_started).await?;
let removed = super::delete_oauth_tokens_locked_with_keyring_store(
&store,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
&initial_tokens.server_name,
&initial_tokens.url,
)
.await?;
refresh_task.await??;
server.verify().await;
assert!(removed);
let stored = super::load_oauth_tokens_with_source_and_keyring_store(
&store,
&initial_tokens.server_name,
&initial_tokens.url,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)?;
assert!(stored.is_none());
Ok(())
}
#[test]
fn save_oauth_tokens_with_secrets_backend_falls_back_to_file_when_keyring_fails() -> Result<()>
{
@@ -2261,57 +2221,6 @@ mod tests {
);
}
async fn persistor_after_rmcp_refresh(
store: &MockKeyringStore,
server: &MockServer,
) -> Result<(
StoredOAuthTokens,
Arc<TokioMutex<AuthorizationManager>>,
OAuthPersistor,
)> {
mount_oauth_metadata(server).await;
let mut initial_tokens = sample_tokens();
initial_tokens.url = format!("{}/mcp", server.uri());
super::save_oauth_tokens_with_keyring_store(
store,
&initial_tokens.server_name,
&initial_tokens,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)?;
let manager = authorization_manager_for(&initial_tokens).await?;
let persistor = OAuthPersistor::new(
initial_tokens.server_name.clone(),
initial_tokens.url.clone(),
manager.clone(),
ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct),
Some(initial_tokens.clone()),
);
let rmcp_tokens = tokens_with_credentials(
initial_tokens.clone(),
"rmcp-access-token",
"rmcp-refresh-token",
);
install_tokens_in_manager(&manager, &rmcp_tokens).await?;
Ok((initial_tokens, manager, persistor))
}
fn tokens_with_credentials(
mut tokens: StoredOAuthTokens,
access_token: &str,
refresh_token: &str,
) -> StoredOAuthTokens {
tokens
.token_response
.0
.set_access_token(AccessToken::new(access_token.to_string()));
tokens
.token_response
.0
.set_refresh_token(Some(RefreshToken::new(refresh_token.to_string())));
tokens
}
async fn authorization_manager_for(
tokens: &StoredOAuthTokens,
) -> Result<Arc<TokioMutex<AuthorizationManager>>> {
codex-rs/rmcp-client/src/oauth/persistor.rs
+190 -260
View File
@@ -8,7 +8,6 @@ use std::time::UNIX_EPOCH;
use anyhow::Context;
use anyhow::Result;
use codex_config::types::AuthKeyringBackendKind;
use codex_keyring_store::DefaultKeyringStore;
use codex_keyring_store::KeyringStore;
use oauth2::TokenResponse;
@@ -26,10 +25,6 @@ use super::ResolvedOAuthCredentialStore;
use super::StoredOAuthTokens;
use super::WrappedOAuthTokenResponse;
use super::compute_expires_at_millis;
use super::compute_store_key;
use super::delete_oauth_tokens_from_direct_keyring;
use super::delete_oauth_tokens_from_file;
use super::delete_oauth_tokens_from_secrets_keyring;
use super::load_oauth_tokens_from_file;
use super::load_oauth_tokens_from_keyring;
use super::refresh_lock::RefreshCredentialLock;
@@ -43,7 +38,6 @@ const REFRESH_REQUEST_TIMEOUT: Duration = Duration::from_secs(45);
pub(crate) struct OAuthPersistor {
inner: Arc<OAuthPersistorInner>,
}
struct OAuthPersistorInner {
server_name: String,
url: String,
@@ -111,79 +105,6 @@ impl OAuthPersistor {
}
}
/// Persists the latest stored credentials if they have changed.
/// Deletes the credentials if they are no longer present.
pub(crate) async fn persist_if_needed(&self) -> Result<()> {
self.persist_if_needed_locked_with_keyring_store(&DefaultKeyringStore)
.await
}
pub(super) async fn persist_if_needed_locked_with_keyring_store<
K: KeyringStore + Clone + 'static,
>(
&self,
keyring_store: &K,
) -> Result<()> {
let (snapshot, unpersisted_refresh) = {
let state = self.inner.credential_state.lock().await;
(state.current.clone(), state.unpersisted_refresh.clone())
};
let (client_id, current_credentials) = self.manager_credentials().await?;
let manager_changed = match (&snapshot, current_credentials.as_ref()) {
(Some(previous), Some(current)) => {
previous.client_id != client_id
|| previous.token_response != WrappedOAuthTokenResponse(current.clone())
}
(None, None) => false,
(Some(_), None) | (None, Some(_)) => true,
};
if !manager_changed {
return Ok(());
}
let _lock =
RefreshCredentialLock::acquire_for_server(&self.inner.server_name, &self.inner.url)
.await?;
let latest = self.load_resolved_credentials(keyring_store)?;
let latest_matches_snapshot = durable_credentials_match_snapshot(&latest, &snapshot)
|| unpersisted_refresh.as_ref().is_some_and(|pending| {
// The failed write never changed durable authority. If storage still contains the
// last known durable snapshot, persist the manager's newer result instead of
// adopting and replaying that stale snapshot.
durable_credentials_match_snapshot(&latest, &pending.previously_persisted)
});
if !latest_matches_snapshot {
// A completed login or logout is authoritative over tokens refreshed inside an RMCP
// operation. Adopt that state instead of allowing delayed post-operation persistence
// to overwrite the login or resurrect the logout.
match latest {
Some(latest) => self.adopt_credentials(latest).await?,
None => {
self.clear_manager_credentials().await;
let mut state = self.inner.credential_state.lock().await;
state.current = None;
state.unpersisted_refresh = None;
}
}
return Ok(());
}
self.persist_if_needed_with_keyring_store(keyring_store)
.await
}
#[expect(
clippy::await_holding_invalid_type,
reason = "AuthorizationManager async access must be serialized through its mutex"
)]
async fn manager_credentials(&self) -> Result<(String, Option<OAuthTokenResponse>)> {
let manager = self.inner.authorization_manager.clone();
let guard = manager.lock().await;
guard.get_credentials().await.map_err(Into::into)
}
fn load_resolved_credentials<K: KeyringStore + Clone + 'static>(
&self,
keyring_store: &K,
@@ -207,128 +128,16 @@ impl OAuthPersistor {
}
}
#[expect(
clippy::await_holding_invalid_type,
reason = "AuthorizationManager async access must be serialized through its mutex"
)]
pub(super) async fn persist_if_needed_with_keyring_store<K: KeyringStore + Clone + 'static>(
&self,
keyring_store: &K,
) -> 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 state = self.inner.credential_state.lock().await;
let new_token_response = WrappedOAuthTokenResponse(credentials.clone());
let same_token = state
.current
.as_ref()
.map(|prev| prev.token_response == new_token_response)
.unwrap_or(false);
let expires_at = if same_token {
state.current.as_ref().and_then(|prev| prev.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 state.current.as_ref() != Some(&stored) {
// The provider may already have consumed the old rotating refresh token. Make
// B authoritative in this process before the fallible save so a later public
// operation cannot reinstall A from the last snapshot.
// Preserve the last snapshot known to be durable across repeated in-memory
// changes. Using the immediately preceding in-memory token here could make an
// unchanged stale store look like an external login or refresh.
let previously_persisted = state
.unpersisted_refresh
.take()
.map(|pending| pending.previously_persisted)
.unwrap_or_else(|| state.current.clone());
state.current = Some(stored.clone());
state.unpersisted_refresh = Some(UnpersistedRefresh {
previously_persisted,
});
debug!("persisting refreshed MCP OAuth credentials to the resolved store");
let persistence_started_at = Instant::now();
let persistence_result = match self.inner.credential_store {
ResolvedOAuthCredentialStore::File => save_oauth_tokens_to_file(&stored),
ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind) => {
save_oauth_tokens_with_keyring(
keyring_store,
keyring_backend_kind,
&self.inner.server_name,
&stored,
)
}
};
if let Err(error) = persistence_result {
warn!(
persistence_elapsed_ms = persistence_started_at.elapsed().as_millis(),
error = %error,
"failed to persist refreshed MCP OAuth credentials; retaining them as the in-process authority without retrying persistence"
);
return Err(error);
}
state.unpersisted_refresh = None;
debug!(
persistence_elapsed_ms = persistence_started_at.elapsed().as_millis(),
"persisted refreshed MCP OAuth credentials"
);
}
}
None => {
let mut state = self.inner.credential_state.lock().await;
if state.current.take().is_some()
&& let Err(error) = match self.inner.credential_store {
ResolvedOAuthCredentialStore::File => {
let key = compute_store_key(&self.inner.server_name, &self.inner.url)?;
delete_oauth_tokens_from_file(&key).map(|_| ())
}
ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct) => {
delete_oauth_tokens_from_direct_keyring(
keyring_store,
&self.inner.server_name,
&self.inner.url,
)
.map(|_| ())
}
ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Secrets) => {
delete_oauth_tokens_from_secrets_keyring(
keyring_store,
&self.inner.server_name,
&self.inner.url,
)
.map(|_| ())
}
}
{
warn!(
"failed to remove OAuth tokens for server {}: {error}",
self.inner.server_name
);
}
state.unpersisted_refresh = None;
}
}
Ok(())
}
pub(crate) async fn refresh_if_needed(&self) -> Result<()> {
self.refresh_if_needed_with_keyring_store(&DefaultKeyringStore)
.await
}
pub(crate) async fn refresh_after_unauthorized(&self) -> Result<()> {
self.refresh_after_unauthorized_with_keyring_store(&DefaultKeyringStore)
.await
}
pub(super) async fn refresh_if_needed_with_keyring_store<K: KeyringStore + Clone + 'static>(
&self,
keyring_store: &K,
@@ -356,13 +165,32 @@ impl OAuthPersistor {
return Ok(());
}
self.run_owned_refresh_transaction(keyring_store.clone(), refresh_request_timeout)
.await
self.run_owned_refresh_transaction(
keyring_store.clone(),
RefreshReason::Expiry,
refresh_request_timeout,
)
.await
}
pub(super) async fn refresh_after_unauthorized_with_keyring_store<
K: KeyringStore + Clone + 'static,
>(
&self,
keyring_store: &K,
) -> Result<()> {
self.run_owned_refresh_transaction(
keyring_store.clone(),
RefreshReason::Unauthorized,
REFRESH_REQUEST_TIMEOUT,
)
.await
}
async fn run_owned_refresh_transaction<K: KeyringStore + Clone + 'static>(
&self,
keyring_store: K,
reason: RefreshReason,
refresh_request_timeout: Duration,
) -> Result<()> {
let persistor = self.clone();
@@ -378,7 +206,7 @@ impl OAuthPersistor {
// token-family-revocation risk rather than holding the lock indefinitely.
tokio::spawn(async move {
persistor
.refresh_transaction(&keyring_store, refresh_request_timeout)
.refresh_transaction(&keyring_store, reason, refresh_request_timeout)
.await
})
.await
@@ -392,16 +220,27 @@ impl OAuthPersistor {
#[tracing::instrument(
level = "debug",
skip_all,
fields(server_name = %self.inner.server_name),
fields(
server_name = %self.inner.server_name,
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<()> {
let transaction_started_at = Instant::now();
let lock_started_at = Instant::now();
let local_access_token = {
let state = self.inner.credential_state.lock().await;
state
.current
.as_ref()
.map(|tokens| tokens.token_response.0.access_token().secret().to_string())
};
debug!("waiting for the MCP OAuth credential transaction lock");
let _lock =
RefreshCredentialLock::acquire_for_server(&self.inner.server_name, &self.inner.url)
@@ -414,10 +253,11 @@ impl OAuthPersistor {
let unpersisted_refresh = {
let state = self.inner.credential_state.lock().await;
if let Some(unpersisted_refresh) = state.unpersisted_refresh.as_ref() {
if state
.current
.as_ref()
.is_some_and(|tokens| !token_needs_refresh(tokens.expires_at))
if matches!(reason, RefreshReason::Expiry)
&& state
.current
.as_ref()
.is_some_and(|tokens| !token_needs_refresh(tokens.expires_at))
{
debug!(
"using the memory-authoritative MCP OAuth credentials from a refresh whose persistence failed"
@@ -464,67 +304,114 @@ impl OAuthPersistor {
);
};
if !token_needs_refresh(latest.expires_at) {
let latest_access_token = latest.token_response.0.access_token().secret();
// Expiry refresh can adopt any reread that is now healthy. A 401 is different: an
// unexpired token is still rejected, so adopt only when another actor has already changed
// the access token; otherwise force one serialized provider refresh.
let should_adopt = !token_needs_refresh(latest.expires_at)
&& match reason {
RefreshReason::Expiry => true,
RefreshReason::Unauthorized => {
local_access_token.as_deref() != Some(latest_access_token)
}
};
if should_adopt {
self.adopt_credentials(latest).await?;
return Ok(());
}
self.adopt_credentials(latest).await?;
let manager = self.inner.authorization_manager.clone();
let mut guard = manager.lock().await;
install_tokens_in_manager_guard(&mut guard, &latest)
.await
.context("failed to stage OAuth credentials for refresh")?;
let provider_started_at = Instant::now();
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 {
Ok(Ok(token_response)) => {
debug!(
provider_elapsed_ms = provider_started_at.elapsed().as_millis(),
"received refreshed MCP OAuth credentials from the provider"
);
refreshed_tokens(token_response, &latest, &self.inner)
}
Ok(Err(error)) => {
warn!(
provider_elapsed_ms = provider_started_at.elapsed().as_millis(),
error = %error,
"MCP OAuth provider refresh failed"
);
return Err(error).with_context(|| {
format!(
"failed to refresh OAuth tokens for server {}",
self.inner.server_name
)
});
}
Err(_) => {
warn!(
provider_elapsed_ms = provider_started_at.elapsed().as_millis(),
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!(
"timed out after {refresh_request_timeout:?} refreshing OAuth tokens for server {}",
self.inner.server_name
);
}
};
install_tokens_in_manager_guard(&mut guard, &refreshed)
.await
.context("failed to install refreshed OAuth credentials")?;
drop(guard);
{
let manager = self.inner.authorization_manager.clone();
let guard = manager.lock().await;
let provider_started_at = Instant::now();
debug!(
timeout_ms = refresh_request_timeout.as_millis(),
"requesting refreshed MCP OAuth credentials from the provider"
);
match timeout(refresh_request_timeout, guard.refresh_token()).await {
Ok(Ok(_token_response)) => {
debug!(
provider_elapsed_ms = provider_started_at.elapsed().as_millis(),
"received refreshed MCP OAuth credentials from the provider"
);
}
Ok(Err(error)) => {
warn!(
provider_elapsed_ms = provider_started_at.elapsed().as_millis(),
error = %error,
"MCP OAuth provider refresh failed"
);
return Err(error).with_context(|| {
format!(
"failed to refresh OAuth tokens for server {}",
self.inner.server_name
)
});
}
Err(_) => {
warn!(
provider_elapsed_ms = provider_started_at.elapsed().as_millis(),
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!(
"timed out after {refresh_request_timeout:?} refreshing OAuth tokens for server {}",
self.inner.server_name
);
}
let mut state = self.inner.credential_state.lock().await;
state.current = Some(refreshed.clone());
state.unpersisted_refresh = Some(UnpersistedRefresh {
previously_persisted: Some(latest.clone()),
});
}
debug!(
"installed refreshed MCP OAuth credentials in memory and marked persistence pending"
);
// Once the provider rotates a refresh token, persistence must complete even if the caller's
// deadline expires in the meantime. Returning early here would lose the only usable token.
// Refresh persistence stays on the source resolved at client startup. In particular, a
// keyring failure must surface instead of writing the rotated token to fallback File.
debug!("persisting refreshed MCP OAuth credentials to the resolved store");
let persistence_started_at = Instant::now();
let persistence_result = match self.inner.credential_store {
ResolvedOAuthCredentialStore::File => save_oauth_tokens_to_file(&refreshed),
ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind) => {
save_oauth_tokens_with_keyring(
keyring_store,
keyring_backend_kind,
&self.inner.server_name,
&refreshed,
)
}
}
// Once the provider returns a rotated token, persistence must finish before the credential
// lock is released. In particular, caller startup deadlines must not cancel this step.
let result = self
.persist_if_needed_with_keyring_store(keyring_store)
.await;
if result.is_ok() {
debug!(
};
if let Err(error) = persistence_result {
warn!(
persistence_elapsed_ms = persistence_started_at.elapsed().as_millis(),
transaction_elapsed_ms = transaction_started_at.elapsed().as_millis(),
"completed the MCP OAuth refresh transaction"
error = %error,
"failed to persist refreshed MCP OAuth credentials; retaining them as the in-process authority without retrying persistence"
);
return Err(error);
}
result
let mut state = self.inner.credential_state.lock().await;
state.unpersisted_refresh = None;
debug!(
persistence_elapsed_ms = persistence_started_at.elapsed().as_millis(),
transaction_elapsed_ms = transaction_started_at.elapsed().as_millis(),
"persisted refreshed MCP OAuth credentials and completed the transaction"
);
Ok(())
}
async fn adopt_credentials(&self, tokens: StoredOAuthTokens) -> Result<()> {
@@ -542,13 +429,37 @@ impl OAuthPersistor {
}
}
#[derive(Clone, Copy)]
enum RefreshReason {
Expiry,
Unauthorized,
}
impl RefreshReason {
fn as_str(&self) -> &'static str {
match self {
Self::Expiry => "expiry",
Self::Unauthorized => "unauthorized",
}
}
}
#[expect(
clippy::await_holding_invalid_type,
reason = "AuthorizationManager async access must be serialized through its mutex"
)]
pub(super) async fn install_tokens_in_manager(
async fn install_tokens_in_manager(
authorization_manager: &Arc<Mutex<AuthorizationManager>>,
tokens: &StoredOAuthTokens,
) -> Result<()> {
let manager = authorization_manager.clone();
let mut guard = manager.lock().await;
install_tokens_in_manager_guard(&mut guard, tokens).await
}
async fn install_tokens_in_manager_guard(
authorization_manager: &mut AuthorizationManager,
tokens: &StoredOAuthTokens,
) -> Result<()> {
let store = InMemoryCredentialStore::new();
store
@@ -556,13 +467,11 @@ pub(super) async fn install_tokens_in_manager(
.await
.context("failed to stage OAuth tokens for authorization manager")?;
let manager = authorization_manager.clone();
let mut guard = manager.lock().await;
guard.set_credential_store(store);
authorization_manager.set_credential_store(store);
// TODO(stevenlee): RMCP's `initialize_from_store` updates the credential store and client ID
// but not its private `current_scopes`. Credential adoption can therefore leave scope-upgrade
// state stale until RMCP exposes an adoption API that synchronizes both.
guard
authorization_manager
.initialize_from_store()
.await
.context("failed to adopt refreshed OAuth tokens")?;
@@ -587,3 +496,24 @@ fn stored_credentials_from_tokens(tokens: &StoredOAuthTokens) -> StoredCredentia
token_received_at,
)
}
fn refreshed_tokens(
mut token_response: OAuthTokenResponse,
previous: &StoredOAuthTokens,
inner: &OAuthPersistorInner,
) -> StoredOAuthTokens {
if token_response.refresh_token().is_none() {
token_response.set_refresh_token(previous.token_response.0.refresh_token().cloned());
}
if token_response.scopes().is_none() {
token_response.set_scopes(previous.token_response.0.scopes().cloned());
}
let expires_at = compute_expires_at_millis(&token_response);
StoredOAuthTokens {
server_name: inner.server_name.clone(),
url: inner.url.clone(),
client_id: previous.client_id.clone(),
token_response: WrappedOAuthTokenResponse(token_response),
expires_at,
}
}
codex-rs/rmcp-client/src/oauth/refresh_lock.rs
+8 -1
View File
@@ -38,8 +38,15 @@ impl RefreshCredentialLock {
}
pub(super) async fn acquire(store_key: &str) -> Result<Self> {
Self::acquire_with_timeout(store_key, REFRESH_LOCK_ACQUIRE_TIMEOUT).await
}
pub(super) async fn acquire_with_timeout(
store_key: &str,
acquire_timeout: Duration,
) -> Result<Self> {
let codex_home = find_codex_home()?;
Self::acquire_in(&codex_home, store_key, REFRESH_LOCK_ACQUIRE_TIMEOUT).await
Self::acquire_in(&codex_home, store_key, acquire_timeout).await
}
async fn acquire_in(
codex-rs/rmcp-client/src/rmcp_client.rs
+75 -62
View File
@@ -453,7 +453,7 @@ impl RmcpClient {
let mut initialize_deadline = timeout.map(|duration| Instant::now() + duration);
let (service, oauth_persistor) = self
.connect_pending_transport_with_initialize_retries(
.connect_pending_transport_with_oauth_recovery(
pending_transport,
client_service.clone(),
timeout,
@@ -486,12 +486,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)
}
@@ -507,7 +501,6 @@ impl RmcpClient {
async move { service.list_tools(params).await }.boxed()
})
.await?;
self.persist_oauth_tokens().await;
Ok(result)
}
@@ -542,7 +535,6 @@ impl RmcpClient {
})
})
.collect::<Result<Vec<_>>>()?;
self.persist_oauth_tokens().await;
Ok(ListToolsWithConnectorIdResult {
next_cursor: result.next_cursor,
tools,
@@ -569,7 +561,6 @@ impl RmcpClient {
async move { service.list_resources(params).await }.boxed()
})
.await?;
self.persist_oauth_tokens().await;
Ok(result)
}
@@ -585,7 +576,6 @@ impl RmcpClient {
async move { service.list_resource_templates(params).await }.boxed()
})
.await?;
self.persist_oauth_tokens().await;
Ok(result)
}
@@ -601,7 +591,6 @@ impl RmcpClient {
async move { service.read_resource(params).await }.boxed()
})
.await?;
self.persist_oauth_tokens().await;
Ok(result)
}
@@ -659,7 +648,6 @@ impl RmcpClient {
.boxed()
})
.await?;
self.persist_oauth_tokens().await;
Ok(result)
}
@@ -689,7 +677,6 @@ impl RmcpClient {
},
)
.await?;
self.persist_oauth_tokens().await;
Ok(())
}
@@ -712,7 +699,6 @@ impl RmcpClient {
.boxed()
})
.await?;
self.persist_oauth_tokens().await;
Ok(response)
}
@@ -752,16 +738,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}");
}
}
async fn refresh_oauth_if_needed(&self) -> Result<()> {
if let Some(runtime) = self.oauth_persistor().await {
runtime.refresh_if_needed().await?;
@@ -968,19 +944,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))
}
@@ -995,38 +959,69 @@ impl RmcpClient {
F: Fn(Arc<RunningService<RoleClient, ElicitationClientService>>) -> Fut,
Fut: std::future::Future<Output = std::result::Result<T, rmcp::service::ServiceError>>,
{
let deadline = timeout.map(|duration| Instant::now() + duration);
let service = self.service().await?;
match Self::run_service_operation_with_transient_retries(
let mut result = Self::run_service_operation_with_transient_retries(
Arc::clone(&service),
label,
timeout,
deadline,
self.elicitation_pause_state.clone(),
&operation,
)
.await
.await;
if result
.as_ref()
.is_err_and(Self::is_unauthorized_operation_error)
&& let Some(oauth_persistor) = self.oauth_persistor().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)
// OAuth recovery is intentionally one-shot. The caller deadline gates whether recovery
// starts and whether the retry runs, but it must not cancel an in-flight provider
// refresh before a rotated token is persisted.
remaining_operation_timeout(label, timeout, deadline)?;
let refresh_result = oauth_persistor.refresh_after_unauthorized().await;
if let Err(error) = refresh_result {
if let Err(timeout_error) = remaining_operation_timeout(label, timeout, deadline) {
return Err(timeout_error.into());
}
return Err(error);
}
Err(error) => Err(error.into()),
result = Self::run_service_operation_with_transient_retries(
Arc::clone(&service),
label,
timeout,
deadline,
self.elicitation_pause_state.clone(),
&operation,
)
.await;
}
if result.as_ref().is_err_and(Self::is_session_expired_404) {
// Session recovery remains one-shot and runs after the optional OAuth retry, so a 401
// followed by the old session's 404 still reconstructs the transport before retrying.
self.reinitialize_after_session_expiry(&service).await?;
let recovered_service = self.service().await?;
result = Self::run_service_operation_with_transient_retries(
recovered_service,
label,
timeout,
deadline,
self.elicitation_pause_state.clone(),
&operation,
)
.await;
}
result.map_err(Into::into)
}
async fn run_service_operation_with_transient_retries<T, F, Fut>(
service: Arc<RunningService<RoleClient, ElicitationClientService>>,
label: &str,
timeout: Option<Duration>,
retry_deadline: Option<Instant>,
pause_state: ElicitationPauseState,
operation: &F,
) -> std::result::Result<T, ClientOperationError>
@@ -1034,7 +1029,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()
@@ -1140,6 +1134,31 @@ impl RmcpClient {
})
}
fn is_unauthorized_operation_error(error: &ClientOperationError) -> bool {
let ClientOperationError::Service(rmcp::service::ServiceError::TransportSend(error)) =
error
else {
return false;
};
error
.error
.downcast_ref::<StreamableHttpError<StreamableHttpClientAdapterError>>()
.is_some_and(Self::is_unauthorized_streamable_http_error)
}
pub(super) fn is_unauthorized_streamable_http_error(
error: &StreamableHttpError<StreamableHttpClientAdapterError>,
) -> bool {
match error {
StreamableHttpError::AuthRequired(_) => true,
StreamableHttpError::UnexpectedServerResponse(message) => {
message.starts_with("HTTP 401")
}
_ => false,
}
}
async fn reinitialize_after_session_expiry(
&self,
failed_service: &Arc<RunningService<RoleClient, ElicitationClientService>>,
@@ -1177,7 +1196,7 @@ impl RmcpClient {
.timeout
.map(|duration| Instant::now() + duration);
let (service, oauth_persistor) = self
.connect_pending_transport_with_initialize_retries(
.connect_pending_transport_with_oauth_recovery(
pending_transport,
initialize_context.client_service,
initialize_context.timeout,
@@ -1196,12 +1215,6 @@ 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(())
}
}
codex-rs/rmcp-client/src/streamable_http_retry.rs
+103 -1
View File
@@ -22,8 +22,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,
@@ -32,6 +37,69 @@ impl RmcpClient {
) -> Result<(
Arc<RunningService<RoleClient, ElicitationClientService>>,
Option<OAuthPersistor>,
)> {
let mut attempt_context = InitializeAttemptContext::default();
match self
.connect_pending_transport_with_initialize_retries(
initial_transport,
client_service.clone(),
timeout,
initialize_deadline,
&mut attempt_context,
)
.await
{
Ok(result) => Ok(result),
Err(error) if Self::is_unauthorized_initialize_error(&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().await;
if let Some(deadline) = initialize_deadline.as_mut() {
*deadline += refresh_started_at.elapsed();
}
refresh_result?;
let remaining = remaining_initialize_timeout(timeout, *initialize_deadline)?;
let transport = match remaining {
Some(remaining) => time::timeout(
remaining,
Self::create_pending_transport(&self.transport_recipe),
)
.await
.map_err(|_| initialize_timeout_error(timeout, remaining))??,
None => Self::create_pending_transport(&self.transport_recipe).await?,
};
let mut retry_context = InitializeAttemptContext::default();
self.connect_pending_transport_with_initialize_retries(
transport,
client_service,
timeout,
initialize_deadline,
&mut retry_context,
)
.await
}
Err(error) => Err(error),
}
}
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,
@@ -62,6 +130,17 @@ impl RmcpClient {
}
}
};
// Keep the persistor paired with the transport attempt that returned 401. Rebuilt
// transports reuse the recipe's lifecycle-pinned credential source, and this pairing
// also keeps the authorization manager and snapshot aligned with the failed attempt.
attempt_context.oauth_persistor = match &transport {
PendingTransport::StreamableHttpWithOAuth {
oauth_persistor, ..
} => Some(oauth_persistor.clone()),
PendingTransport::InProcess { .. }
| PendingTransport::Stdio { .. }
| PendingTransport::StreamableHttp { .. } => None,
};
match Self::connect_pending_transport(
transport,
client_service.clone(),
@@ -108,6 +187,29 @@ impl RmcpClient {
})
}
fn is_unauthorized_initialize_error(error: &anyhow::Error) -> bool {
error.chain().any(|source| {
source
.downcast_ref::<HandshakeError>()
.is_some_and(|error| Self::is_unauthorized_client_initialize_error(&error.source))
|| source
.downcast_ref::<rmcp::service::ClientInitializeError>()
.is_some_and(Self::is_unauthorized_client_initialize_error)
})
}
fn is_unauthorized_client_initialize_error(
error: &rmcp::service::ClientInitializeError,
) -> bool {
match error {
rmcp::service::ClientInitializeError::TransportError { error, .. } => error
.error
.downcast_ref::<StreamableHttpError<StreamableHttpClientAdapterError>>()
.is_some_and(Self::is_unauthorized_streamable_http_error),
_ => false,
}
}
fn is_retryable_client_initialize_error(error: &rmcp::service::ClientInitializeError) -> bool {
match error {
rmcp::service::ClientInitializeError::TransportError { error, context }