mirror of
https://github.com/openai/codex.git
synced 2026-10-26 18:38:47 +00:00
fix(rmcp-client): retain failed OAuth refresh barriers
This commit is contained in:
Adam Perry
@@ -41,6 +41,7 @@ use std::io::ErrorKind;
|
||||
use std::path::Path;
|
||||
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;
|
||||
@@ -399,12 +400,22 @@ struct OAuthPersistorInner {
|
||||
credential_store: InMemoryCredentialStore,
|
||||
store_mode: OAuthCredentialsStoreMode,
|
||||
last_credentials: Mutex<Option<StoredOAuthTokens>>,
|
||||
refresh_attempt: Arc<Mutex<()>>,
|
||||
refresh_persistence_failure: Mutex<Option<RefreshPersistenceFailure>>,
|
||||
}
|
||||
|
||||
struct RefreshPersistenceFailure {
|
||||
message: String,
|
||||
_refresh_locks: Vec<fs::File>,
|
||||
}
|
||||
|
||||
static FAILED_REFRESH_LOCKS: OnceLock<std::sync::Mutex<Vec<fs::File>>> = OnceLock::new();
|
||||
|
||||
fn retain_failed_refresh_locks(refresh_locks: Vec<fs::File>) {
|
||||
FAILED_REFRESH_LOCKS
|
||||
.get_or_init(std::sync::Mutex::default)
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.extend(refresh_locks);
|
||||
}
|
||||
|
||||
impl OAuthPersistor {
|
||||
@@ -424,6 +435,7 @@ impl OAuthPersistor {
|
||||
credential_store,
|
||||
store_mode,
|
||||
last_credentials: Mutex::new(initial_credentials),
|
||||
refresh_attempt: Arc::new(Mutex::new(())),
|
||||
refresh_persistence_failure: Mutex::new(None),
|
||||
}),
|
||||
}
|
||||
@@ -505,6 +517,7 @@ impl OAuthPersistor {
|
||||
reason = "AuthorizationManager async access must be serialized through its mutex"
|
||||
)]
|
||||
pub(crate) async fn refresh_if_needed(&self) -> Result<()> {
|
||||
let refresh_attempt_guard = Arc::clone(&self.inner.refresh_attempt).lock_owned().await;
|
||||
let previous_failure = self
|
||||
.inner
|
||||
.refresh_persistence_failure
|
||||
@@ -578,6 +591,7 @@ impl OAuthPersistor {
|
||||
|
||||
let persistor = self.clone();
|
||||
tokio::spawn(async move {
|
||||
let _refresh_attempt_guard = refresh_attempt_guard;
|
||||
let manager = persistor.inner.authorization_manager.clone();
|
||||
let guard = manager.lock().await;
|
||||
match guard.refresh_token().await {
|
||||
@@ -604,10 +618,10 @@ impl OAuthPersistor {
|
||||
persistor.inner.server_name
|
||||
);
|
||||
warn!("{message}");
|
||||
retain_failed_refresh_locks(refresh_locks);
|
||||
*persistor.inner.refresh_persistence_failure.lock().await =
|
||||
Some(RefreshPersistenceFailure {
|
||||
message: message.clone(),
|
||||
_refresh_locks: refresh_locks,
|
||||
});
|
||||
return Err(anyhow!(message));
|
||||
}
|
||||
|
||||
@@ -966,18 +966,8 @@ 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 {
|
||||
if let Err(error) = runtime.persist_if_needed().await {
|
||||
warn!("{error_context}: {error}");
|
||||
}
|
||||
});
|
||||
@@ -1129,13 +1119,83 @@ async fn create_oauth_transport_and_runtime(
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::ffi::OsString;
|
||||
use std::fs::OpenOptions;
|
||||
use std::path::Path;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use std::time::SystemTime;
|
||||
use std::time::UNIX_EPOCH;
|
||||
|
||||
use anyhow::Context;
|
||||
use codex_exec_server::Environment;
|
||||
use oauth2::AccessToken;
|
||||
use oauth2::RefreshToken;
|
||||
use oauth2::basic::BasicTokenType;
|
||||
use pretty_assertions::assert_eq;
|
||||
use tokio::sync::oneshot;
|
||||
use rmcp::model::ClientCapabilities;
|
||||
use rmcp::model::Implementation;
|
||||
use rmcp::model::ProtocolVersion;
|
||||
use rmcp::transport::auth::CredentialStore;
|
||||
use rmcp::transport::auth::OAuthTokenResponse;
|
||||
use rmcp::transport::auth::StoredCredentials;
|
||||
use rmcp::transport::auth::VendorExtraTokenFields;
|
||||
use serde_json::Value;
|
||||
use serde_json::json;
|
||||
use sha2::Digest;
|
||||
use sha2::Sha256;
|
||||
use tempfile::TempDir;
|
||||
use tokio::time;
|
||||
use wiremock::Mock;
|
||||
use wiremock::MockServer;
|
||||
use wiremock::Request;
|
||||
use wiremock::ResponseTemplate;
|
||||
use wiremock::matchers::method;
|
||||
use wiremock::matchers::path;
|
||||
|
||||
use super::*;
|
||||
use crate::oauth::WrappedOAuthTokenResponse;
|
||||
use crate::oauth::load_oauth_tokens;
|
||||
use crate::oauth::save_oauth_tokens;
|
||||
|
||||
const OPERATION_TEST_SERVER_NAME: &str = "test-completed-operation-persistence";
|
||||
const OPERATION_TEST_CLIENT_ID: &str = "test-operation-client-id";
|
||||
const OPERATION_TEST_OLD_ACCESS_TOKEN: &str = "old-operation-access-token";
|
||||
const OPERATION_TEST_OLD_REFRESH_TOKEN: &str = "old-operation-refresh-token";
|
||||
const OPERATION_TEST_NEW_ACCESS_TOKEN: &str = "new-operation-access-token";
|
||||
const OPERATION_TEST_NEW_REFRESH_TOKEN: &str = "new-operation-refresh-token";
|
||||
|
||||
struct CodexHomeGuard {
|
||||
original: Option<OsString>,
|
||||
}
|
||||
|
||||
impl CodexHomeGuard {
|
||||
fn set(path: &Path) -> Self {
|
||||
let original = std::env::var_os("CODEX_HOME");
|
||||
// SAFETY: this test is serialized and restores the process environment on drop.
|
||||
unsafe {
|
||||
std::env::set_var("CODEX_HOME", path);
|
||||
}
|
||||
Self { original }
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for CodexHomeGuard {
|
||||
fn drop(&mut self) {
|
||||
if let Some(original) = &self.original {
|
||||
// SAFETY: this test is serialized and restores the prior value.
|
||||
unsafe {
|
||||
std::env::set_var("CODEX_HOME", original);
|
||||
}
|
||||
} else {
|
||||
// SAFETY: this test is serialized and removes only the value it set.
|
||||
unsafe {
|
||||
std::env::remove_var("CODEX_HOME");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn active_time_timeout_pauses_while_elicitation_is_pending() {
|
||||
@@ -1156,36 +1216,273 @@ 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();
|
||||
#[cfg(unix)]
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial_test::serial]
|
||||
async fn completed_operation_returns_while_oauth_persistence_is_contended() -> 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("/mcp"))
|
||||
.respond_with(|request: &Request| {
|
||||
let Ok(body) = request.body_json::<Value>() else {
|
||||
return ResponseTemplate::new(400).set_body_string("invalid 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": "2025-06-18",
|
||||
"capabilities": {},
|
||||
"serverInfo": {
|
||||
"name": "completed-operation-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;
|
||||
|
||||
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 codex_home = TempDir::new()?;
|
||||
let _codex_home_guard = CodexHomeGuard::set(codex_home.path());
|
||||
let server_url = format!("{}/mcp", server.uri());
|
||||
let old_response = operation_test_token_response(
|
||||
OPERATION_TEST_OLD_ACCESS_TOKEN,
|
||||
OPERATION_TEST_OLD_REFRESH_TOKEN,
|
||||
);
|
||||
let completed_result = "completed tool response";
|
||||
let initial_tokens = StoredOAuthTokens {
|
||||
server_name: OPERATION_TEST_SERVER_NAME.to_string(),
|
||||
url: server_url.clone(),
|
||||
client_id: OPERATION_TEST_CLIENT_ID.to_string(),
|
||||
token_response: WrappedOAuthTokenResponse(old_response.clone()),
|
||||
expires_at: Some(unix_millis().saturating_add(3_600_000)),
|
||||
};
|
||||
save_oauth_tokens(
|
||||
OPERATION_TEST_SERVER_NAME,
|
||||
&initial_tokens,
|
||||
OAuthCredentialsStoreMode::File,
|
||||
)?;
|
||||
|
||||
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");
|
||||
let client = RmcpClient::new_streamable_http_client(
|
||||
OPERATION_TEST_SERVER_NAME,
|
||||
&server_url,
|
||||
Some("test-bearer".to_string()),
|
||||
/*http_headers*/ None,
|
||||
/*env_http_headers*/ None,
|
||||
OAuthCredentialsStoreMode::File,
|
||||
Environment::default_for_tests().get_http_client(),
|
||||
/*auth_provider*/ None,
|
||||
)
|
||||
.await?;
|
||||
client
|
||||
.initialize(
|
||||
InitializeRequestParams::new(
|
||||
ClientCapabilities::default(),
|
||||
Implementation::new("codex-test", "0.0.0-test"),
|
||||
)
|
||||
.with_protocol_version(ProtocolVersion::V_2025_06_18),
|
||||
Some(Duration::from_secs(2)),
|
||||
Box::new(|_, _| {
|
||||
async {
|
||||
Ok(ElicitationResponse {
|
||||
action: ElicitationAction::Accept,
|
||||
content: Some(json!({})),
|
||||
meta: None,
|
||||
})
|
||||
}
|
||||
.boxed()
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
let service = client.service().await?;
|
||||
|
||||
let credential_store = InMemoryCredentialStore::new();
|
||||
let mut oauth_state =
|
||||
OAuthState::new(server_url.clone(), Some(reqwest::Client::new())).await?;
|
||||
match &mut oauth_state {
|
||||
OAuthState::Unauthorized(manager) => {
|
||||
manager.set_credential_store(credential_store.clone());
|
||||
}
|
||||
_ => {
|
||||
return Err(anyhow!(
|
||||
"unexpected OAuth state during operation test setup"
|
||||
));
|
||||
}
|
||||
}
|
||||
oauth_state
|
||||
.set_credentials(OPERATION_TEST_CLIENT_ID, old_response)
|
||||
.await?;
|
||||
let authorization_manager = match oauth_state {
|
||||
OAuthState::Authorized(manager) => manager,
|
||||
_ => return Err(anyhow!("OAuth operation test setup did not authorize")),
|
||||
};
|
||||
let oauth_persistor = OAuthPersistor::new(
|
||||
OPERATION_TEST_SERVER_NAME.to_string(),
|
||||
server_url.clone(),
|
||||
Arc::new(Mutex::new(authorization_manager)),
|
||||
credential_store.clone(),
|
||||
OAuthCredentialsStoreMode::File,
|
||||
Some(initial_tokens.clone()),
|
||||
);
|
||||
|
||||
let fallback_lock = OpenOptions::new()
|
||||
.read(true)
|
||||
.write(true)
|
||||
.create(true)
|
||||
.truncate(false)
|
||||
.open(operation_test_fallback_lock_path(codex_home.path()))?;
|
||||
fallback_lock.lock_shared()?;
|
||||
|
||||
let new_response = operation_test_token_response(
|
||||
OPERATION_TEST_NEW_ACCESS_TOKEN,
|
||||
OPERATION_TEST_NEW_REFRESH_TOKEN,
|
||||
);
|
||||
let operation_store = credential_store.clone();
|
||||
let operation_response = new_response.clone();
|
||||
let operation = move |_service| {
|
||||
let operation_store = operation_store.clone();
|
||||
let operation_response = operation_response.clone();
|
||||
async move {
|
||||
if let Err(error) = operation_store
|
||||
.save(StoredCredentials::new(
|
||||
OPERATION_TEST_CLIENT_ID.to_string(),
|
||||
Some(operation_response),
|
||||
Vec::new(),
|
||||
Some(unix_seconds()),
|
||||
))
|
||||
.await
|
||||
{
|
||||
panic!("failed to simulate completed OAuth side effect: {error}");
|
||||
}
|
||||
Ok::<_, rmcp::service::ServiceError>("side effect completed".to_string())
|
||||
}
|
||||
.boxed()
|
||||
};
|
||||
|
||||
let result = RmcpClient::run_service_operation_once(
|
||||
service,
|
||||
Some(oauth_persistor),
|
||||
"tools/call",
|
||||
Some(Duration::from_millis(100)),
|
||||
ElicitationPauseState::new(),
|
||||
&operation,
|
||||
)
|
||||
.await?;
|
||||
assert_eq!(result, "side effect completed");
|
||||
time::sleep(Duration::from_millis(200)).await;
|
||||
let persisted_while_locked = load_oauth_tokens(
|
||||
OPERATION_TEST_SERVER_NAME,
|
||||
&server_url,
|
||||
OAuthCredentialsStoreMode::File,
|
||||
)?
|
||||
.context("missing persisted OAuth credentials while fallback lock is held")?;
|
||||
assert_eq!(
|
||||
(
|
||||
persisted_while_locked
|
||||
.token_response
|
||||
.0
|
||||
.access_token()
|
||||
.secret()
|
||||
.as_str(),
|
||||
persisted_while_locked
|
||||
.token_response
|
||||
.0
|
||||
.refresh_token()
|
||||
.map(|token| token.secret().as_str()),
|
||||
),
|
||||
(
|
||||
OPERATION_TEST_OLD_ACCESS_TOKEN,
|
||||
Some(OPERATION_TEST_OLD_REFRESH_TOKEN),
|
||||
)
|
||||
);
|
||||
|
||||
drop(fallback_lock);
|
||||
let persisted = time::timeout(Duration::from_secs(2), async {
|
||||
loop {
|
||||
let loaded = load_oauth_tokens(
|
||||
OPERATION_TEST_SERVER_NAME,
|
||||
&server_url,
|
||||
OAuthCredentialsStoreMode::File,
|
||||
)?;
|
||||
if loaded.as_ref().is_some_and(|tokens| {
|
||||
tokens.token_response.0.access_token().secret()
|
||||
== OPERATION_TEST_NEW_ACCESS_TOKEN
|
||||
}) {
|
||||
return Ok::<_, anyhow::Error>(loaded);
|
||||
}
|
||||
time::sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.context("timed out waiting for post-operation OAuth persistence")??;
|
||||
assert_eq!(
|
||||
persisted.as_ref().map(|tokens| tokens
|
||||
.token_response
|
||||
.0
|
||||
.access_token()
|
||||
.secret()
|
||||
.as_str()),
|
||||
Some(OPERATION_TEST_NEW_ACCESS_TOKEN)
|
||||
);
|
||||
|
||||
server.verify().await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn operation_test_token_response(
|
||||
access_token: &str,
|
||||
refresh_token: &str,
|
||||
) -> OAuthTokenResponse {
|
||||
let mut response = OAuthTokenResponse::new(
|
||||
AccessToken::new(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)));
|
||||
response
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn operation_test_fallback_lock_path(codex_home: &Path) -> PathBuf {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(codex_home.as_os_str().to_string_lossy().as_bytes());
|
||||
let digest = format!("{:x}", hasher.finalize());
|
||||
// SAFETY: getuid has no preconditions and returns the real UID.
|
||||
let user_namespace = format!("uid-{}", unsafe { libc::getuid() });
|
||||
PathBuf::from("/tmp").join(format!(
|
||||
"codex-mcp-oauth-fallback-{user_namespace}-{}.lock",
|
||||
&digest[..16]
|
||||
))
|
||||
}
|
||||
|
||||
fn unix_millis() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_millis() as u64
|
||||
}
|
||||
|
||||
fn unix_seconds() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -63,6 +63,7 @@ const SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_RACE_SERVER_URL";
|
||||
const READY_PATH_ENV: &str = "MCP_TEST_OAUTH_RACE_READY_PATH";
|
||||
const GO_PATH_ENV: &str = "MCP_TEST_OAUTH_RACE_GO_PATH";
|
||||
const RESULT_PATH_ENV: &str = "MCP_TEST_OAUTH_RACE_RESULT_PATH";
|
||||
const ATTEMPT_PATH_ENV: &str = "MCP_TEST_OAUTH_RACE_ATTEMPT_PATH";
|
||||
const ACCESS_TOKEN_ENV: &str = "MCP_TEST_OAUTH_RACE_ACCESS_TOKEN";
|
||||
const REFRESH_TOKEN_ENV: &str = "MCP_TEST_OAUTH_RACE_REFRESH_TOKEN";
|
||||
const EXPIRES_AT_ENV: &str = "MCP_TEST_OAUTH_RACE_EXPIRES_AT";
|
||||
@@ -510,6 +511,101 @@ async fn provider_refresh_persists_after_startup_timeout_before_next_refresh() -
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
|
||||
async fn canceled_startup_persistence_failure_retains_lock_until_process_exit() -> 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_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,
|
||||
})),
|
||||
_ => 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*/ 0).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 credentials_path = codex_home.path().join(".credentials.json");
|
||||
fs::set_permissions(&credentials_path, fs::Permissions::from_mode(0o400))?;
|
||||
|
||||
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"
|
||||
);
|
||||
sleep(Duration::from_secs(2)).await;
|
||||
assert_eq!(refresh_attempts.load(Ordering::SeqCst), 1);
|
||||
|
||||
let second_ready_path = control_dir.path().join("client-b.ready");
|
||||
let second_go_path = control_dir.path().join("client-b.go");
|
||||
let second_attempt_path = control_dir.path().join("client-b.attempt");
|
||||
let second_result_path = control_dir.path().join("client-b.result");
|
||||
let second_child = spawn_client_child_with_attempt(
|
||||
&codex_home,
|
||||
&server_url,
|
||||
&second_ready_path,
|
||||
&second_go_path,
|
||||
&second_attempt_path,
|
||||
&second_result_path,
|
||||
)?;
|
||||
wait_for_paths(std::slice::from_ref(&second_ready_path)).await?;
|
||||
fs::write(&second_go_path, "go")?;
|
||||
wait_for_paths(std::slice::from_ref(&second_attempt_path)).await?;
|
||||
sleep(Duration::from_millis(250)).await;
|
||||
assert_eq!(refresh_attempts.load(Ordering::SeqCst), 1);
|
||||
assert!(!second_result_path.exists());
|
||||
|
||||
fs::set_permissions(&credentials_path, fs::Permissions::from_mode(0o600))?;
|
||||
fs::write(first_release_path, "release")?;
|
||||
wait_for_children(vec![first_child]).await?;
|
||||
wait_for_children(vec![second_child]).await?;
|
||||
assert_eq!(
|
||||
fs::read_to_string(&second_result_path)?,
|
||||
"error:Auth required: OAuth refresh token was rejected; reauthentication required"
|
||||
);
|
||||
assert_eq!(refresh_attempts.load(Ordering::SeqCst), 2);
|
||||
|
||||
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;
|
||||
@@ -831,7 +927,7 @@ async fn rotated_token_storage_failure_blocks_another_process_refresh() -> anyho
|
||||
})),
|
||||
}
|
||||
})
|
||||
.expect(2)
|
||||
.expect(3)
|
||||
.mount(&server)
|
||||
.await;
|
||||
mount_mcp_server(&server, /*expected_requests*/ 2).await;
|
||||
@@ -871,27 +967,107 @@ async fn rotated_token_storage_failure_blocks_another_process_refresh() -> anyho
|
||||
|
||||
let second_ready_path = control_dir.path().join("client-b.ready");
|
||||
let second_go_path = control_dir.path().join("client-b.go");
|
||||
let second_attempt_path = control_dir.path().join("client-b.attempt");
|
||||
let second_result_path = control_dir.path().join("client-b.result");
|
||||
let mut second_child = spawn_client_child(
|
||||
let second_child = spawn_client_child_with_attempt(
|
||||
&codex_home,
|
||||
&server_url,
|
||||
&second_ready_path,
|
||||
&second_go_path,
|
||||
&second_attempt_path,
|
||||
&second_result_path,
|
||||
)?;
|
||||
wait_for_paths(std::slice::from_ref(&second_ready_path)).await?;
|
||||
fs::write(&second_go_path, "go")?;
|
||||
wait_for_paths(std::slice::from_ref(&second_attempt_path)).await?;
|
||||
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?;
|
||||
wait_for_children(vec![second_child]).await?;
|
||||
assert_eq!(
|
||||
fs::read_to_string(&second_result_path)?,
|
||||
"error:Auth required: OAuth refresh token was rejected; reauthentication required"
|
||||
);
|
||||
assert_eq!(refresh_attempts.load(Ordering::SeqCst), 3);
|
||||
|
||||
server.verify().await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
|
||||
async fn concurrent_waiter_observes_recorded_persistence_failure() -> anyhow::Result<()> {
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
|
||||
let server = MockServer::start().await;
|
||||
mount_oauth_metadata(&server, /*expected_requests*/ 1).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,
|
||||
})),
|
||||
_ => ResponseTemplate::new(200).set_body_json(json!({
|
||||
"access_token": FINAL_ACCESS_TOKEN,
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 7200,
|
||||
"refresh_token": FINAL_REFRESH_TOKEN,
|
||||
})),
|
||||
}
|
||||
})
|
||||
.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 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_concurrent_waiter_child(
|
||||
&codex_home,
|
||||
&server_url,
|
||||
&ready_path,
|
||||
&go_path,
|
||||
&result_path,
|
||||
)?;
|
||||
wait_for_paths(std::slice::from_ref(&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(&go_path, "go")?;
|
||||
wait_for_children(vec![child]).await?;
|
||||
fs::set_permissions(&credentials_path, fs::Permissions::from_mode(0o600))?;
|
||||
|
||||
let outcomes = fs::read_to_string(result_path)?
|
||||
.lines()
|
||||
.map(str::to_string)
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(outcomes.len(), 2);
|
||||
assert_eq!(outcomes[0], outcomes[1]);
|
||||
assert!(
|
||||
outcomes[0].starts_with(&format!(
|
||||
"error:OAuth refresh succeeded but failed to persist rotated credentials for server {SERVER_NAME}:"
|
||||
)),
|
||||
"unexpected concurrent waiter outcome: {}",
|
||||
outcomes[0]
|
||||
);
|
||||
assert_eq!(refresh_attempts.load(Ordering::SeqCst), 2);
|
||||
|
||||
server.verify().await;
|
||||
Ok(())
|
||||
@@ -1239,6 +1415,9 @@ async fn oauth_refresh_race_client_child() -> anyhow::Result<()> {
|
||||
|
||||
fs::write(ready_path, "ready")?;
|
||||
wait_for_paths(std::slice::from_ref(&go_path)).await?;
|
||||
if let Ok(attempt_path) = std::env::var(ATTEMPT_PATH_ENV) {
|
||||
fs::write(attempt_path, "attempt")?;
|
||||
}
|
||||
let outcome = match initialize_client(&client).await {
|
||||
Ok(()) => "success".to_string(),
|
||||
Err(error) => format!("error:{error:#}"),
|
||||
@@ -1283,6 +1462,51 @@ async fn oauth_refresh_operation_timeout_child() -> anyhow::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
|
||||
#[ignore = "spawned by the OAuth refresh race integration tests"]
|
||||
async fn oauth_refresh_concurrent_waiter_child() -> anyhow::Result<()> {
|
||||
let server_url = std::env::var(SERVER_URL_ENV)?;
|
||||
let ready_path = PathBuf::from(std::env::var(READY_PATH_ENV)?);
|
||||
let go_path = PathBuf::from(std::env::var(GO_PATH_ENV)?);
|
||||
let result_path = PathBuf::from(std::env::var(RESULT_PATH_ENV)?);
|
||||
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?;
|
||||
|
||||
fs::write(ready_path, "ready")?;
|
||||
wait_for_paths(std::slice::from_ref(&go_path)).await?;
|
||||
let outcomes = timeout(Duration::from_secs(3), async {
|
||||
tokio::join!(
|
||||
client.list_tools(None, /*timeout*/ None),
|
||||
client.list_tools(None, /*timeout*/ None),
|
||||
)
|
||||
})
|
||||
.await
|
||||
.map(|(first, second)| {
|
||||
[first, second].map(|outcome| match outcome {
|
||||
Ok(_) => "success".to_string(),
|
||||
Err(error) => format!("error:{error:#}"),
|
||||
})
|
||||
})
|
||||
.unwrap_or_else(|_| {
|
||||
[
|
||||
"error:concurrent refresh waiters timed out".to_string(),
|
||||
"error:concurrent refresh waiters timed out".to_string(),
|
||||
]
|
||||
});
|
||||
fs::write(result_path, outcomes.join("\n"))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
|
||||
#[ignore = "spawned by the OAuth refresh race integration tests"]
|
||||
async fn oauth_refresh_read_only_child() -> anyhow::Result<()> {
|
||||
@@ -1441,6 +1665,32 @@ fn spawn_client_child_held_open(
|
||||
Ok(command.spawn()?)
|
||||
}
|
||||
|
||||
fn spawn_client_child_with_attempt(
|
||||
codex_home: &TempDir,
|
||||
server_url: &str,
|
||||
ready_path: &Path,
|
||||
go_path: &Path,
|
||||
attempt_path: &Path,
|
||||
result_path: &Path,
|
||||
) -> anyhow::Result<Child> {
|
||||
let mut command = Command::new(std::env::current_exe()?);
|
||||
command
|
||||
.args([
|
||||
"oauth_refresh_race_client_child",
|
||||
"--exact",
|
||||
"--ignored",
|
||||
"--nocapture",
|
||||
])
|
||||
.env("CODEX_HOME", codex_home.path())
|
||||
.env(SERVER_URL_ENV, server_url)
|
||||
.env(READY_PATH_ENV, ready_path)
|
||||
.env(GO_PATH_ENV, go_path)
|
||||
.env(ATTEMPT_PATH_ENV, attempt_path)
|
||||
.env(RESULT_PATH_ENV, result_path)
|
||||
.kill_on_drop(true);
|
||||
Ok(command.spawn()?)
|
||||
}
|
||||
|
||||
fn spawn_operation_timeout_child(
|
||||
codex_home: &TempDir,
|
||||
server_url: &str,
|
||||
@@ -1491,6 +1741,30 @@ fn spawn_operation_timeout_child_held_open(
|
||||
Ok(command.spawn()?)
|
||||
}
|
||||
|
||||
fn spawn_concurrent_waiter_child(
|
||||
codex_home: &TempDir,
|
||||
server_url: &str,
|
||||
ready_path: &Path,
|
||||
go_path: &Path,
|
||||
result_path: &Path,
|
||||
) -> anyhow::Result<Child> {
|
||||
let mut command = Command::new(std::env::current_exe()?);
|
||||
command
|
||||
.args([
|
||||
"oauth_refresh_concurrent_waiter_child",
|
||||
"--exact",
|
||||
"--ignored",
|
||||
"--nocapture",
|
||||
])
|
||||
.env("CODEX_HOME", codex_home.path())
|
||||
.env(SERVER_URL_ENV, server_url)
|
||||
.env(READY_PATH_ENV, ready_path)
|
||||
.env(GO_PATH_ENV, go_path)
|
||||
.env(RESULT_PATH_ENV, result_path)
|
||||
.kill_on_drop(true);
|
||||
Ok(command.spawn()?)
|
||||
}
|
||||
|
||||
fn spawn_read_only_refresh_child(
|
||||
codex_home: &TempDir,
|
||||
server_url: &str,
|
||||
|
||||
Reference in New Issue
Block a user