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

fix(rmcp-client): preserve rotated OAuth credentials

Browse Source
This commit is contained in:
Adam Perry
2026-06-05 08:00:23 +00:00
parent cb1a7928ad
commit 935794b236
3 changed files with 334 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/rmcp-client/src/oauth.rs
+44 -19
View File
@@ -399,6 +399,12 @@ struct OAuthPersistorInner {
credential_store: InMemoryCredentialStore,
store_mode: OAuthCredentialsStoreMode,
last_credentials: Mutex<Option<StoredOAuthTokens>>,
refresh_persistence_failure: Mutex<Option<RefreshPersistenceFailure>>,
}
struct RefreshPersistenceFailure {
message: String,
_refresh_locks: Vec<fs::File>,
}
impl OAuthPersistor {
@@ -418,6 +424,7 @@ impl OAuthPersistor {
credential_store,
store_mode,
last_credentials: Mutex::new(initial_credentials),
refresh_persistence_failure: Mutex::new(None),
}),
}
}
@@ -425,20 +432,10 @@ 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_with_refresh_locks(Vec::new()).await
}
async fn persist_if_needed_with_refresh_locks(
&self,
refresh_locks: Vec<fs::File>,
) -> Result<()> {
let persistor = self.clone();
tokio::spawn(async move {
let _refresh_locks = refresh_locks;
persistor.persist_if_needed_inner().await
})
.await
.context("OAuth token persistence task failed")?
tokio::spawn(async move { persistor.persist_if_needed_inner().await })
.await
.context("OAuth token persistence task failed")?
}
#[expect(
@@ -508,6 +505,17 @@ impl OAuthPersistor {
reason = "AuthorizationManager async access must be serialized through its mutex"
)]
pub(crate) async fn refresh_if_needed(&self) -> Result<()> {
let previous_failure = self
.inner
.refresh_persistence_failure
.lock()
.await
.as_ref()
.map(|failure| failure.message.clone());
if let Some(message) = previous_failure {
return Err(anyhow!(message));
}
let initial_expires_at = {
let guard = self.inner.last_credentials.lock().await;
guard.as_ref().and_then(|tokens| tokens.expires_at)
@@ -568,8 +576,9 @@ impl OAuthPersistor {
return Ok(());
}
{
let manager = self.inner.authorization_manager.clone();
let persistor = self.clone();
tokio::spawn(async move {
let manager = persistor.inner.authorization_manager.clone();
let guard = manager.lock().await;
match guard.refresh_token().await {
Ok(_) => {}
@@ -582,15 +591,31 @@ impl OAuthPersistor {
return Err(error).with_context(|| {
format!(
"failed to refresh OAuth tokens for server {}",
self.inner.server_name
persistor.inner.server_name
)
});
}
}
}
drop(guard);
self.persist_if_needed_with_refresh_locks(refresh_locks)
.await
if let Err(error) = persistor.persist_if_needed_inner().await {
let message = format!(
"OAuth refresh succeeded but failed to persist rotated credentials for server {}: {error:#}",
persistor.inner.server_name
);
warn!("{message}");
*persistor.inner.refresh_persistence_failure.lock().await =
Some(RefreshPersistenceFailure {
message: message.clone(),
_refresh_locks: refresh_locks,
});
return Err(anyhow!(message));
}
Ok(())
})
.await
.context("OAuth refresh task failed")?
}
}
codex-rs/rmcp-client/src/rmcp_client.rs
+66 -12
View File
@@ -225,6 +225,8 @@ where
enum ClientOperationError {
#[error(transparent)]
Service(#[from] rmcp::service::ServiceError),
#[error(transparent)]
OAuth(anyhow::Error),
#[error("timed out awaiting {label} after {duration:?}")]
Timeout { label: String, duration: Duration },
}
@@ -863,10 +865,11 @@ impl RmcpClient {
let service = transport
.await
.map_err(|err| anyhow!("handshaking with MCP server failed: {err}"))?;
if let Some(runtime) = &oauth_persistor
&& let Err(error) = runtime.persist_if_needed().await
{
warn!("failed to persist OAuth tokens after initialize: {error}");
if let Some(runtime) = oauth_persistor.clone() {
Self::persist_oauth_tokens_in_background(
runtime,
"failed to persist OAuth tokens after initialize",
);
}
Ok((Arc::new(service), oauth_persistor))
};
@@ -934,18 +937,17 @@ impl RmcpClient {
Fut: std::future::Future<Output = std::result::Result<T, rmcp::service::ServiceError>>,
{
let refresh_and_operation = async move {
if let Some(runtime) = &oauth_persistor
&& let Err(error) = runtime.refresh_if_needed().await
{
warn!("failed to refresh OAuth tokens: {error}");
if let Some(runtime) = &oauth_persistor {
runtime
.refresh_if_needed()
.await
.map_err(ClientOperationError::OAuth)?;
}
let result = operation(service)
.await
.map_err(ClientOperationError::from)?;
if let Some(runtime) = oauth_persistor
&& let Err(error) = runtime.persist_if_needed().await
{
warn!("failed to persist OAuth tokens: {error}");
if let Some(runtime) = oauth_persistor {
Self::persist_oauth_tokens_in_background(runtime, "failed to persist OAuth tokens");
}
Ok(result)
};
@@ -963,6 +965,24 @@ impl RmcpClient {
}
}
fn persist_oauth_tokens_in_background(runtime: OAuthPersistor, error_context: &'static str) {
Self::spawn_oauth_persistence(
async move { runtime.persist_if_needed().await },
error_context,
);
}
fn spawn_oauth_persistence(
persistence: impl Future<Output = Result<()>> + Send + 'static,
error_context: &'static str,
) {
tokio::spawn(async move {
if let Err(error) = persistence.await {
warn!("{error_context}: {error}");
}
});
}
fn is_session_expired_404(error: &ClientOperationError) -> bool {
let ClientOperationError::Service(rmcp::service::ServiceError::TransportSend(error)) =
error
@@ -1112,6 +1132,7 @@ mod tests {
use std::time::Duration;
use pretty_assertions::assert_eq;
use tokio::sync::oneshot;
use tokio::time;
use super::*;
@@ -1134,4 +1155,37 @@ mod tests {
assert_eq!(Ok("done"), result);
}
#[tokio::test]
async fn completed_operation_is_not_delayed_by_oauth_persistence() {
let (started_tx, started_rx) = oneshot::channel();
let (release_tx, release_rx) = oneshot::channel();
let (finished_tx, mut finished_rx) = oneshot::channel();
RmcpClient::spawn_oauth_persistence(
async move {
started_tx.send(()).expect("signal persistence start");
release_rx.await.expect("release persistence");
finished_tx.send(()).expect("signal persistence finish");
Ok(())
},
"test OAuth persistence failure",
);
let completed_result = "completed tool response";
time::timeout(Duration::from_secs(1), started_rx)
.await
.expect("persistence task did not start")
.expect("persistence start sender dropped");
assert_eq!(completed_result, "completed tool response");
assert!(matches!(
finished_rx.try_recv(),
Err(oneshot::error::TryRecvError::Empty)
));
release_tx.send(()).expect("release persistence task");
time::timeout(Duration::from_secs(1), finished_rx)
.await
.expect("persistence task did not finish")
.expect("persistence finish sender dropped");
}
}
codex-rs/rmcp-client/tests/streamable_http_oauth_refresh_race.rs
+224 -13
View File
@@ -429,6 +429,87 @@ async fn rotated_credentials_persist_after_startup_timeout_before_next_refresh()
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn provider_refresh_persists_after_startup_timeout_before_next_refresh() -> anyhow::Result<()>
{
let server = MockServer::start().await;
mount_oauth_metadata(&server, /*expected_requests*/ 2).await;
let refresh_attempts = Arc::new(AtomicUsize::new(0));
let responder_attempts = Arc::clone(&refresh_attempts);
Mock::given(method("POST"))
.and(path("/oauth/token"))
.respond_with(move |_request: &Request| {
responder_attempts.fetch_add(1, Ordering::SeqCst);
ResponseTemplate::new(200)
.set_delay(Duration::from_secs(6))
.set_body_json(json!({
"access_token": NEW_ACCESS_TOKEN,
"token_type": "Bearer",
"expires_in": 7200,
"refresh_token": NEW_REFRESH_TOKEN,
}))
})
.expect(1)
.mount(&server)
.await;
mount_mcp_server(&server, /*expected_requests*/ 2).await;
let codex_home = TempDir::new()?;
let control_dir = TempDir::new()?;
let server_url = format!("{}/mcp", server.uri());
seed_tokens(&codex_home, &server_url, /*expires_at*/ 0).await?;
let first_ready_path = control_dir.path().join("client-a.ready");
let first_go_path = control_dir.path().join("client-a.go");
let first_result_path = control_dir.path().join("client-a.result");
let first_release_path = control_dir.path().join("client-a.release");
let first_child = spawn_client_child_held_open(
&codex_home,
&server_url,
&first_ready_path,
&first_go_path,
&first_result_path,
&first_release_path,
)?;
wait_for_paths(std::slice::from_ref(&first_ready_path)).await?;
fs::write(&first_go_path, "go")?;
wait_for_paths(std::slice::from_ref(&first_result_path)).await?;
assert_eq!(
fs::read_to_string(&first_result_path)?,
"error:timed out handshaking with MCP server after 5s"
);
assert_eq!(refresh_attempts.load(Ordering::SeqCst), 1);
wait_for_persisted_token_snapshot(
&codex_home,
PersistedTokenSnapshot {
access_token: NEW_ACCESS_TOKEN.to_string(),
refresh_token: Some(NEW_REFRESH_TOKEN.to_string()),
},
)
.await?;
let second_ready_path = control_dir.path().join("client-b.ready");
let second_go_path = control_dir.path().join("client-b.go");
let second_result_path = control_dir.path().join("client-b.result");
let second_child = spawn_client_child(
&codex_home,
&server_url,
&second_ready_path,
&second_go_path,
&second_result_path,
)?;
wait_for_paths(std::slice::from_ref(&second_ready_path)).await?;
fs::write(&second_go_path, "go")?;
wait_for_children(vec![second_child]).await?;
assert_eq!(fs::read_to_string(second_result_path)?, "success");
assert_eq!(refresh_attempts.load(Ordering::SeqCst), 1);
fs::write(first_release_path, "release")?;
wait_for_children(vec![first_child]).await?;
server.verify().await;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn fallback_read_lock_wait_uses_operation_timeout_budget() -> anyhow::Result<()> {
let server = MockServer::start().await;
@@ -655,16 +736,21 @@ async fn refresh_provider_call_uses_operation_timeout_budget() -> anyhow::Result
Mock::given(method("POST"))
.and(path("/oauth/token"))
.respond_with(move |_request: &Request| {
let response = ResponseTemplate::new(200).set_body_json(json!({
"access_token": NEW_ACCESS_TOKEN,
"token_type": "Bearer",
"expires_in": 31,
"refresh_token": NEW_REFRESH_TOKEN,
}));
if responder_attempts.fetch_add(1, Ordering::SeqCst) == 0 {
response
} else {
response.set_delay(Duration::from_secs(1))
match responder_attempts.fetch_add(1, Ordering::SeqCst) {
0 => ResponseTemplate::new(200).set_body_json(json!({
"access_token": NEW_ACCESS_TOKEN,
"token_type": "Bearer",
"expires_in": 31,
"refresh_token": NEW_REFRESH_TOKEN,
})),
_ => ResponseTemplate::new(200)
.set_delay(Duration::from_secs(1))
.set_body_json(json!({
"access_token": FINAL_ACCESS_TOKEN,
"token_type": "Bearer",
"expires_in": 7200,
"refresh_token": FINAL_REFRESH_TOKEN,
})),
}
})
.expect(2)
@@ -680,23 +766,132 @@ async fn refresh_provider_call_uses_operation_timeout_budget() -> anyhow::Result
let ready_path = control_dir.path().join("client.ready");
let go_path = control_dir.path().join("client.go");
let result_path = control_dir.path().join("client.result");
let child = spawn_operation_timeout_child(
let release_path = control_dir.path().join("client.release");
let child = spawn_operation_timeout_child_held_open(
&codex_home,
&server_url,
&ready_path,
&go_path,
&result_path,
&release_path,
)?;
wait_for_paths(std::slice::from_ref(&ready_path)).await?;
sleep(Duration::from_secs(2)).await;
fs::write(&go_path, "go")?;
wait_for_children(vec![child]).await?;
wait_for_paths(std::slice::from_ref(&result_path)).await?;
assert_eq!(
fs::read_to_string(result_path)?,
fs::read_to_string(&result_path)?,
format!("error:timed out awaiting tools/list after {OPERATION_TIMEOUT:?}")
);
assert_eq!(refresh_attempts.load(Ordering::SeqCst), 2);
wait_for_persisted_token_snapshot(
&codex_home,
PersistedTokenSnapshot {
access_token: FINAL_ACCESS_TOKEN.to_string(),
refresh_token: Some(FINAL_REFRESH_TOKEN.to_string()),
},
)
.await?;
fs::write(release_path, "release")?;
wait_for_children(vec![child]).await?;
server.verify().await;
Ok(())
}
#[cfg(unix)]
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn rotated_token_storage_failure_blocks_another_process_refresh() -> anyhow::Result<()> {
use std::os::unix::fs::PermissionsExt;
let server = MockServer::start().await;
mount_oauth_metadata(&server, /*expected_requests*/ 2).await;
let refresh_attempts = Arc::new(AtomicUsize::new(0));
let responder_attempts = Arc::clone(&refresh_attempts);
Mock::given(method("POST"))
.and(path("/oauth/token"))
.respond_with(move |_request: &Request| {
match responder_attempts.fetch_add(1, Ordering::SeqCst) {
0 => ResponseTemplate::new(200).set_body_json(json!({
"access_token": NEW_ACCESS_TOKEN,
"token_type": "Bearer",
"expires_in": 31,
"refresh_token": NEW_REFRESH_TOKEN,
})),
1 => ResponseTemplate::new(200).set_body_json(json!({
"access_token": FINAL_ACCESS_TOKEN,
"token_type": "Bearer",
"expires_in": 7200,
"refresh_token": FINAL_REFRESH_TOKEN,
})),
_ => ResponseTemplate::new(400).set_body_json(json!({
"error": "invalid_grant",
"error_description": "refresh token already used",
})),
}
})
.expect(2)
.mount(&server)
.await;
mount_mcp_server(&server, /*expected_requests*/ 2).await;
let codex_home = TempDir::new()?;
let control_dir = TempDir::new()?;
let server_url = format!("{}/mcp", server.uri());
seed_tokens(&codex_home, &server_url, /*expires_at*/ 0).await?;
let first_ready_path = control_dir.path().join("client-a.ready");
let first_go_path = control_dir.path().join("client-a.go");
let first_result_path = control_dir.path().join("client-a.result");
let first_release_path = control_dir.path().join("client-a.release");
let first_child = spawn_operation_timeout_child_held_open(
&codex_home,
&server_url,
&first_ready_path,
&first_go_path,
&first_result_path,
&first_release_path,
)?;
wait_for_paths(std::slice::from_ref(&first_ready_path)).await?;
let credentials_path = codex_home.path().join(".credentials.json");
fs::set_permissions(&credentials_path, fs::Permissions::from_mode(0o400))?;
sleep(Duration::from_secs(2)).await;
fs::write(&first_go_path, "go")?;
wait_for_paths(std::slice::from_ref(&first_result_path)).await?;
let first_outcome = fs::read_to_string(&first_result_path)?;
assert!(
first_outcome.starts_with(&format!(
"error:OAuth refresh succeeded but failed to persist rotated credentials for server {SERVER_NAME}:"
)),
"unexpected first-client outcome: {first_outcome}"
);
assert_eq!(refresh_attempts.load(Ordering::SeqCst), 2);
let second_ready_path = control_dir.path().join("client-b.ready");
let second_go_path = control_dir.path().join("client-b.go");
let second_result_path = control_dir.path().join("client-b.result");
let mut second_child = spawn_client_child(
&codex_home,
&server_url,
&second_ready_path,
&second_go_path,
&second_result_path,
)?;
wait_for_paths(std::slice::from_ref(&second_ready_path)).await?;
fs::write(&second_go_path, "go")?;
sleep(Duration::from_millis(250)).await;
assert_eq!(refresh_attempts.load(Ordering::SeqCst), 2);
assert!(!second_result_path.exists());
second_child.kill().await?;
timeout(CHILD_WAIT_TIMEOUT, second_child.wait())
.await
.context("timed out stopping blocked OAuth client child")??;
fs::set_permissions(&credentials_path, fs::Permissions::from_mode(0o600))?;
fs::write(first_release_path, "release")?;
wait_for_children(vec![first_child]).await?;
server.verify().await;
Ok(())
@@ -1348,6 +1543,22 @@ async fn wait_for_paths(paths: &[PathBuf]) -> anyhow::Result<()> {
.context("timed out waiting for OAuth race process barrier")
}
async fn wait_for_persisted_token_snapshot(
codex_home: &TempDir,
expected: PersistedTokenSnapshot,
) -> anyhow::Result<()> {
timeout(Duration::from_secs(10), async {
loop {
if persisted_token_snapshot(codex_home).ok().as_ref() == Some(&expected) {
return;
}
sleep(Duration::from_millis(10)).await;
}
})
.await
.context("timed out waiting for rotated OAuth credentials to persist")
}
fn read_outcomes(paths: &[PathBuf]) -> anyhow::Result<Vec<String>> {
paths
.iter()