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

fix(rmcp-client): serialize OAuth credential generations

Browse Source
This commit is contained in:
Adam Perry
2026-06-05 05:23:49 +00:00
parent e4a137185c
commit a525c9cf4b
3 changed files with 554 additions and 357 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/rmcp-client/src/oauth.rs
+279 -67
View File
@@ -39,6 +39,7 @@ use std::fs::OpenOptions;
use std::io::ErrorKind;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::OnceLock;
use std::time::Duration;
use std::time::SystemTime;
use std::time::UNIX_EPOCH;
@@ -51,7 +52,7 @@ use rmcp::transport::auth::CredentialStore;
use rmcp::transport::auth::InMemoryCredentialStore;
use rmcp::transport::auth::StoredCredentials;
use tokio::sync::Mutex;
use tokio::sync::Semaphore;
use tokio::sync::OwnedMutexGuard;
use codex_utils_home_dir::find_codex_home;
@@ -172,6 +173,15 @@ pub fn save_oauth_tokens(
save_oauth_tokens_locked(server_name, tokens, store_mode)
}
pub(crate) async fn save_oauth_tokens_async(
server_name: &str,
tokens: &StoredOAuthTokens,
store_mode: OAuthCredentialsStoreMode,
) -> Result<()> {
let _lock = acquire_oauth_server_lock_async(server_name, &tokens.url).await?;
save_oauth_tokens_locked(server_name, tokens, store_mode)
}
fn save_oauth_tokens_locked(
server_name: &str,
tokens: &StoredOAuthTokens,
@@ -237,6 +247,15 @@ pub fn delete_oauth_tokens(
server_name: &str,
url: &str,
store_mode: OAuthCredentialsStoreMode,
) -> Result<bool> {
let _lock = acquire_oauth_server_lock(server_name, url)?;
delete_oauth_tokens_locked(server_name, url, store_mode)
}
fn delete_oauth_tokens_locked(
server_name: &str,
url: &str,
store_mode: OAuthCredentialsStoreMode,
) -> Result<bool> {
let keyring_store = DefaultKeyringStore;
delete_oauth_tokens_from_keyring_and_file(&keyring_store, store_mode, server_name, url)
@@ -280,7 +299,12 @@ struct OAuthPersistorInner {
authorization_manager: Arc<Mutex<AuthorizationManager>>,
store_mode: OAuthCredentialsStoreMode,
last_credentials: Mutex<Option<StoredOAuthTokens>>,
update_lock: Semaphore,
}
enum CredentialReload {
Unchanged,
Replaced,
Removed,
}
impl OAuthPersistor {
@@ -298,7 +322,6 @@ impl OAuthPersistor {
authorization_manager,
store_mode,
last_credentials: Mutex::new(initial_credentials),
update_lock: Semaphore::new(1),
}),
}
}
@@ -306,7 +329,14 @@ 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<()> {
let _permit = self.inner.update_lock.acquire().await?;
let _server_lock =
acquire_oauth_server_lock_async(&self.inner.server_name, &self.inner.url).await?;
if !matches!(
self.reload_persisted_credentials_locked().await?,
CredentialReload::Unchanged
) {
return Ok(());
}
self.persist_if_needed_locked().await
}
@@ -341,36 +371,22 @@ impl OAuthPersistor {
token_response: new_token_response,
expires_at,
};
if last_credentials
if !last_credentials
.as_ref()
.is_some_and(|previous| same_persisted_generation(previous, &stored))
{
return Ok(());
save_oauth_tokens_locked(
&self.inner.server_name,
&stored,
self.inner.store_mode,
)?;
*self.inner.last_credentials.lock().await = Some(stored);
}
let _server_lock =
acquire_oauth_server_lock_async(&self.inner.server_name, &self.inner.url)
.await?;
if let Some(persisted) = load_oauth_tokens(
&self.inner.server_name,
&self.inner.url,
self.inner.store_mode,
)? && !last_credentials
.as_ref()
.is_some_and(|previous| same_persisted_generation(previous, &persisted))
{
self.replace_manager_credentials(&persisted).await?;
*self.inner.last_credentials.lock().await = Some(persisted);
return Ok(());
}
save_oauth_tokens_locked(&self.inner.server_name, &stored, self.inner.store_mode)?;
*self.inner.last_credentials.lock().await = Some(stored);
}
None => {
let mut last_serialized = self.inner.last_credentials.lock().await;
if last_serialized.take().is_some()
&& let Err(error) = delete_oauth_tokens(
&& let Err(error) = delete_oauth_tokens_locked(
&self.inner.server_name,
&self.inner.url,
self.inner.store_mode,
@@ -389,11 +405,20 @@ impl OAuthPersistor {
#[expect(
clippy::await_holding_invalid_type,
reason = "AuthorizationManager async access and persisted credential adoption must be serialized"
reason = "the server credential lock must span reload, refresh, and persistence"
)]
pub(crate) async fn refresh_if_needed(&self) -> Result<()> {
let _permit = self.inner.update_lock.acquire().await?;
self.adopt_persisted_credentials_if_changed().await?;
let _server_lock =
acquire_oauth_server_lock_async(&self.inner.server_name, &self.inner.url).await?;
if matches!(
self.reload_persisted_credentials_locked().await?,
CredentialReload::Removed
) {
return Err(anyhow::anyhow!(
"OAuth credentials were removed for server {}",
self.inner.server_name
));
}
let expires_at = {
let guard = self.inner.last_credentials.lock().await;
@@ -418,29 +443,41 @@ impl OAuthPersistor {
self.persist_if_needed_locked().await
}
async fn adopt_persisted_credentials_if_changed(&self) -> Result<()> {
let _server_lock =
acquire_oauth_server_lock_async(&self.inner.server_name, &self.inner.url).await?;
let Some(persisted) = load_oauth_tokens(
async fn reload_persisted_credentials_locked(&self) -> Result<CredentialReload> {
let loaded = load_oauth_tokens(
&self.inner.server_name,
&self.inner.url,
self.inner.store_mode,
)?
else {
return Ok(());
};
)
.with_context(|| {
format!(
"failed to reload OAuth tokens for server {}",
self.inner.server_name
)
})?;
let last_credentials = self.inner.last_credentials.lock().await.clone();
if last_credentials
.as_ref()
.is_some_and(|previous| same_persisted_generation(previous, &persisted))
{
return Ok(());
let unchanged = match (&loaded, &last_credentials) {
(Some(loaded), Some(previous)) => same_persisted_generation(loaded, previous),
(None, None) => true,
_ => false,
};
if unchanged {
return Ok(CredentialReload::Unchanged);
}
self.replace_manager_credentials(&persisted).await?;
*self.inner.last_credentials.lock().await = Some(persisted);
Ok(())
match loaded {
Some(tokens) => {
self.replace_manager_credentials(&tokens).await?;
*self.inner.last_credentials.lock().await = Some(tokens);
Ok(CredentialReload::Replaced)
}
None => {
self.clear_manager_credentials().await;
*self.inner.last_credentials.lock().await = None;
Ok(CredentialReload::Removed)
}
}
}
async fn replace_manager_credentials(&self, tokens: &StoredOAuthTokens) -> Result<()> {
@@ -469,33 +506,112 @@ impl OAuthPersistor {
guard.set_credential_store(store);
Ok(())
}
async fn clear_manager_credentials(&self) {
let manager = self.inner.authorization_manager.clone();
let mut guard = manager.lock().await;
guard.set_credential_store(InMemoryCredentialStore::new());
}
}
fn oauth_server_lock_for(server_name: &str, url: &str) -> Arc<Mutex<()>> {
static OAUTH_SERVER_LOCKS: OnceLock<std::sync::Mutex<BTreeMap<String, Arc<Mutex<()>>>>> =
OnceLock::new();
let mut locks = OAUTH_SERVER_LOCKS
.get_or_init(std::sync::Mutex::default)
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
locks
.entry(format!("{server_name}\n{url}"))
.or_insert_with(|| Arc::new(Mutex::new(())))
.clone()
}
struct OAuthServerLock {
_file: fs::File,
_in_process_guard: Option<OwnedMutexGuard<()>>,
}
fn acquire_oauth_server_lock(server_name: &str, url: &str) -> Result<OAuthServerLock> {
let digest = sha_256_prefix(&Value::String(format!("{server_name}\n{url}")))?;
let path = find_codex_home()?.join(format!(".mcp-oauth-{digest}.lock"));
let file = OpenOptions::new()
let path = oauth_server_lock_path(server_name, url)?;
let file = open_oauth_lock_file(&path)?;
file.lock()
.with_context(|| format!("failed to lock OAuth credentials at {}", path.display()))?;
Ok(OAuthServerLock {
_file: file,
_in_process_guard: None,
})
}
async fn acquire_oauth_server_lock_async(server_name: &str, url: &str) -> Result<OAuthServerLock> {
let in_process_guard = oauth_server_lock_for(server_name, url).lock_owned().await;
let path = oauth_server_lock_path(server_name, url)?;
let file = tokio::task::spawn_blocking({
let path = path.clone();
move || {
let file = open_oauth_lock_file(&path)?;
file.lock().with_context(|| {
format!("failed to lock OAuth credentials at {}", path.display())
})?;
Ok::<_, anyhow::Error>(file)
}
})
.await
.context("OAuth credential lock task failed")??;
Ok(OAuthServerLock {
_file: file,
_in_process_guard: Some(in_process_guard),
})
}
struct FallbackStoreLock {
_in_process_guard: std::sync::MutexGuard<'static, ()>,
_file: fs::File,
}
fn acquire_fallback_store_lock() -> Result<FallbackStoreLock> {
static FALLBACK_STORE_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
let in_process_guard = FALLBACK_STORE_LOCK
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let path = find_codex_home()?.join(".mcp-oauth-locks/fallback-store.lock");
let file = open_oauth_lock_file(&path)?;
file.lock().with_context(|| {
format!(
"failed to lock fallback OAuth credentials at {}",
path.display()
)
})?;
Ok(FallbackStoreLock {
_in_process_guard: in_process_guard,
_file: file,
})
}
fn open_oauth_lock_file(path: &std::path::Path) -> Result<fs::File> {
if let Some(parent) = path.parent() {
fs::create_dir_all(parent)?;
}
OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(&path)
.with_context(|| format!("failed to open OAuth credential lock at {}", path.display()))?;
file.lock()
.with_context(|| format!("failed to lock OAuth credentials at {}", path.display()))?;
Ok(OAuthServerLock { _file: file })
.open(path)
.with_context(|| format!("failed to open OAuth credential lock at {}", path.display()))
}
async fn acquire_oauth_server_lock_async(server_name: &str, url: &str) -> Result<OAuthServerLock> {
let server_name = server_name.to_string();
let url = url.to_string();
tokio::task::spawn_blocking(move || acquire_oauth_server_lock(&server_name, &url))
.await
.context("OAuth credential lock task failed")?
fn oauth_server_lock_path(server_name: &str, url: &str) -> Result<PathBuf> {
let server_digest = sha_256_prefix(&Value::String(format!("{server_name}\n{url}")))?;
let user_identity = std::env::var_os("HOME")
.or_else(|| std::env::var_os("USERPROFILE"))
.unwrap_or_else(|| std::env::temp_dir().into_os_string());
let user_digest = sha_256_prefix(&Value::String(user_identity.to_string_lossy().into_owned()))?;
Ok(std::env::temp_dir()
.join(format!("codex-mcp-oauth-{user_digest}"))
.join(format!("server-{server_digest}.lock")))
}
fn same_persisted_generation(left: &StoredOAuthTokens, right: &StoredOAuthTokens) -> bool {
@@ -507,16 +623,15 @@ fn same_persisted_generation(left: &StoredOAuthTokens, right: &StoredOAuthTokens
return false;
}
let normalize = |response: &WrappedOAuthTokenResponse| {
let mut value = serde_json::to_value(response).ok()?;
value.as_object_mut()?.remove("expires_in");
Some(value)
};
let mut left_response = left.token_response.0.clone();
let mut right_response = right.token_response.0.clone();
left_response.set_expires_in(None);
right_response.set_expires_in(None);
match (
normalize(&left.token_response),
normalize(&right.token_response),
serde_json::to_string(&left_response),
serde_json::to_string(&right_response),
) {
(Some(left), Some(right)) => left == right,
(Ok(left), Ok(right)) => left == right,
_ => false,
}
}
@@ -541,6 +656,7 @@ struct FallbackTokenEntry {
}
fn load_oauth_tokens_from_file(server_name: &str, url: &str) -> Result<Option<StoredOAuthTokens>> {
let _lock = acquire_fallback_store_lock()?;
let Some(store) = read_fallback_file()? else {
return Ok(None);
};
@@ -584,6 +700,7 @@ fn load_oauth_tokens_from_file(server_name: &str, url: &str) -> Result<Option<St
}
fn save_oauth_tokens_to_file(tokens: &StoredOAuthTokens) -> Result<()> {
let _lock = acquire_fallback_store_lock()?;
let key = compute_store_key(&tokens.server_name, &tokens.url)?;
let mut store = read_fallback_file()?.unwrap_or_default();
@@ -613,6 +730,7 @@ fn save_oauth_tokens_to_file(tokens: &StoredOAuthTokens) -> Result<()> {
}
fn delete_oauth_tokens_from_file(key: &str) -> Result<bool> {
let _lock = acquire_fallback_store_lock()?;
let mut store = match read_fallback_file()? {
Some(store) => store,
None => return Ok(false),
@@ -750,10 +868,13 @@ mod tests {
use anyhow::Result;
use keyring::Error as KeyringError;
use pretty_assertions::assert_eq;
use std::sync::Barrier;
use std::sync::Mutex;
use std::sync::MutexGuard;
use std::sync::OnceLock;
use std::sync::PoisonError;
use std::sync::mpsc;
use std::thread;
use tempfile::tempdir;
use codex_keyring_store::tests::MockKeyringStore;
@@ -999,6 +1120,97 @@ mod tests {
assert_eq!(tokens.token_response.0.expires_in(), Some(Duration::ZERO));
}
#[test]
fn persisted_generation_ignores_reload_expires_in_drift() -> Result<()> {
let _env = TempCodexHome::new();
let mut tokens = sample_tokens();
tokens
.token_response
.0
.set_expires_in(Some(&Duration::from_secs(7200)));
super::save_oauth_tokens_to_file(&tokens)?;
let loaded = super::load_oauth_tokens_from_file(&tokens.server_name, &tokens.url)?
.expect("saved tokens should reload");
assert_ne!(
loaded.token_response.0.expires_in(),
tokens.token_response.0.expires_in()
);
assert_ne!(loaded, tokens);
assert!(super::same_persisted_generation(&loaded, &tokens));
Ok(())
}
#[test]
fn fallback_store_serializes_writes_for_different_servers() -> Result<()> {
let _env = TempCodexHome::new();
let first = sample_tokens();
let mut second = sample_tokens();
second.server_name = "other-server".to_string();
second.url = "https://other.example.test".to_string();
second.token_response.0 = OAuthTokenResponse::new(
AccessToken::new("other-access-token".to_string()),
BasicTokenType::Bearer,
VendorExtraTokenFields::default(),
);
let fallback_lock = super::acquire_fallback_store_lock()?;
let barrier = Arc::new(Barrier::new(3));
let (tx, rx) = mpsc::channel();
let handles = [first.clone(), second.clone()]
.into_iter()
.map(|tokens| {
let barrier = Arc::clone(&barrier);
let tx = tx.clone();
thread::spawn(move || {
barrier.wait();
tx.send(super::save_oauth_tokens_to_file(&tokens))
.expect("send save result");
})
})
.collect::<Vec<_>>();
drop(tx);
barrier.wait();
thread::sleep(Duration::from_millis(50));
assert!(matches!(rx.try_recv(), Err(mpsc::TryRecvError::Empty)));
drop(fallback_lock);
for _ in 0..2 {
rx.recv_timeout(Duration::from_secs(5))??;
}
for handle in handles {
handle.join().expect("fallback writer should not panic");
}
let store = super::read_fallback_file()?.expect("fallback file should load");
let first_key = super::compute_store_key(&first.server_name, &first.url)?;
let second_key = super::compute_store_key(&second.server_name, &second.url)?;
assert_eq!(
store
.keys()
.cloned()
.collect::<std::collections::BTreeSet<_>>(),
[first_key, second_key].into_iter().collect()
);
Ok(())
}
#[test]
fn server_generation_lock_is_shared_across_codex_homes() -> Result<()> {
let _env = TempCodexHome::new();
let first = super::oauth_server_lock_path("server", "https://example.test")?;
let other_codex_home = tempdir()?;
unsafe {
std::env::set_var("CODEX_HOME", other_codex_home.path());
}
let second = super::oauth_server_lock_path("server", "https://example.test")?;
assert_eq!(first, second);
Ok(())
}
fn assert_tokens_match_without_expiry(
actual: &StoredOAuthTokens,
expected: &StoredOAuthTokens,
codex-rs/rmcp-client/src/perform_oauth_login.rs
+2 -2
View File
@@ -26,7 +26,7 @@ use urlencoding::decode;
use crate::StoredOAuthTokens;
use crate::WrappedOAuthTokenResponse;
use crate::oauth::compute_expires_at_millis;
use crate::save_oauth_tokens;
use crate::oauth::save_oauth_tokens_async;
use crate::utils::apply_default_headers;
use crate::utils::build_default_headers;
use codex_config::types::OAuthCredentialsStoreMode;
@@ -571,7 +571,7 @@ impl OauthLoginFlow {
token_response: WrappedOAuthTokenResponse(credentials),
expires_at,
};
save_oauth_tokens(&self.server_name, &stored, self.store_mode)?;
save_oauth_tokens_async(&self.server_name, &stored, self.store_mode).await?;
Ok(())
}
codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs
+273 -288
View File
@@ -9,6 +9,7 @@ use codex_exec_server::Environment;
use codex_rmcp_client::RmcpClient;
use codex_rmcp_client::StoredOAuthTokens;
use codex_rmcp_client::WrappedOAuthTokenResponse;
use codex_rmcp_client::delete_oauth_tokens;
use codex_rmcp_client::save_oauth_tokens;
use oauth2::AccessToken;
use oauth2::RefreshToken;
@@ -24,18 +25,16 @@ use wiremock::MockServer;
use wiremock::Request;
use wiremock::ResponseTemplate;
use wiremock::matchers::body_string_contains;
use wiremock::matchers::header;
use wiremock::matchers::method;
use wiremock::matchers::path;
use streamable_http_test_support::initialize_client;
const SERVER_NAME: &str = "test-streamable-http-oauth-startup";
const EXPIRED_ACCESS_TOKEN: &str = "expired-access-token";
const REFRESH_TOKEN: &str = "valid-refresh-token";
const REFRESHED_ACCESS_TOKEN: &str = "refreshed-access-token";
const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_STARTUP_SERVER_URL";
const EXTERNAL_UPDATE_CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_EXTERNAL_UPDATE_SERVER_URL";
const REFRESH_CONTENTION_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_REFRESH_CONTENTION_SERVER_URL";
const REFRESH_CONTENTION_ATTEMPT_ENV: &str = "MCP_TEST_OAUTH_REFRESH_CONTENTION_ATTEMPT";
const REFRESH_CONTENTION_DONE_ENV: &str = "MCP_TEST_OAUTH_REFRESH_CONTENTION_DONE";
const OLD_ACCESS_TOKEN: &str = "old-access-token";
const OLD_REFRESH_TOKEN: &str = "old-refresh-token";
const EXTERNAL_ACCESS_TOKEN: &str = "external-access-token";
@@ -44,186 +43,22 @@ const STALE_REFRESHED_ACCESS_TOKEN: &str = "stale-refreshed-access-token";
const STALE_ROTATED_REFRESH_TOKEN: &str = "stale-rotated-refresh-token";
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result<()> {
async fn running_client_adopts_external_oauth_update_and_deletion() -> anyhow::Result<()> {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/.well-known/oauth-authorization-server/mcp"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"authorization_endpoint": format!("{}/oauth/authorize", server.uri()),
"token_endpoint": format!("{}/oauth/token", server.uri()),
"scopes_supported": [""],
})))
.expect(1)
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/oauth/token"))
.and(body_string_contains("grant_type=refresh_token"))
.and(body_string_contains(format!(
"refresh_token={REFRESH_TOKEN}"
)))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"access_token": REFRESHED_ACCESS_TOKEN,
"token_type": "Bearer",
"expires_in": 7200,
"refresh_token": REFRESH_TOKEN,
})))
.expect(1)
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/mcp"))
.and(header(
"authorization",
format!("Bearer {REFRESHED_ACCESS_TOKEN}"),
))
.respond_with(|request: &Request| {
let body: Value = request.body_json().expect("valid JSON-RPC request");
match body.get("method").and_then(Value::as_str) {
Some("initialize") => ResponseTemplate::new(200).set_body_json(json!({
"jsonrpc": "2.0",
"id": body.get("id").cloned().unwrap_or(Value::Null),
"result": {
"protocolVersion": body
.pointer("/params/protocolVersion")
.cloned()
.unwrap_or_else(|| json!("2025-06-18")),
"capabilities": {},
"serverInfo": {
"name": "oauth-startup-test",
"version": "0.0.0-test",
},
},
})),
Some("notifications/initialized") => ResponseTemplate::new(202),
method => ResponseTemplate::new(400)
.set_body_string(format!("unexpected JSON-RPC method: {method:?}")),
}
})
.expect(2)
.mount(&server)
.await;
let codex_home = TempDir::new()?;
let server_url = format!("{}/mcp", server.uri());
// Credential storage resolves CODEX_HOME from the process environment.
// Run the client half of the test in an ignored helper test so it can use
// an isolated home without mutating the parent test runner's environment.
let status = Command::new(std::env::current_exe()?)
.args(["oauth_startup_child", "--exact", "--ignored", "--nocapture"])
.env("CODEX_HOME", codex_home.path())
.env(CHILD_SERVER_URL_ENV, server_url)
.status()
.await?;
assert!(status.success(), "OAuth startup child failed: {status}");
server.verify().await;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
#[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)?;
// Save an expired access token with a valid refresh token so startup must
// refresh before sending the initialize request.
let mut response = OAuthTokenResponse::new(
AccessToken::new(EXPIRED_ACCESS_TOKEN.to_string()),
BasicTokenType::Bearer,
VendorExtraTokenFields::default(),
);
response.set_refresh_token(Some(RefreshToken::new(REFRESH_TOKEN.to_string())));
response.set_expires_in(Some(&Duration::from_secs(7200)));
let tokens = StoredOAuthTokens {
server_name: SERVER_NAME.to_string(),
url: server_url.clone(),
client_id: "test-client-id".to_string(),
token_response: WrappedOAuthTokenResponse(response),
expires_at: Some(0),
};
save_oauth_tokens(SERVER_NAME, &tokens, OAuthCredentialsStoreMode::File)?;
// This mirrors create_client's transport and initialization setup, except
// it omits the direct bearer token. Supplying that token would bypass the
// persisted OAuth credentials and the startup refresh under test.
let client = RmcpClient::new_streamable_http_client(
SERVER_NAME,
&server_url,
/*bearer_token*/ None,
/*http_headers*/ None,
/*env_http_headers*/ None,
OAuthCredentialsStoreMode::File,
Environment::default_for_tests().get_http_client(),
/*auth_provider*/ None,
)
.await?;
initialize_client(&client).await?;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn running_client_adopts_external_oauth_update_before_next_operation() -> anyhow::Result<()> {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/.well-known/oauth-authorization-server/mcp"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"authorization_endpoint": format!("{}/oauth/authorize", server.uri()),
"token_endpoint": format!("{}/oauth/token", server.uri()),
"scopes_supported": [""],
})))
.expect(1)
.mount(&server)
.await;
mount_oauth_metadata(&server).await;
Mock::given(method("POST"))
.and(path("/mcp"))
.respond_with(|request: &Request| {
let body: Value = request.body_json().expect("valid JSON-RPC request");
let Ok(body) = request.body_json::<Value>() else {
return ResponseTemplate::new(400);
};
let request_method = body.get("method").and_then(Value::as_str);
let expected_access_token = match request_method {
Some("tools/call") => EXTERNAL_ACCESS_TOKEN,
_ => OLD_ACCESS_TOKEN,
};
let authorization = request
.headers
.get("authorization")
.and_then(|value| value.to_str().ok());
if authorization != Some(format!("Bearer {expected_access_token}").as_str()) {
return ResponseTemplate::new(401);
}
match request_method {
Some("initialize") => ResponseTemplate::new(200).set_body_json(json!({
"jsonrpc": "2.0",
"id": body.get("id").cloned().unwrap_or(Value::Null),
"result": {
"protocolVersion": body
.pointer("/params/protocolVersion")
.cloned()
.unwrap_or_else(|| json!("2025-06-18")),
"capabilities": {
"tools": {},
},
"serverInfo": {
"name": "oauth-external-update-test",
"version": "0.0.0-test",
},
},
})),
Some("notifications/initialized") => ResponseTemplate::new(202),
Some("tools/call") => ResponseTemplate::new(200).set_body_json(json!({
"jsonrpc": "2.0",
"id": body.get("id").cloned().unwrap_or(Value::Null),
"result": {
"content": [],
"isError": false,
},
})),
method => ResponseTemplate::new(400)
.set_body_string(format!("unexpected JSON-RPC method: {method:?}")),
}
oauth_mcp_response(request, expected_access_token)
})
.expect(3)
.mount(&server)
@@ -251,7 +86,7 @@ async fn running_client_adopts_external_oauth_update_before_next_operation() ->
}
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
#[ignore = "spawned by running_client_adopts_external_oauth_update_before_next_operation"]
#[ignore = "spawned by running_client_adopts_external_oauth_update_and_deletion"]
async fn oauth_external_update_child() -> anyhow::Result<()> {
let server_url = std::env::var(EXTERNAL_UPDATE_CHILD_SERVER_URL_ENV)?;
save_test_tokens(
@@ -261,18 +96,7 @@ async fn oauth_external_update_child() -> anyhow::Result<()> {
future_expiry()?,
)?;
let client = RmcpClient::new_streamable_http_client(
SERVER_NAME,
&server_url,
/*bearer_token*/ None,
/*http_headers*/ None,
/*env_http_headers*/ None,
OAuthCredentialsStoreMode::File,
Environment::default_for_tests().get_http_client(),
/*auth_provider*/ None,
)
.await?;
initialize_client(&client).await?;
let client = initialized_oauth_client(&server_url).await?;
save_test_tokens(
&server_url,
@@ -291,13 +115,105 @@ async fn oauth_external_update_child() -> anyhow::Result<()> {
.await?;
assert_persisted_tokens(EXTERNAL_ACCESS_TOKEN, EXTERNAL_REFRESH_TOKEN)?;
let status = Command::new(std::env::current_exe()?)
.args([
"oauth_external_delete_writer_child",
"--exact",
"--ignored",
"--nocapture",
])
.env("CODEX_HOME", std::env::var("CODEX_HOME")?)
.env(EXTERNAL_UPDATE_CHILD_SERVER_URL_ENV, &server_url)
.status()
.await?;
assert!(
status.success(),
"OAuth external delete writer failed: {status}"
);
client
.call_tool(
"call-after-external-delete".to_string(),
/*arguments*/ None,
/*meta*/ None,
Some(Duration::from_secs(5)),
)
.await
.expect_err("externally deleted credentials must not be reused");
assert!(
!std::path::Path::new(&std::env::var("CODEX_HOME")?)
.join(".credentials.json")
.exists()
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn refresh_does_not_overwrite_external_oauth_update() -> anyhow::Result<()> {
#[ignore = "spawned by oauth_external_update_child"]
async fn oauth_external_delete_writer_child() -> anyhow::Result<()> {
let server_url = std::env::var(EXTERNAL_UPDATE_CHILD_SERVER_URL_ENV)?;
assert!(delete_oauth_tokens(
SERVER_NAME,
&server_url,
OAuthCredentialsStoreMode::File,
)?);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn refresh_holds_generation_lock_against_external_oauth_update() -> anyhow::Result<()> {
let server = MockServer::start().await;
let server_url = format!("{}/mcp", server.uri());
mount_oauth_metadata(&server).await;
Mock::given(method("POST"))
.and(path("/oauth/token"))
.and(body_string_contains("grant_type=refresh_token"))
.and(body_string_contains(format!(
"refresh_token={OLD_REFRESH_TOKEN}"
)))
.respond_with(
ResponseTemplate::new(200)
.set_delay(Duration::from_millis(1500))
.set_body_json(json!({
"access_token": STALE_REFRESHED_ACCESS_TOKEN,
"token_type": "Bearer",
"expires_in": 7200,
"refresh_token": STALE_ROTATED_REFRESH_TOKEN,
})),
)
.expect(1)
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/mcp"))
.respond_with(|request: &Request| {
let Ok(body) = request.body_json::<Value>() else {
return ResponseTemplate::new(400);
};
let request_method = body.get("method").and_then(Value::as_str);
let expected_access_token = match (
request_method,
body.pointer("/params/name").and_then(Value::as_str),
) {
(Some("initialize" | "notifications/initialized"), _) => OLD_ACCESS_TOKEN,
(Some("tools/call"), Some("trigger-refresh-contention")) => {
STALE_REFRESHED_ACCESS_TOKEN
}
(Some("tools/call"), Some("call-after-refresh-contention")) => {
EXTERNAL_ACCESS_TOKEN
}
_ => OLD_ACCESS_TOKEN,
};
oauth_mcp_response(request, expected_access_token)
})
.expect(4)
.mount(&server)
.await;
let codex_home = TempDir::new()?;
let status = Command::new(std::env::current_exe()?)
let attempt_path = codex_home.path().join("external-update-attempt");
let done_path = codex_home.path().join("external-update-done");
let mut live_client = Command::new(std::env::current_exe()?)
.args([
"oauth_external_update_during_refresh_child",
"--exact",
@@ -305,20 +221,104 @@ async fn refresh_does_not_overwrite_external_oauth_update() -> anyhow::Result<()
"--nocapture",
])
.env("CODEX_HOME", codex_home.path())
.status()
.await?;
.env(REFRESH_CONTENTION_SERVER_URL_ENV, &server_url)
.env(REFRESH_CONTENTION_DONE_ENV, &done_path)
.spawn()?;
wait_for_request_path(&server, "/oauth/token").await?;
let mut external_writer = Command::new(std::env::current_exe()?)
.args([
"oauth_external_update_writer_child",
"--exact",
"--ignored",
"--nocapture",
])
.env("CODEX_HOME", codex_home.path())
.env(REFRESH_CONTENTION_SERVER_URL_ENV, &server_url)
.env(REFRESH_CONTENTION_ATTEMPT_ENV, &attempt_path)
.env(REFRESH_CONTENTION_DONE_ENV, &done_path)
.spawn()?;
wait_for_path(&attempt_path).await?;
tokio::time::sleep(Duration::from_millis(100)).await;
assert!(
status.success(),
"OAuth external update during refresh child failed: {status}"
!done_path.exists(),
"external writer should block while refresh owns the generation lock"
);
assert!(
external_writer.try_wait()?.is_none(),
"external writer exited before the refresh released its generation lock"
);
let writer_status = external_writer.wait().await?;
assert!(
writer_status.success(),
"OAuth external update writer failed: {writer_status}"
);
let live_status = live_client.wait().await?;
assert!(
live_status.success(),
"OAuth refresh contention child failed: {live_status}"
);
server.verify().await;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
#[ignore = "spawned by refresh_does_not_overwrite_external_oauth_update"]
#[ignore = "spawned by refresh_holds_generation_lock_against_external_oauth_update"]
async fn oauth_external_update_during_refresh_child() -> anyhow::Result<()> {
let server = MockServer::start().await;
let server_url = format!("{}/mcp", server.uri());
let server_url = std::env::var(REFRESH_CONTENTION_SERVER_URL_ENV)?;
let done_path = std::path::PathBuf::from(std::env::var(REFRESH_CONTENTION_DONE_ENV)?);
save_test_tokens(
&server_url,
OLD_ACCESS_TOKEN,
OLD_REFRESH_TOKEN,
future_expiry()?,
)?;
let client = initialized_oauth_client(&server_url).await?;
save_test_tokens(&server_url, OLD_ACCESS_TOKEN, OLD_REFRESH_TOKEN, 0)?;
client
.call_tool(
"trigger-refresh-contention".to_string(),
/*arguments*/ None,
/*meta*/ None,
Some(Duration::from_secs(5)),
)
.await?;
wait_for_path(&done_path).await?;
client
.call_tool(
"call-after-refresh-contention".to_string(),
/*arguments*/ None,
/*meta*/ None,
Some(Duration::from_secs(5)),
)
.await?;
assert_persisted_tokens(EXTERNAL_ACCESS_TOKEN, EXTERNAL_REFRESH_TOKEN)?;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
#[ignore = "spawned by refresh_holds_generation_lock_against_external_oauth_update"]
async fn oauth_external_update_writer_child() -> anyhow::Result<()> {
let server_url = std::env::var(REFRESH_CONTENTION_SERVER_URL_ENV)?;
let attempt_path = std::env::var(REFRESH_CONTENTION_ATTEMPT_ENV)?;
let done_path = std::env::var(REFRESH_CONTENTION_DONE_ENV)?;
std::fs::write(attempt_path, b"attempting")?;
save_test_tokens(
&server_url,
EXTERNAL_ACCESS_TOKEN,
EXTERNAL_REFRESH_TOKEN,
future_expiry()?,
)?;
std::fs::write(done_path, b"done")?;
Ok(())
}
async fn mount_oauth_metadata(server: &MockServer) {
Mock::given(method("GET"))
.and(path("/.well-known/oauth-authorization-server/mcp"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
@@ -327,91 +327,14 @@ async fn oauth_external_update_during_refresh_child() -> anyhow::Result<()> {
"scopes_supported": [""],
})))
.expect(1)
.mount(&server)
.mount(server)
.await;
Mock::given(method("POST"))
.and(path("/oauth/token"))
.and(body_string_contains("grant_type=refresh_token"))
.and(body_string_contains(format!(
"refresh_token={OLD_REFRESH_TOKEN}"
)))
.respond_with({
let server_url = server_url.clone();
move |_request: &Request| {
save_test_tokens(
&server_url,
EXTERNAL_ACCESS_TOKEN,
EXTERNAL_REFRESH_TOKEN,
future_expiry().expect("future expiry"),
)
.expect("save externally updated credentials");
ResponseTemplate::new(200).set_body_json(json!({
"access_token": STALE_REFRESHED_ACCESS_TOKEN,
"token_type": "Bearer",
"expires_in": 7200,
"refresh_token": STALE_ROTATED_REFRESH_TOKEN,
}))
}
})
.expect(1)
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/mcp"))
.respond_with(|request: &Request| {
let body: Value = request.body_json().expect("valid JSON-RPC request");
let request_method = body.get("method").and_then(Value::as_str);
let expected_access_token = match request_method {
Some("tools/call") => EXTERNAL_ACCESS_TOKEN,
_ => STALE_REFRESHED_ACCESS_TOKEN,
};
let authorization = request
.headers
.get("authorization")
.and_then(|value| value.to_str().ok());
if authorization != Some(format!("Bearer {expected_access_token}").as_str()) {
return ResponseTemplate::new(401);
}
}
match request_method {
Some("initialize") => ResponseTemplate::new(200).set_body_json(json!({
"jsonrpc": "2.0",
"id": body.get("id").cloned().unwrap_or(Value::Null),
"result": {
"protocolVersion": body
.pointer("/params/protocolVersion")
.cloned()
.unwrap_or_else(|| json!("2025-06-18")),
"capabilities": {
"tools": {},
},
"serverInfo": {
"name": "oauth-external-refresh-test",
"version": "0.0.0-test",
},
},
})),
Some("notifications/initialized") => ResponseTemplate::new(202),
Some("tools/call") => ResponseTemplate::new(200).set_body_json(json!({
"jsonrpc": "2.0",
"id": body.get("id").cloned().unwrap_or(Value::Null),
"result": {
"content": [],
"isError": false,
},
})),
method => ResponseTemplate::new(400)
.set_body_string(format!("unexpected JSON-RPC method: {method:?}")),
}
})
.expect(3)
.mount(&server)
.await;
save_test_tokens(&server_url, OLD_ACCESS_TOKEN, OLD_REFRESH_TOKEN, 0)?;
async fn initialized_oauth_client(server_url: &str) -> anyhow::Result<RmcpClient> {
let client = RmcpClient::new_streamable_http_client(
SERVER_NAME,
&server_url,
server_url,
/*bearer_token*/ None,
/*http_headers*/ None,
/*env_http_headers*/ None,
@@ -421,17 +344,79 @@ async fn oauth_external_update_during_refresh_child() -> anyhow::Result<()> {
)
.await?;
initialize_client(&client).await?;
client
.call_tool(
"call-after-refresh-race".to_string(),
/*arguments*/ None,
/*meta*/ None,
Some(Duration::from_secs(5)),
)
.await?;
Ok(client)
}
assert_persisted_tokens(EXTERNAL_ACCESS_TOKEN, EXTERNAL_REFRESH_TOKEN)?;
server.verify().await;
fn oauth_mcp_response(request: &Request, expected_access_token: &str) -> ResponseTemplate {
let authorization = request
.headers
.get("authorization")
.and_then(|value| value.to_str().ok());
if authorization != Some(format!("Bearer {expected_access_token}").as_str()) {
return ResponseTemplate::new(401);
}
let Ok(body) = request.body_json::<Value>() else {
return ResponseTemplate::new(400);
};
match body.get("method").and_then(Value::as_str) {
Some("initialize") => ResponseTemplate::new(200).set_body_json(json!({
"jsonrpc": "2.0",
"id": body.get("id").cloned().unwrap_or(Value::Null),
"result": {
"protocolVersion": body
.pointer("/params/protocolVersion")
.cloned()
.unwrap_or_else(|| json!("2025-06-18")),
"capabilities": {
"tools": {},
},
"serverInfo": {
"name": "oauth-external-credential-test",
"version": "0.0.0-test",
},
},
})),
Some("notifications/initialized") => ResponseTemplate::new(202),
Some("tools/call") => ResponseTemplate::new(200).set_body_json(json!({
"jsonrpc": "2.0",
"id": body.get("id").cloned().unwrap_or(Value::Null),
"result": {
"content": [],
"isError": false,
},
})),
method => ResponseTemplate::new(400)
.set_body_string(format!("unexpected JSON-RPC method: {method:?}")),
}
}
async fn wait_for_request_path(server: &MockServer, expected_path: &str) -> anyhow::Result<()> {
tokio::time::timeout(Duration::from_secs(5), async {
loop {
if server
.received_requests()
.await
.unwrap_or_default()
.iter()
.any(|request| request.url.path() == expected_path)
{
return;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await?;
Ok(())
}
async fn wait_for_path(path: &std::path::Path) -> anyhow::Result<()> {
tokio::time::timeout(Duration::from_secs(5), async {
while !path.exists() {
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await?;
Ok(())
}