mirror of
https://github.com/openai/codex.git
synced 2026-11-05 20:10:33 +00:00
Address OAuth refresh lock review feedback
This commit is contained in:
@@ -16,6 +16,8 @@
|
||||
//!
|
||||
//! If the keyring is not available or fails, we fall back to CODEX_HOME/.credentials.json which is consistent with other coding CLI agents.
|
||||
|
||||
mod refresh_lock;
|
||||
|
||||
use anyhow::Context;
|
||||
use anyhow::Error;
|
||||
use anyhow::Result;
|
||||
@@ -41,8 +43,6 @@ use sha2::Digest;
|
||||
use sha2::Sha256;
|
||||
use std::collections::BTreeMap;
|
||||
use std::fs;
|
||||
use std::fs::File;
|
||||
use std::fs::OpenOptions;
|
||||
use std::io::ErrorKind;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
@@ -58,17 +58,14 @@ use rmcp::transport::auth::CredentialStore as _;
|
||||
use rmcp::transport::auth::InMemoryCredentialStore;
|
||||
use rmcp::transport::auth::StoredCredentials;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio::time::sleep;
|
||||
use tokio::time::timeout;
|
||||
|
||||
use self::refresh_lock::RefreshCredentialLock;
|
||||
use codex_utils_home_dir::find_codex_home;
|
||||
|
||||
const KEYRING_SERVICE: &str = "Codex MCP Credentials";
|
||||
const MCP_OAUTH_SECRET_PREFIX: &str = "MCP_OAUTH";
|
||||
const REFRESH_SKEW_MILLIS: u64 = 30_000;
|
||||
const REFRESH_LOCK_DIR: &str = "mcp-oauth-refresh-locks";
|
||||
const REFRESH_LOCK_ACQUIRE_TIMEOUT: Duration = Duration::from_secs(60);
|
||||
const REFRESH_LOCK_RETRY_SLEEP: Duration = Duration::from_millis(50);
|
||||
const REFRESH_REQUEST_TIMEOUT: Duration = Duration::from_secs(45);
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
@@ -832,61 +829,6 @@ impl OAuthPersistor {
|
||||
}
|
||||
}
|
||||
|
||||
struct RefreshCredentialLock {
|
||||
_file: File,
|
||||
}
|
||||
|
||||
impl RefreshCredentialLock {
|
||||
async fn acquire(store_key: &str) -> Result<Self> {
|
||||
Self::acquire_with_timeout(store_key, REFRESH_LOCK_ACQUIRE_TIMEOUT).await
|
||||
}
|
||||
|
||||
async fn acquire_with_timeout(store_key: &str, acquire_timeout: Duration) -> Result<Self> {
|
||||
let path = refresh_lock_path(store_key)?;
|
||||
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)
|
||||
.with_context(|| format!("failed to open OAuth refresh lock {}", path.display()))?;
|
||||
|
||||
// Bound every contender, but keep the acquired lock for the full provider request and
|
||||
// persistence transaction. Releasing it while awaiting the provider would allow concurrent
|
||||
// use of a rotating refresh token.
|
||||
match timeout(acquire_timeout, async {
|
||||
loop {
|
||||
match file.try_lock() {
|
||||
Ok(()) => return Ok(()),
|
||||
Err(std::fs::TryLockError::WouldBlock) => {
|
||||
sleep(REFRESH_LOCK_RETRY_SLEEP).await;
|
||||
}
|
||||
Err(error) => return Err(std::io::Error::from(error)),
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
{
|
||||
Ok(Ok(())) => {}
|
||||
Ok(Err(error)) => {
|
||||
return Err(error).with_context(|| {
|
||||
format!("failed to lock OAuth refresh lock {}", path.display())
|
||||
});
|
||||
}
|
||||
Err(_) => anyhow::bail!(
|
||||
"timed out after {acquire_timeout:?} waiting for OAuth refresh lock {}",
|
||||
path.display()
|
||||
),
|
||||
}
|
||||
|
||||
Ok(Self { _file: file })
|
||||
}
|
||||
}
|
||||
|
||||
#[expect(
|
||||
clippy::await_holding_invalid_type,
|
||||
reason = "AuthorizationManager async access must be serialized through its mutex"
|
||||
@@ -1111,21 +1053,6 @@ fn fallback_file_path() -> Result<PathBuf> {
|
||||
Ok(find_codex_home()?.join(FALLBACK_FILENAME).to_path_buf())
|
||||
}
|
||||
|
||||
fn refresh_lock_path(store_key: &str) -> Result<PathBuf> {
|
||||
// Credential coordination is deliberately scoped to the active CODEX_HOME, alongside File
|
||||
// and Secrets state. Coordinating the process-global Direct keyring across distinct homes
|
||||
// would require a separately defined global lock namespace and is outside this transaction.
|
||||
// TODO(stevenlee): define a safe per-user, cross-platform rendezvous before extending Direct
|
||||
// keyring coordination across distinct CODEX_HOME values.
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(store_key.as_bytes());
|
||||
let digest = hasher.finalize();
|
||||
Ok(find_codex_home()?
|
||||
.join(REFRESH_LOCK_DIR)
|
||||
.join(format!("{digest:x}.lock"))
|
||||
.to_path_buf())
|
||||
}
|
||||
|
||||
fn read_fallback_file() -> Result<Option<FallbackFile>> {
|
||||
let path = fallback_file_path()?;
|
||||
let contents = match fs::read_to_string(&path) {
|
||||
@@ -1479,35 +1406,6 @@ mod tests {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_lock_acquisition_times_out_without_stealing() -> Result<()> {
|
||||
let _env = TempCodexHome::new();
|
||||
let store_key = "test-store-key";
|
||||
let held_lock =
|
||||
RefreshCredentialLock::acquire_with_timeout(store_key, Duration::from_millis(100))
|
||||
.await?;
|
||||
|
||||
let error =
|
||||
match RefreshCredentialLock::acquire_with_timeout(store_key, Duration::from_millis(50))
|
||||
.await
|
||||
{
|
||||
Ok(_) => panic!("contending lock acquisition should time out"),
|
||||
Err(error) => error,
|
||||
};
|
||||
assert!(
|
||||
error
|
||||
.to_string()
|
||||
.contains("timed out after 50ms waiting for OAuth refresh lock"),
|
||||
"unexpected error: {error:#}"
|
||||
);
|
||||
|
||||
drop(held_lock);
|
||||
let _reacquired =
|
||||
RefreshCredentialLock::acquire_with_timeout(store_key, Duration::from_millis(100))
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_transaction_preserves_credentials_when_resolved_keyring_reread_fails()
|
||||
-> Result<()> {
|
||||
|
||||
@@ -0,0 +1,113 @@
|
||||
//! Cross-process serialization for one MCP OAuth credential's refresh transaction.
|
||||
//!
|
||||
//! The guard is intentionally acquired before the authoritative credential reread and retained
|
||||
//! through provider refresh and persistence. This prevents two processes from replaying the same
|
||||
//! rotating refresh token or observing a partially persisted transaction.
|
||||
|
||||
use anyhow::Context;
|
||||
use anyhow::Result;
|
||||
use codex_utils_home_dir::find_codex_home;
|
||||
use sha2::Digest;
|
||||
use sha2::Sha256;
|
||||
use std::fs;
|
||||
use std::fs::File;
|
||||
use std::fs::OpenOptions;
|
||||
use std::path::Path;
|
||||
use std::path::PathBuf;
|
||||
use std::time::Duration;
|
||||
use tokio::time::sleep;
|
||||
use tokio::time::timeout;
|
||||
|
||||
const REFRESH_LOCK_DIR: &str = "mcp-oauth-refresh-locks";
|
||||
const REFRESH_LOCK_ACQUIRE_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 60);
|
||||
const REFRESH_LOCK_RETRY_SLEEP: Duration = Duration::from_millis(/*millis*/ 50);
|
||||
// Keep this internal target stable so diagnostics and cross-process tests can distinguish actual
|
||||
// WouldBlock contention from a contender that merely started late and observed persisted tokens.
|
||||
const LOCK_CONTENTION_EVENT_TARGET: &str = "codex_rmcp_client::oauth::refresh_lock::contention";
|
||||
|
||||
pub(super) struct RefreshCredentialLock {
|
||||
_file: File,
|
||||
}
|
||||
|
||||
impl RefreshCredentialLock {
|
||||
pub(super) async fn acquire(store_key: &str) -> Result<Self> {
|
||||
let codex_home = find_codex_home()?;
|
||||
Self::acquire_in(&codex_home, store_key, REFRESH_LOCK_ACQUIRE_TIMEOUT).await
|
||||
}
|
||||
|
||||
async fn acquire_in(
|
||||
codex_home: &Path,
|
||||
store_key: &str,
|
||||
acquire_timeout: Duration,
|
||||
) -> Result<Self> {
|
||||
let path = refresh_lock_path(codex_home, store_key);
|
||||
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)
|
||||
.with_context(|| format!("failed to open OAuth refresh lock {}", path.display()))?;
|
||||
|
||||
// Bound every contender, but keep the acquired lock for the full provider request and
|
||||
// persistence transaction. Releasing it while awaiting the provider would allow concurrent
|
||||
// use of a rotating refresh token.
|
||||
let mut reported_contention = false;
|
||||
match timeout(acquire_timeout, async {
|
||||
loop {
|
||||
match file.try_lock() {
|
||||
Ok(()) => return Ok(()),
|
||||
Err(std::fs::TryLockError::WouldBlock) => {
|
||||
if !reported_contention {
|
||||
tracing::debug!(
|
||||
target: LOCK_CONTENTION_EVENT_TARGET,
|
||||
lock_path = %path.display(),
|
||||
"waiting for another process to finish refreshing MCP OAuth credentials"
|
||||
);
|
||||
reported_contention = true;
|
||||
}
|
||||
sleep(REFRESH_LOCK_RETRY_SLEEP).await;
|
||||
}
|
||||
Err(error) => return Err(std::io::Error::from(error)),
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
{
|
||||
Ok(Ok(())) => {}
|
||||
Ok(Err(error)) => {
|
||||
return Err(error).with_context(|| {
|
||||
format!("failed to lock OAuth refresh lock {}", path.display())
|
||||
});
|
||||
}
|
||||
Err(_) => anyhow::bail!(
|
||||
"timed out after {acquire_timeout:?} waiting for OAuth refresh lock {}",
|
||||
path.display()
|
||||
),
|
||||
}
|
||||
|
||||
Ok(Self { _file: file })
|
||||
}
|
||||
}
|
||||
|
||||
fn refresh_lock_path(codex_home: &Path, store_key: &str) -> PathBuf {
|
||||
// Credential coordination is deliberately scoped to the active CODEX_HOME, alongside File
|
||||
// and Secrets state. Coordinating the process-global Direct keyring across distinct homes
|
||||
// would require a separately defined global lock namespace and is outside this transaction.
|
||||
// TODO(stevenlee): define a safe per-user, cross-platform rendezvous before extending Direct
|
||||
// keyring coordination across distinct CODEX_HOME values.
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(store_key.as_bytes());
|
||||
let digest = hasher.finalize();
|
||||
codex_home
|
||||
.join(REFRESH_LOCK_DIR)
|
||||
.join(format!("{digest:x}.lock"))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "refresh_lock_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,42 @@
|
||||
use super::RefreshCredentialLock;
|
||||
use anyhow::Result;
|
||||
use std::time::Duration;
|
||||
use tempfile::tempdir;
|
||||
|
||||
#[tokio::test]
|
||||
async fn acquisition_times_out_without_stealing() -> Result<()> {
|
||||
let codex_home = tempdir()?;
|
||||
let store_key = "test-store-key";
|
||||
let held_lock = RefreshCredentialLock::acquire_in(
|
||||
codex_home.path(),
|
||||
store_key,
|
||||
Duration::from_millis(/*millis*/ 100),
|
||||
)
|
||||
.await?;
|
||||
|
||||
let error = match RefreshCredentialLock::acquire_in(
|
||||
codex_home.path(),
|
||||
store_key,
|
||||
Duration::from_millis(/*millis*/ 50),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(_) => panic!("contending lock acquisition should time out"),
|
||||
Err(error) => error,
|
||||
};
|
||||
assert!(
|
||||
error
|
||||
.to_string()
|
||||
.contains("timed out after 50ms waiting for OAuth refresh lock"),
|
||||
"unexpected error: {error:#}"
|
||||
);
|
||||
|
||||
drop(held_lock);
|
||||
let _reacquired = RefreshCredentialLock::acquire_in(
|
||||
codex_home.path(),
|
||||
store_key,
|
||||
Duration::from_millis(/*millis*/ 100),
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,11 +1,19 @@
|
||||
mod streamable_http_test_support;
|
||||
|
||||
use std::fs;
|
||||
use std::path::Path;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::AtomicUsize;
|
||||
use std::sync::atomic::Ordering;
|
||||
use std::time::Duration;
|
||||
use std::time::SystemTime;
|
||||
use std::time::UNIX_EPOCH;
|
||||
|
||||
use axum::Json;
|
||||
use axum::Router;
|
||||
use axum::extract::State;
|
||||
use axum::routing::post;
|
||||
use codex_config::types::AuthKeyringBackendKind;
|
||||
use codex_config::types::OAuthCredentialsStoreMode;
|
||||
use codex_exec_server::Environment;
|
||||
@@ -24,7 +32,18 @@ use rmcp::transport::auth::VendorExtraTokenFields;
|
||||
use serde_json::Value;
|
||||
use serde_json::json;
|
||||
use tempfile::TempDir;
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::process::Command;
|
||||
use tokio::sync::Notify;
|
||||
use tokio::sync::Semaphore;
|
||||
use tokio::task::JoinHandle;
|
||||
use tracing::Event;
|
||||
use tracing::Id;
|
||||
use tracing::Metadata;
|
||||
use tracing::Subscriber;
|
||||
use tracing::span::Attributes;
|
||||
use tracing::span::Record;
|
||||
use tracing::subscriber::Interest;
|
||||
use wiremock::Mock;
|
||||
use wiremock::MockServer;
|
||||
use wiremock::Request;
|
||||
@@ -42,13 +61,170 @@ const EXPIRED_ACCESS_TOKEN: &str = "expired-access-token";
|
||||
const REFRESH_TOKEN: &str = "valid-refresh-token";
|
||||
const REFRESHED_ACCESS_TOKEN: &str = "refreshed-access-token";
|
||||
const ROTATED_REFRESH_TOKEN: &str = "rotated-refresh-token";
|
||||
const CHILD_CONTENTION_FILE_ENV: &str = "MCP_TEST_OAUTH_STARTUP_CONTENTION_FILE";
|
||||
const CHILD_READY_FILE_ENV: &str = "MCP_TEST_OAUTH_STARTUP_READY_FILE";
|
||||
const CHILD_RELEASE_FILE_ENV: &str = "MCP_TEST_OAUTH_STARTUP_RELEASE_FILE";
|
||||
const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_STARTUP_SERVER_URL";
|
||||
// This mirrors the private event target in oauth::refresh_lock without exposing test-only crate API.
|
||||
const LOCK_CONTENTION_EVENT_TARGET: &str = "codex_rmcp_client::oauth::refresh_lock::contention";
|
||||
const UNREFRESHABLE_SERVER_URL: &str = "https://unrefreshable.example/mcp";
|
||||
const UNEXPIRED_SERVER_URL: &str = "https://unexpired.example/mcp";
|
||||
const REFRESHABLE_SERVER_URL: &str = "https://refreshable.example/mcp";
|
||||
|
||||
#[derive(Clone)]
|
||||
struct GatedRefreshState {
|
||||
request_count: Arc<AtomicUsize>,
|
||||
request_started: Arc<Notify>,
|
||||
response_release: Arc<Semaphore>,
|
||||
}
|
||||
|
||||
struct GatedRefreshServer {
|
||||
token_endpoint: String,
|
||||
state: GatedRefreshState,
|
||||
task: JoinHandle<()>,
|
||||
}
|
||||
|
||||
impl GatedRefreshServer {
|
||||
async fn start() -> anyhow::Result<Self> {
|
||||
let state = GatedRefreshState {
|
||||
request_count: Arc::new(AtomicUsize::new(/*v*/ 0)),
|
||||
request_started: Arc::new(Notify::new()),
|
||||
response_release: Arc::new(Semaphore::new(/*permits*/ 0)),
|
||||
};
|
||||
let router = Router::new()
|
||||
.route("/oauth/token", post(gated_refresh_response))
|
||||
.with_state(state.clone());
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await?;
|
||||
let address = listener.local_addr()?;
|
||||
let task = tokio::spawn(async move {
|
||||
if let Err(error) = axum::serve(listener, router).await {
|
||||
panic!("gated refresh server failed: {error}");
|
||||
}
|
||||
});
|
||||
Ok(Self {
|
||||
token_endpoint: format!("http://{address}/oauth/token"),
|
||||
state,
|
||||
task,
|
||||
})
|
||||
}
|
||||
|
||||
async fn wait_until_request_started(&self) -> anyhow::Result<()> {
|
||||
let notified = self.state.request_started.notified();
|
||||
if self.request_count() == 0 {
|
||||
tokio::time::timeout(Duration::from_secs(/*secs*/ 10), notified)
|
||||
.await
|
||||
.map_err(|_| anyhow::anyhow!("provider refresh request did not start"))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn release_responses(&self) {
|
||||
// Two permits also let the no-lock negative control exit cleanly after issuing two refresh
|
||||
// requests. The passing path consumes one permit because only the lock owner calls out.
|
||||
self.state.response_release.add_permits(/*n*/ 2);
|
||||
}
|
||||
|
||||
fn request_count(&self) -> usize {
|
||||
self.state.request_count.load(Ordering::SeqCst)
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for GatedRefreshServer {
|
||||
fn drop(&mut self) {
|
||||
self.task.abort();
|
||||
}
|
||||
}
|
||||
|
||||
async fn gated_refresh_response(
|
||||
State(state): State<GatedRefreshState>,
|
||||
body: String,
|
||||
) -> Json<Value> {
|
||||
assert!(body.contains("grant_type=refresh_token"));
|
||||
assert!(body.contains(&format!("refresh_token={REFRESH_TOKEN}")));
|
||||
state.request_count.fetch_add(/*val*/ 1, Ordering::SeqCst);
|
||||
state.request_started.notify_one();
|
||||
let Ok(permit) = state.response_release.acquire().await else {
|
||||
panic!("gated refresh server closed its response semaphore");
|
||||
};
|
||||
permit.forget();
|
||||
Json(json!({
|
||||
"access_token": REFRESHED_ACCESS_TOKEN,
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 7200,
|
||||
"refresh_token": ROTATED_REFRESH_TOKEN,
|
||||
}))
|
||||
}
|
||||
|
||||
struct LockContentionMarkerSubscriber {
|
||||
marker_file: PathBuf,
|
||||
}
|
||||
|
||||
impl Subscriber for LockContentionMarkerSubscriber {
|
||||
fn enabled(&self, metadata: &Metadata<'_>) -> bool {
|
||||
metadata.target() == LOCK_CONTENTION_EVENT_TARGET
|
||||
}
|
||||
|
||||
fn register_callsite(&self, metadata: &'static Metadata<'static>) -> Interest {
|
||||
if self.enabled(metadata) {
|
||||
Interest::always()
|
||||
} else {
|
||||
Interest::never()
|
||||
}
|
||||
}
|
||||
|
||||
fn max_level_hint(&self) -> Option<tracing::level_filters::LevelFilter> {
|
||||
Some(tracing::level_filters::LevelFilter::DEBUG)
|
||||
}
|
||||
|
||||
fn new_span(&self, _span: &Attributes<'_>) -> Id {
|
||||
// This subscriber enables only the contention event callsite, so it never observes spans.
|
||||
Id::from_u64(/*u*/ 1)
|
||||
}
|
||||
|
||||
fn record(&self, _span: &Id, _values: &Record<'_>) {}
|
||||
|
||||
fn record_follows_from(&self, _span: &Id, _follows: &Id) {}
|
||||
|
||||
fn event(&self, event: &Event<'_>) {
|
||||
if event.metadata().target() == LOCK_CONTENTION_EVENT_TARGET
|
||||
&& let Err(error) = fs::write(&self.marker_file, b"contended")
|
||||
{
|
||||
panic!("failed to write refresh-lock contention marker: {error}");
|
||||
}
|
||||
}
|
||||
|
||||
fn enter(&self, _span: &Id) {}
|
||||
|
||||
fn exit(&self, _span: &Id) {}
|
||||
}
|
||||
|
||||
async fn wait_for_marker(path: &Path, timeout_message: &str) -> anyhow::Result<()> {
|
||||
tokio::time::timeout(Duration::from_secs(/*secs*/ 10), async {
|
||||
while !path.exists() {
|
||||
tokio::time::sleep(Duration::from_millis(/*millis*/ 10)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.map_err(|_| anyhow::anyhow!("{timeout_message}"))
|
||||
}
|
||||
|
||||
fn oauth_concurrency_child_command(
|
||||
codex_home: &Path,
|
||||
server_url: &str,
|
||||
ready_file: &Path,
|
||||
release_file: &Path,
|
||||
) -> anyhow::Result<Command> {
|
||||
let mut command = Command::new(std::env::current_exe()?);
|
||||
command
|
||||
.args(["oauth_concurrency_client_child", "--exact", "--ignored"])
|
||||
.env("CODEX_HOME", codex_home)
|
||||
.env(CHILD_SERVER_URL_ENV, server_url)
|
||||
.env(CHILD_READY_FILE_ENV, ready_file)
|
||||
.env(CHILD_RELEASE_FILE_ENV, release_file)
|
||||
.kill_on_drop(true);
|
||||
Ok(command)
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
|
||||
async fn startup_refresh_does_not_consume_handshake_timeout() -> anyhow::Result<()> {
|
||||
let server = MockServer::start().await;
|
||||
@@ -133,39 +309,21 @@ async fn startup_refresh_does_not_consume_handshake_timeout() -> anyhow::Result<
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn concurrent_file_mode_startup_refreshes_once() -> anyhow::Result<()> {
|
||||
let server = MockServer::start().await;
|
||||
let refresh_server = GatedRefreshServer::start().await?;
|
||||
let codex_home = TempDir::new()?;
|
||||
let server_url = format!("{}/mcp", server.uri());
|
||||
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()),
|
||||
"token_endpoint": refresh_server.token_endpoint.clone(),
|
||||
"scopes_supported": [""],
|
||||
})))
|
||||
.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_delay(Duration::from_millis(250))
|
||||
.set_body_json(json!({
|
||||
"access_token": REFRESHED_ACCESS_TOKEN,
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 7200,
|
||||
"refresh_token": ROTATED_REFRESH_TOKEN,
|
||||
})),
|
||||
)
|
||||
// This exact expectation proves that the two child processes issue one serialized provider
|
||||
// refresh rather than both replaying the same rotating refresh token.
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/mcp"))
|
||||
.and(header(
|
||||
@@ -199,8 +357,6 @@ async fn concurrent_file_mode_startup_refreshes_once() -> anyhow::Result<()> {
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let codex_home = TempDir::new()?;
|
||||
let server_url = format!("{}/mcp", server.uri());
|
||||
let seed_status = Command::new(std::env::current_exe()?)
|
||||
.args(["oauth_concurrency_seed_child", "--exact", "--ignored"])
|
||||
.env("CODEX_HOME", codex_home.path())
|
||||
@@ -212,34 +368,51 @@ async fn concurrent_file_mode_startup_refreshes_once() -> anyhow::Result<()> {
|
||||
"OAuth concurrency seed child failed: {seed_status}"
|
||||
);
|
||||
|
||||
let release_file = codex_home.path().join("oauth-clients.release");
|
||||
let first_ready_file = codex_home.path().join("oauth-client-first.ready");
|
||||
let first_release_file = codex_home.path().join("oauth-client-first.release");
|
||||
let second_ready_file = codex_home.path().join("oauth-client-second.ready");
|
||||
let mut first_child = Command::new(std::env::current_exe()?)
|
||||
.args(["oauth_concurrency_client_child", "--exact", "--ignored"])
|
||||
.env("CODEX_HOME", codex_home.path())
|
||||
.env(CHILD_SERVER_URL_ENV, &server_url)
|
||||
.env(CHILD_READY_FILE_ENV, &first_ready_file)
|
||||
.env(CHILD_RELEASE_FILE_ENV, &release_file)
|
||||
.kill_on_drop(true)
|
||||
.spawn()?;
|
||||
let mut second_child = Command::new(std::env::current_exe()?)
|
||||
.args(["oauth_concurrency_client_child", "--exact", "--ignored"])
|
||||
.env("CODEX_HOME", codex_home.path())
|
||||
.env(CHILD_SERVER_URL_ENV, &server_url)
|
||||
.env(CHILD_READY_FILE_ENV, &second_ready_file)
|
||||
.env(CHILD_RELEASE_FILE_ENV, &release_file)
|
||||
.kill_on_drop(true)
|
||||
.spawn()?;
|
||||
let second_release_file = codex_home.path().join("oauth-client-second.release");
|
||||
let contention_file = codex_home.path().join("oauth-client-second.contended");
|
||||
let mut first_child = oauth_concurrency_child_command(
|
||||
codex_home.path(),
|
||||
&server_url,
|
||||
&first_ready_file,
|
||||
&first_release_file,
|
||||
)?
|
||||
.spawn()?;
|
||||
let mut second_command = oauth_concurrency_child_command(
|
||||
codex_home.path(),
|
||||
&server_url,
|
||||
&second_ready_file,
|
||||
&second_release_file,
|
||||
)?;
|
||||
second_command.env(CHILD_CONTENTION_FILE_ENV, &contention_file);
|
||||
let mut second_child = second_command.spawn()?;
|
||||
|
||||
tokio::time::timeout(Duration::from_secs(10), async {
|
||||
while !first_ready_file.exists() || !second_ready_file.exists() {
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.map_err(|_| anyhow::anyhow!("OAuth concurrency children did not become ready"))?;
|
||||
fs::write(&release_file, b"release")?;
|
||||
wait_for_marker(
|
||||
&first_ready_file,
|
||||
"first OAuth concurrency child did not become ready",
|
||||
)
|
||||
.await?;
|
||||
wait_for_marker(
|
||||
&second_ready_file,
|
||||
"second OAuth concurrency child did not become ready",
|
||||
)
|
||||
.await?;
|
||||
|
||||
fs::write(&first_release_file, b"release")?;
|
||||
refresh_server.wait_until_request_started().await?;
|
||||
|
||||
// The first child is now inside the provider request while retaining the credential lock.
|
||||
// Releasing the second child must make its first try_lock call observe WouldBlock. Keep the
|
||||
// provider response gated until that exact branch emits the contention marker.
|
||||
fs::write(&second_release_file, b"release")?;
|
||||
let contention_result = wait_for_marker(
|
||||
&contention_file,
|
||||
"second OAuth concurrency child did not observe refresh-lock contention",
|
||||
)
|
||||
.await;
|
||||
refresh_server.release_responses();
|
||||
|
||||
let (first_status, second_status) = tokio::try_join!(first_child.wait(), second_child.wait())?;
|
||||
assert!(
|
||||
@@ -252,6 +425,8 @@ async fn concurrent_file_mode_startup_refreshes_once() -> anyhow::Result<()> {
|
||||
);
|
||||
|
||||
server.verify().await;
|
||||
contention_result?;
|
||||
assert_eq!(refresh_server.request_count(), 1);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -428,6 +603,12 @@ async fn oauth_concurrency_client_child() -> anyhow::Result<()> {
|
||||
let server_url = std::env::var(CHILD_SERVER_URL_ENV)?;
|
||||
let ready_file = PathBuf::from(std::env::var(CHILD_READY_FILE_ENV)?);
|
||||
let release_file = PathBuf::from(std::env::var(CHILD_RELEASE_FILE_ENV)?);
|
||||
if let Ok(marker_file) = std::env::var(CHILD_CONTENTION_FILE_ENV) {
|
||||
tracing::subscriber::set_global_default(LockContentionMarkerSubscriber {
|
||||
marker_file: PathBuf::from(marker_file),
|
||||
})
|
||||
.map_err(|error| anyhow::anyhow!("failed to install contention subscriber: {error}"))?;
|
||||
}
|
||||
let client = RmcpClient::new_streamable_http_client(
|
||||
SERVER_NAME,
|
||||
&server_url,
|
||||
@@ -441,11 +622,11 @@ async fn oauth_concurrency_client_child() -> anyhow::Result<()> {
|
||||
)
|
||||
.await?;
|
||||
|
||||
// Both processes must construct their OAuth client from the same expired snapshot before
|
||||
// either begins refresh. This removes process scheduling from the contention setup.
|
||||
// Both processes must construct their OAuth client from the same expired snapshot before the
|
||||
// parent releases either one. The parent then gates them separately to force lock contention.
|
||||
fs::write(ready_file, b"ready")?;
|
||||
while !release_file.exists() {
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
tokio::time::sleep(Duration::from_millis(/*millis*/ 10)).await;
|
||||
}
|
||||
initialize_client(&client).await?;
|
||||
Ok(())
|
||||
|
||||
Reference in New Issue
Block a user