openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-03 19:53:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(rmcp): bound OAuth refresh coordination

Browse Source
This commit is contained in:
Adam Perry
2026-06-05 05:15:59 +00:00
parent a7b485f127
commit 0d5e92d6fa
3 changed files with 224 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/rmcp-client/src/oauth.rs
+101 -22
View File
@@ -37,6 +37,7 @@ use std::collections::BTreeMap;
use std::fs;
use std::fs::OpenOptions;
use std::io::ErrorKind;
use std::path::Path;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
@@ -55,7 +56,8 @@ use tokio::sync::Mutex;
use codex_utils_home_dir::find_codex_home;
const KEYRING_SERVICE: &str = "Codex MCP Credentials";
const OAUTH_REFRESH_LOCK_FILENAME: &str = ".credentials.json.refresh.lock";
const FILE_OAUTH_REFRESH_LOCK_FILENAME: &str = ".credentials.json.refresh.lock";
const KEYRING_OAUTH_REFRESH_LOCK_PREFIX: &str = "codex-mcp-oauth-refresh";
const REFRESH_SKEW_MILLIS: u64 = 30_000;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
@@ -369,13 +371,14 @@ impl OAuthPersistor {
// A different process may have rotated the one-time refresh token.
// Reload its durable result while all refreshers share the same lock.
let _refresh_lock = acquire_oauth_refresh_lock().await?;
if let Some(stored) = load_oauth_tokens(
let (_refresh_locks, stored) = lock_and_load_oauth_tokens(
&self.inner.server_name,
&self.inner.url,
self.inner.store_mode,
)? {
)
.await?;
if let Some(stored) = stored {
let credentials_changed = {
let last_credentials = self.inner.last_credentials.lock().await;
last_credentials.as_ref() != Some(&stored)
@@ -432,27 +435,96 @@ impl OAuthPersistor {
}
}
async fn acquire_oauth_refresh_lock() -> Result<fs::File> {
let path = find_codex_home()?.join(OAUTH_REFRESH_LOCK_FILENAME);
async fn lock_and_load_oauth_tokens(
server_name: &str,
url: &str,
store_mode: OAuthCredentialsStoreMode,
) -> Result<(Vec<fs::File>, Option<StoredOAuthTokens>)> {
let keyring_store = DefaultKeyringStore;
match store_mode {
OAuthCredentialsStoreMode::Auto => {
let keyring_lock =
acquire_oauth_refresh_lock(&keyring_oauth_refresh_lock_path()?).await?;
match load_oauth_tokens_from_keyring(&keyring_store, server_name, url) {
Ok(Some(tokens)) => Ok((vec![keyring_lock], Some(tokens))),
Ok(None) => {
let file_lock =
acquire_oauth_refresh_lock(&file_oauth_refresh_lock_path()?).await?;
let tokens = load_oauth_tokens_from_file(server_name, url)?;
Ok((vec![keyring_lock, file_lock], tokens))
}
Err(error) => {
warn!("failed to read OAuth tokens from keyring: {error}");
let file_lock =
acquire_oauth_refresh_lock(&file_oauth_refresh_lock_path()?).await?;
let tokens =
load_oauth_tokens_from_file(server_name, url).with_context(|| {
format!("failed to read OAuth tokens from keyring: {error}")
})?;
Ok((vec![keyring_lock, file_lock], tokens))
}
}
}
OAuthCredentialsStoreMode::File => {
let lock = acquire_oauth_refresh_lock(&file_oauth_refresh_lock_path()?).await?;
let tokens = load_oauth_tokens_from_file(server_name, url)?;
Ok((vec![lock], tokens))
}
OAuthCredentialsStoreMode::Keyring => {
let lock = acquire_oauth_refresh_lock(&keyring_oauth_refresh_lock_path()?).await?;
let tokens = load_oauth_tokens_from_keyring(&keyring_store, server_name, url)
.with_context(|| "failed to read OAuth tokens from keyring".to_string())?;
Ok((vec![lock], tokens))
}
}
}
fn file_oauth_refresh_lock_path() -> Result<PathBuf> {
Ok(find_codex_home()?
.join(FILE_OAUTH_REFRESH_LOCK_FILENAME)
.to_path_buf())
}
fn keyring_oauth_refresh_lock_path() -> Result<PathBuf> {
let user_namespace = ["HOME", "USERPROFILE", "USER", "USERNAME"]
.into_iter()
.find_map(std::env::var_os)
.unwrap_or_default();
let lock_id = sha_256_prefix(&Value::String(
user_namespace.to_string_lossy().into_owned(),
))?;
Ok(std::env::temp_dir().join(format!(
"{KEYRING_OAUTH_REFRESH_LOCK_PREFIX}-{lock_id}.lock"
)))
}
async fn acquire_oauth_refresh_lock(path: &Path) -> Result<fs::File> {
if let Some(parent) = path.parent() {
fs::create_dir_all(parent)?;
}
let file = OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(&path)
let mut options = OpenOptions::new();
options.read(true).write(true).create(true).truncate(false);
#[cfg(unix)]
{
use std::os::unix::fs::OpenOptionsExt;
options.mode(0o600);
}
let file = options
.open(path)
.with_context(|| format!("failed to open OAuth refresh lock at {}", path.display()))?;
tokio::task::spawn_blocking(move || {
file.lock().with_context(|| {
format!("failed to acquire OAuth refresh lock at {}", path.display())
})?;
Ok(file)
})
.await
.context("failed to join OAuth refresh lock task")?
loop {
match file.try_lock() {
Ok(()) => return Ok(file),
Err(fs::TryLockError::WouldBlock) => {
tokio::time::sleep(Duration::from_millis(10)).await;
}
Err(error) => {
return Err(error).with_context(|| {
format!("failed to acquire OAuth refresh lock at {}", path.display())
});
}
}
}
}
const FALLBACK_FILENAME: &str = ".credentials.json";
@@ -723,6 +795,13 @@ mod tests {
}
}
#[test]
fn keyring_refresh_lock_does_not_use_codex_home() -> Result<()> {
let lock_path = keyring_oauth_refresh_lock_path()?;
assert_eq!(lock_path.parent(), Some(std::env::temp_dir().as_path()));
Ok(())
}
#[test]
fn load_oauth_tokens_reads_from_keyring_when_available() -> Result<()> {
let _env = TempCodexHome::new();
codex-rs/rmcp-client/src/rmcp_client.rs
+33 -30
View File
@@ -824,42 +824,45 @@ impl RmcpClient {
Arc<RunningService<RoleClient, ElicitationClientService>>,
Option<OAuthPersistor>,
)> {
let (transport, oauth_persistor) = match pending_transport {
PendingTransport::InProcess { transport } => (
service::serve_client(client_service, transport).boxed(),
None,
),
PendingTransport::Stdio { transport } => (
service::serve_client(client_service, transport).boxed(),
None,
),
PendingTransport::StreamableHttp { transport } => (
service::serve_client(client_service, transport).boxed(),
None,
),
PendingTransport::StreamableHttpWithOAuth {
transport,
oauth_persistor,
} => {
oauth_persistor.refresh_if_needed().await?;
(
let connect = async move {
let (transport, oauth_persistor) = match pending_transport {
PendingTransport::InProcess { transport } => (
service::serve_client(client_service, transport).boxed(),
None,
),
PendingTransport::Stdio { transport } => (
service::serve_client(client_service, transport).boxed(),
None,
),
PendingTransport::StreamableHttp { transport } => (
service::serve_client(client_service, transport).boxed(),
None,
),
PendingTransport::StreamableHttpWithOAuth {
transport,
oauth_persistor,
} => (
service::serve_client(client_service, transport).boxed(),
Some(oauth_persistor),
)
),
};
if let Some(runtime) = &oauth_persistor {
runtime.refresh_if_needed().await?;
}
let service = transport
.await
.map_err(|err| anyhow!("handshaking with MCP server failed: {err}"))?;
Ok((Arc::new(service), oauth_persistor))
};
let service = match timeout {
Some(duration) => time::timeout(duration, transport)
match timeout {
Some(duration) => time::timeout(duration, connect)
.await
.map_err(|_| anyhow!("timed out handshaking with MCP server after {duration:?}"))?
.map_err(|err| anyhow!("handshaking with MCP server failed: {err}"))?,
None => transport
.await
.map_err(|err| anyhow!("handshaking with MCP server failed: {err}"))?,
};
Ok((Arc::new(service), oauth_persistor))
.map_err(|_| anyhow!("timed out handshaking with MCP server after {duration:?}"))?,
None => connect.await,
}
}
async fn run_service_operation<T, F, Fut>(
codex-rs/rmcp-client/tests/streamable_http_oauth_refresh_race.rs
+90 -10
View File
@@ -2,6 +2,7 @@ mod streamable_http_test_support;
use std::collections::BTreeMap;
use std::fs;
use std::fs::OpenOptions;
use std::path::Path;
use std::path::PathBuf;
use std::sync::Arc;
@@ -47,6 +48,7 @@ const OLD_ACCESS_TOKEN: &str = "old-access-token";
const OLD_REFRESH_TOKEN: &str = "old-one-time-refresh-token";
const NEW_ACCESS_TOKEN: &str = "new-access-token";
const NEW_REFRESH_TOKEN: &str = "new-one-time-refresh-token";
const FILE_OAUTH_REFRESH_LOCK_FILENAME: &str = ".credentials.json.refresh.lock";
const SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_RACE_SERVER_URL";
const READY_PATH_ENV: &str = "MCP_TEST_OAUTH_RACE_READY_PATH";
@@ -68,12 +70,14 @@ async fn concurrent_processes_coordinate_one_time_refresh_and_reload_rotated_tok
.and(path("/oauth/token"))
.respond_with(move |_request: &Request| {
if 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": 7200,
"refresh_token": NEW_REFRESH_TOKEN,
}))
ResponseTemplate::new(200)
.set_delay(Duration::from_secs(1))
.set_body_json(json!({
"access_token": NEW_ACCESS_TOKEN,
"token_type": "Bearer",
"expires_in": 7200,
"refresh_token": NEW_REFRESH_TOKEN,
}))
} else {
ResponseTemplate::new(400).set_body_json(json!({
"error": "invalid_grant",
@@ -91,7 +95,10 @@ async fn concurrent_processes_coordinate_one_time_refresh_and_reload_rotated_tok
let server_url = format!("{}/mcp", server.uri());
seed_tokens(&codex_home, &server_url, /*expires_at*/ 0).await?;
let go_path = control_dir.path().join("go");
let go_paths = [
control_dir.path().join("client-a.go"),
control_dir.path().join("client-b.go"),
];
let ready_paths = [
control_dir.path().join("client-a.ready"),
control_dir.path().join("client-b.ready"),
@@ -101,18 +108,30 @@ async fn concurrent_processes_coordinate_one_time_refresh_and_reload_rotated_tok
control_dir.path().join("client-b.result"),
];
let mut children = Vec::new();
for (ready_path, result_path) in ready_paths.iter().zip(result_paths.iter()) {
for ((ready_path, go_path), result_path) in ready_paths
.iter()
.zip(go_paths.iter())
.zip(result_paths.iter())
{
children.push(spawn_client_child(
&codex_home,
&server_url,
ready_path,
&go_path,
go_path,
result_path,
)?);
}
wait_for_paths(&ready_paths).await?;
fs::write(&go_path, "go")?;
fs::write(&go_paths[0], "go")?;
timeout(Duration::from_secs(5), async {
while refresh_attempts.load(Ordering::SeqCst) == 0 {
sleep(Duration::from_millis(10)).await;
}
})
.await
.context("timed out waiting for the first process to hold the refresh lock")?;
fs::write(&go_paths[1], "go")?;
wait_for_children(children).await?;
let outcomes = read_outcomes(&result_paths)?;
@@ -130,6 +149,23 @@ async fn concurrent_processes_coordinate_one_time_refresh_and_reload_rotated_tok
.iter()
.all(|body| body.contains(&format!("refresh_token={OLD_REFRESH_TOKEN}")))
);
let authorization_headers = requests
.iter()
.filter(|request| request.method.as_str() == "POST" && request.url.path() == "/mcp")
.map(|request| {
request
.headers
.get("authorization")
.expect("authorization header")
.to_str()
.expect("ASCII authorization header")
.to_string()
})
.collect::<Vec<_>>();
assert_eq!(
authorization_headers,
vec![format!("Bearer {NEW_ACCESS_TOKEN}"); 4]
);
assert_eq!(
persisted_token_snapshot(&codex_home)?,
PersistedTokenSnapshot {
@@ -142,6 +178,50 @@ async fn concurrent_processes_coordinate_one_time_refresh_and_reload_rotated_tok
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn refresh_lock_wait_uses_startup_timeout_budget() -> anyhow::Result<()> {
let server = MockServer::start().await;
mount_oauth_metadata(&server, /*expected_requests*/ 1).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 lock_path = codex_home.path().join(FILE_OAUTH_REFRESH_LOCK_FILENAME);
let refresh_lock = OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(lock_path)?;
refresh_lock.lock()?;
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_client_child(
&codex_home,
&server_url,
&ready_path,
&go_path,
&result_path,
)?;
wait_for_paths(std::slice::from_ref(&ready_path)).await?;
fs::write(&go_path, "go")?;
timeout(Duration::from_secs(7), wait_for_children(vec![child]))
.await
.context("startup did not time out while waiting for the refresh lock")??;
assert_eq!(
fs::read_to_string(result_path)?,
"error:timed out handshaking with MCP server after 5s"
);
server.verify().await;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn process_keeps_stale_in_memory_tokens_but_does_not_overwrite_rotated_tokens()
-> anyhow::Result<()> {