openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-14 21:40:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

Instrument MCP OAuth credential store locks (#61430)

Browse Source 
## What changed

- Record acquisition counts, wait durations, and hold durations for file and secrets store locks, tagged by store, lock mode, operation, and originating client. Acquisition metrics also report outcome and contention.
- Preserve the resolved credential store's originating client across storage operations and lock release, and exclude nonblocking probes from lock metrics.
- Add `record_global_duration` for standalone durations with the existing buffering limits and metrics opt-out behavior.

## Testing

Extend tests to verify lock metrics for shared and exclusive access, contention, timeouts, probe exclusion, and originating-client attribution. Add coverage for standalone duration buffering, validation, and opt-out behavior.

GitOrigin-RevId: 96de8bd195477d61bc8cd966c5fbf72ae7156273
This commit is contained in:
Celia Chen
2026-11-14 01:21:54 +00:00
committed by copyberry
parent 6d919e65e2
commit de83e7674a
8 changed files with 471 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/otel/src/metrics/buffered.rs
+24 -10
View File
@@ -1,5 +1,5 @@
//! Buffers bounded operation counts and durations until process metrics are configured.
//! Each pair is admitted together; lifecycle transitions serialize with local SDK recording.
//! Standalone durations and count/duration pairs share admission limits and lifecycle handling.
//! Shutting down the active installation resumes buffering until its replacement is ready.
use super::MetricsClient;
@@ -17,7 +17,7 @@ const MAX_METADATA_BYTES: usize = 1024;
const MAX_TAGS: usize = 16;
struct PendingOperation {
count_name: Box<str>,
count_name: Option<Box<str>>,
duration_name: Box<str>,
duration: Duration,
tags: Vec<(Box<str>, Box<str>)>,
@@ -46,7 +46,17 @@ pub fn record_global_operation(
duration: Duration,
tags: &[(&str, &str)],
) -> Result<()> {
GLOBAL.record(count_name, duration_name, duration, tags)
GLOBAL.record(Some(count_name), duration_name, duration, tags)
}
/// Records a duration in milliseconds with the same startup limits and opt-out behavior as
/// `record_global_operation`, without an accompanying counter. Supply only bounded tags.
pub fn record_global_duration(
duration_name: &str,
duration: Duration,
tags: &[(&str, &str)],
) -> Result<()> {
GLOBAL.record(/*count_name*/ None, duration_name, duration, tags)
}
impl BufferedMetrics {
@@ -58,21 +68,23 @@ impl BufferedMetrics {
fn record(
&self,
count_name: &str,
count_name: Option<&str>,
duration_name: &str,
duration: Duration,
tags: &[(&str, &str)],
) -> Result<()> {
if tags.len() > MAX_TAGS
|| tags.iter().fold(
count_name.len().saturating_add(duration_name.len()),
count_name
.map_or(0, str::len)
.saturating_add(duration_name.len()),
|bytes, (key, value)| bytes.saturating_add(key.len()).saturating_add(value.len()),
) > MAX_METADATA_BYTES
{
return Err(MetricsError::OperationMetadataTooLarge);
}
// Validate both instruments before accepting either half of the pair.
for name in [count_name, duration_name] {
for name in count_name.into_iter().chain([duration_name]) {
validate_metric_name(name)?;
// The SDK otherwise silently builds a no-op instrument for these names.
if name.len() > 255 || !name.as_bytes()[0].is_ascii_alphabetic() {
@@ -93,7 +105,7 @@ impl BufferedMetrics {
State::Startup(pending) => {
if pending.len() < MAX_PENDING_OPERATIONS {
pending.push(PendingOperation {
count_name: count_name.into(),
count_name: count_name.map(Into::into),
duration_name: duration_name.into(),
duration,
tags: tags
@@ -129,7 +141,7 @@ impl BufferedMetrics {
.collect();
let _ = record(
&metrics,
&operation.count_name,
operation.count_name.as_deref(),
&operation.duration_name,
operation.duration,
&tags,
@@ -163,12 +175,14 @@ impl BufferedMetrics {
fn record(
metrics: &MetricsClient,
count_name: &str,
count_name: Option<&str>,
duration_name: &str,
duration: Duration,
tags: &[(&str, &str)],
) -> Result<()> {
metrics.counter(count_name, /*inc*/ 1, tags)?;
if let Some(count_name) = count_name {
metrics.counter(count_name, /*inc*/ 1, tags)?;
}
metrics.record_duration(duration_name, duration, tags)
}
codex-rs/otel/src/metrics/buffered_tests.rs
+69 -8
View File
@@ -16,7 +16,7 @@ fn client() -> MetricsClient {
fn observe(buffer: &BufferedMetrics) {
buffer
.record(
"test.operations",
Some("test.operations"),
"test.duration",
Duration::from_millis(/*millis*/ 7),
&[("component", "test")],
@@ -55,12 +55,17 @@ fn startup_limits_and_validation_preserve_whole_pairs() {
let buffer = BufferedMetrics::new();
let oversized = "x".repeat(MAX_METADATA_BYTES + 1);
assert!(matches!(
buffer.record("count", "duration", Duration::ZERO, &[("tag", &oversized)]),
buffer.record(
Some("count"),
"duration",
Duration::ZERO,
&[("tag", &oversized)]
),
Err(MetricsError::OperationMetadataTooLarge)
));
assert!(matches!(
buffer.record(
"count",
Some("count"),
"duration",
Duration::ZERO,
&[("tag", "value"); MAX_TAGS + 1]
@@ -68,17 +73,26 @@ fn startup_limits_and_validation_preserve_whole_pairs() {
Err(MetricsError::OperationMetadataTooLarge)
));
for name in ["invalid duration", "1duration", &"d".repeat(256)] {
assert!(buffer.record("count", name, Duration::ZERO, &[]).is_err());
assert!(
buffer
.record(name, "duration", Duration::ZERO, &[])
.record(/*count_name*/ None, name, Duration::ZERO, &[])
.is_err()
);
assert!(
buffer
.record(Some("count"), name, Duration::ZERO, &[])
.is_err()
);
assert!(
buffer
.record(Some(name), "duration", Duration::ZERO, &[])
.is_err()
);
}
assert!(
buffer
.record(
"count",
Some("count"),
"duration",
Duration::ZERO,
&[("tag", "invalid value")]
@@ -94,16 +108,63 @@ fn startup_limits_and_validation_preserve_whole_pairs() {
buffer.enable(&metrics);
assert_eq!(totals(&metrics), (0, 0, 0.0));
for name in ["invalid duration", "1duration", &"d".repeat(256)] {
assert!(buffer.record("count", name, Duration::ZERO, &[]).is_err());
assert!(
buffer
.record(name, "duration", Duration::ZERO, &[])
.record(/*count_name*/ None, name, Duration::ZERO, &[])
.is_err()
);
assert!(
buffer
.record(Some("count"), name, Duration::ZERO, &[])
.is_err()
);
assert!(
buffer
.record(Some(name), "duration", Duration::ZERO, &[])
.is_err()
);
}
assert_eq!(totals(&metrics), (0, 0, 0.0));
}
#[test]
fn standalone_durations_share_startup_limits_and_opt_out() {
let buffer = BufferedMetrics::new();
for _ in 0..MAX_PENDING_OPERATIONS + 1 {
buffer
.record(
/*count_name*/ None,
"test.duration",
Duration::from_millis(/*millis*/ 7),
&[],
)
.unwrap();
}
let metrics = client();
buffer.enable(&metrics);
assert_eq!(totals(&metrics), (0, 256, 1792.0));
buffer.disable();
buffer
.record(
/*count_name*/ None,
"test.duration",
Duration::from_millis(/*millis*/ 7),
&[],
)
.unwrap();
buffer.enable(&metrics);
assert_eq!(totals(&metrics), (0, 0, 0.0));
buffer
.record(
/*count_name*/ None,
"test.duration",
Duration::from_millis(/*millis*/ 7),
&[],
)
.unwrap();
assert_eq!(totals(&metrics), (0, 1, 7.0));
}
#[test]
fn opt_out_discards_pending_and_live_observations_until_reenabled() {
let buffer = BufferedMetrics::new();
codex-rs/otel/src/metrics/mod.rs
+1
View File
@@ -10,6 +10,7 @@ pub(crate) mod timer;
pub(crate) mod validation;
use crate::config::StatsigMetricsSettings;
pub use crate::metrics::buffered::record_global_duration;
pub use crate::metrics::buffered::record_global_operation;
pub use crate::metrics::client::MetricsClient;
pub use crate::metrics::config::MetricsConfig;
codex-rs/rmcp-client/src/oauth.rs
+9 -4
View File
@@ -72,6 +72,7 @@ use tracing::warn;
use self::store_lock::OAuthStore;
use self::store_lock::OAuthStoreLock;
use self::store_lock::OAuthStoreLockFailure;
use self::store_lock::OAuthStoreWriteOperation;
use codex_config::McpServerOAuthConfig;
use codex_keyring_store::DefaultKeyringStore;
@@ -568,7 +569,8 @@ fn save_oauth_tokens_to_secrets_keyring<K: KeyringStore + Clone + 'static>(
tokens: &StoredOAuthTokens,
) -> Result<()> {
let serialized = serde_json::to_string(tokens).context("failed to serialize OAuth tokens")?;
let _store_lock = OAuthStoreLock::acquire_for_write(OAuthStore::Secrets)?;
let _store_lock =
OAuthStoreLock::acquire_for_write(OAuthStore::Secrets, OAuthStoreWriteOperation::Save)?;
save_oauth_tokens_to_secrets_keyring_with_lock_held(
keyring_store,
server_name,
@@ -751,7 +753,8 @@ fn delete_oauth_tokens_from_secrets_keyring<K: KeyringStore + Clone + 'static>(
server_name: &str,
url: &str,
) -> Result<bool> {
let _store_lock = OAuthStoreLock::acquire_for_write(OAuthStore::Secrets)?;
let _store_lock =
OAuthStoreLock::acquire_for_write(OAuthStore::Secrets, OAuthStoreWriteOperation::Delete)?;
let codex_home = find_codex_home()?;
let manager = SecretsManager::new_with_keyring_store_and_namespace(
codex_home.to_path_buf(),
@@ -1073,7 +1076,8 @@ fn load_oauth_tokens_from_file_with_lock_held(
/// Saves one credential while holding the File aggregate-store lock across the full
/// read-modify-write operation.
fn save_oauth_tokens_to_file(tokens: &StoredOAuthTokens) -> Result<()> {
let _store_lock = OAuthStoreLock::acquire_for_write(OAuthStore::File)?;
let _store_lock =
OAuthStoreLock::acquire_for_write(OAuthStore::File, OAuthStoreWriteOperation::Save)?;
save_oauth_tokens_to_file_with_lock_held(tokens)
}
@@ -1116,7 +1120,8 @@ impl From<&StoredOAuthTokens> for FallbackTokenEntry {
}
fn delete_oauth_tokens_from_file(server_name: &str, url: &str, key: &str) -> Result<bool> {
let _store_lock = OAuthStoreLock::acquire_for_write(OAuthStore::File)?;
let _store_lock =
OAuthStoreLock::acquire_for_write(OAuthStore::File, OAuthStoreWriteOperation::Delete)?;
let mut store = match read_fallback_file_unlocked()? {
Some(store) => store,
None => return Ok(false),
codex-rs/rmcp-client/src/oauth/resolved_store.rs
+23 -14
View File
@@ -17,6 +17,7 @@ use super::OAuthKeyringLoadError;
use super::OAuthStore;
use super::OAuthStoreLock;
use super::OAuthStoreLockFailure;
use super::OAuthStoreWriteOperation;
use super::StoredOAuthTokens;
use super::compute_store_key;
use super::delete_oauth_tokens_from_direct_keyring;
@@ -103,7 +104,7 @@ impl ResolvedOAuthCredentialStore {
url: &str,
) -> Result<Option<StoredOAuthTokens>> {
let mut observation = telemetry::resolved(self, Operation::Load);
let result = match self.backend {
let result = self.originator.sync_scope(|| match self.backend {
Backend::File => load_oauth_tokens_from_file(server_name, url)
.context("failed to reread OAuth tokens from resolved file storage"),
Backend::Keyring(keyring_backend_kind) => load_oauth_tokens_from_keyring(
@@ -116,7 +117,7 @@ impl ResolvedOAuthCredentialStore {
.context(
"failed to reread OAuth tokens from resolved keyring storage; refusing file fallback",
),
};
});
observation.record_load_attempt(telemetry::store(self), &result);
if matches!(self.backend, Backend::Keyring(_))
&& let Err(error) = &result
@@ -166,7 +167,7 @@ impl ResolvedOAuthCredentialStore {
tokens: &StoredOAuthTokens,
) -> Result<()> {
let mut observation = telemetry::resolved(self, Operation::Save);
let result = match self.backend {
let result = self.originator.sync_scope(|| match self.backend {
Backend::File => save_oauth_tokens_to_file(tokens),
Backend::Keyring(keyring_backend_kind) => save_oauth_tokens_with_keyring(
keyring_store,
@@ -174,7 +175,7 @@ impl ResolvedOAuthCredentialStore {
server_name,
tokens,
),
};
});
observation.record_save_attempt(telemetry::store(self), &result);
if matches!(self.backend, Backend::Keyring(_))
&& let Err(error) = &result
@@ -199,8 +200,11 @@ impl ResolvedOAuthCredentialStore {
let result = match self.backend {
Backend::File => {
let mut save = telemetry::resolved(self, Operation::Save);
let result = (|| {
let _store_lock = OAuthStoreLock::acquire_for_write(OAuthStore::File)?;
let result = self.originator.sync_scope(|| {
let _store_lock = OAuthStoreLock::acquire_for_write(
OAuthStore::File,
OAuthStoreWriteOperation::Save,
)?;
let (key, current) =
load_oauth_tokens_from_file_with_lock_held(server_name, &tokens.url)?
.ok_or(AuthError::AuthorizationRequired)?;
@@ -223,7 +227,7 @@ impl ResolvedOAuthCredentialStore {
..FallbackTokenEntry::from(tokens)
};
write_fallback_file(&store)
})();
});
save.record_save_attempt(Store::File, &result);
result
}
@@ -253,10 +257,13 @@ impl ResolvedOAuthCredentialStore {
let mut observation = telemetry::resolved(self, Operation::Load);
let mut store_lock = None;
let mut file_key = None;
let current = (|| -> Result<_> {
let current = self.originator.sync_scope(|| -> Result<_> {
match self.backend {
Backend::File => {
store_lock = Some(OAuthStoreLock::acquire_for_write(OAuthStore::File)?);
store_lock = Some(OAuthStoreLock::acquire_for_write(
OAuthStore::File,
OAuthStoreWriteOperation::RetireRefreshToken,
)?);
Ok(
load_oauth_tokens_from_file_with_lock_held(server_name, url)?.map(
|(key, tokens)| {
@@ -267,8 +274,10 @@ impl ResolvedOAuthCredentialStore {
)
}
Backend::Keyring(AuthKeyringBackendKind::Secrets) => {
let guard =
store_lock.insert(OAuthStoreLock::acquire_for_write(OAuthStore::Secrets)?);
let guard = store_lock.insert(OAuthStoreLock::acquire_for_write(
OAuthStore::Secrets,
OAuthStoreWriteOperation::RetireRefreshToken,
)?);
load_oauth_tokens_from_secrets_keyring(keyring_store, server_name, url, guard)
.map_err(anyhow::Error::from)
}
@@ -280,7 +289,7 @@ impl ResolvedOAuthCredentialStore {
)
.map_err(anyhow::Error::from),
}
})();
});
observation.record_load_attempt(telemetry::store(self), &current);
if matches!(self.backend, Backend::Keyring(_))
&& let Err(error) = &current
@@ -339,7 +348,7 @@ impl ResolvedOAuthCredentialStore {
url: &str,
) -> Result<bool> {
let mut observation = telemetry::resolved(self, Operation::Delete);
let result = match self.backend {
let result = self.originator.sync_scope(|| match self.backend {
Backend::File => compute_store_key(server_name, url)
.and_then(|key| delete_oauth_tokens_from_file(server_name, url, &key)),
Backend::Keyring(AuthKeyringBackendKind::Direct) => {
@@ -348,7 +357,7 @@ impl ResolvedOAuthCredentialStore {
Backend::Keyring(AuthKeyringBackendKind::Secrets) => {
delete_oauth_tokens_from_secrets_keyring(keyring_store, server_name, url)
}
};
});
observation.record_delete_attempt(telemetry::store(self), &result);
if matches!(self.backend, Backend::Keyring(_))
&& let Err(error) = &result
codex-rs/rmcp-client/src/oauth/store_lock.rs
+109 -8
View File
@@ -4,6 +4,9 @@
//! Writers hold an exclusive lock across the complete read-modify-write operation. Readers share
//! the same lock so concurrent MCP startup and status checks do not serialize behind one another.
//! Direct keyring entries are already stored independently per credential and do not use this lock.
//! Blocking acquisitions report contention and wait time; released guards report hold time.
//! Guards retain the operation and originating client captured at acquisition.
//! Nonblocking probes are excluded, and telemetry never changes the lock policy or deadline.
use std::fs;
use std::fs::File;
@@ -14,6 +17,7 @@ use std::path::PathBuf;
use std::time::Duration;
use std::time::Instant;
use codex_otel::auth_storage::AuthStorageOriginator;
use codex_utils_home_dir::find_codex_home;
const OAUTH_LOCK_DIR: &str = "mcp-oauth-locks";
@@ -48,21 +52,49 @@ impl std::fmt::Display for OAuthStore {
/// Serializes one complete operation on an aggregate OAuth credential store.
pub(super) struct OAuthStoreLock {
// Fields drop in declaration order: close the file before recording hold telemetry.
_file: File,
_hold: Option<LockHoldTelemetry>,
}
struct LockHoldTelemetry {
acquired_at: Instant,
tags: [(&'static str, &'static str); 4],
}
impl Drop for LockHoldTelemetry {
fn drop(&mut self) {
let _ = codex_otel::record_global_duration(
"codex.auth_storage.lock_hold.duration",
self.acquired_at.elapsed(),
&self.tags,
);
}
}
#[derive(Clone, Copy, Debug)]
enum OAuthStoreLockMode {
Shared,
Exclusive,
Exclusive(OAuthStoreWriteOperation),
}
/// Identifies the mutation owning an exclusive aggregate-store lock.
#[derive(Clone, Copy, Debug)]
pub(super) enum OAuthStoreWriteOperation {
Save,
Delete,
RetireRefreshToken,
}
impl OAuthStoreLock {
pub(super) fn acquire_for_write(store: OAuthStore) -> Result<Self, OAuthStoreLockFailure> {
pub(super) fn acquire_for_write(
store: OAuthStore,
operation: OAuthStoreWriteOperation,
) -> Result<Self, OAuthStoreLockFailure> {
Self::acquire_with_timeout(
store,
STORE_LOCK_ACQUIRE_TIMEOUT,
OAuthStoreLockMode::Exclusive,
OAuthStoreLockMode::Exclusive(operation),
)
}
@@ -119,18 +151,58 @@ impl OAuthStoreLock {
path: path.clone(),
source,
})?;
let tags = [
(
"store",
match store {
OAuthStore::File => "file",
OAuthStore::Secrets => "secrets",
},
),
(
"lock_mode",
match mode {
OAuthStoreLockMode::Shared => "shared",
OAuthStoreLockMode::Exclusive(_) => "exclusive",
},
),
(
"operation",
match mode {
OAuthStoreLockMode::Shared => "load",
OAuthStoreLockMode::Exclusive(OAuthStoreWriteOperation::Save) => "save",
OAuthStoreLockMode::Exclusive(OAuthStoreWriteOperation::Delete) => "delete",
OAuthStoreLockMode::Exclusive(OAuthStoreWriteOperation::RetireRefreshToken) => {
"retire_refresh_token"
}
},
),
(
codex_otel::ORIGINATOR_TAG,
AuthStorageOriginator::current().as_str(),
),
];
let started = Instant::now();
let mut reported_contention = false;
loop {
let result = loop {
let result = match mode {
OAuthStoreLockMode::Shared => file.try_lock_shared(),
OAuthStoreLockMode::Exclusive => file.try_lock(),
OAuthStoreLockMode::Exclusive(_) => file.try_lock(),
};
match result {
Ok(()) => return Ok(Self { _file: file }),
Ok(()) => {
break Ok(Self {
_file: file,
_hold: (!acquire_timeout.is_zero()).then(|| LockHoldTelemetry {
acquired_at: Instant::now(),
tags,
}),
});
}
Err(std::fs::TryLockError::WouldBlock) if started.elapsed() >= acquire_timeout => {
return Err(OAuthStoreLockFailure::Timeout {
reported_contention = true;
break Err(OAuthStoreLockFailure::Timeout {
store,
path,
acquire_timeout,
@@ -149,14 +221,43 @@ impl OAuthStoreLock {
std::thread::sleep(STORE_LOCK_RETRY_SLEEP.min(acquire_timeout));
}
Err(error) => {
return Err(OAuthStoreLockFailure::Lock {
break Err(OAuthStoreLockFailure::Lock {
store,
path,
source: io::Error::from(error),
});
}
}
};
if !acquire_timeout.is_zero() {
let outcome = match &result {
Ok(_) => "acquired",
Err(OAuthStoreLockFailure::Timeout { .. }) => "timeout",
Err(
OAuthStoreLockFailure::CodexHome { .. }
| OAuthStoreLockFailure::CreateDir { .. }
| OAuthStoreLockFailure::Open { .. }
| OAuthStoreLockFailure::Lock { .. },
) => "error",
};
let _ = codex_otel::record_global_operation(
"codex.auth_storage.lock_acquire",
"codex.auth_storage.lock_wait.duration",
started.elapsed(),
&[
tags[0],
tags[1],
tags[2],
tags[3],
("outcome", outcome),
(
"contended",
if reported_contention { "true" } else { "false" },
),
],
);
}
result
}
}
codex-rs/rmcp-client/src/oauth/tests/credential_store_tests.rs
+7 -2
View File
@@ -1,3 +1,5 @@
#![cfg(test)]
//! Verifies guarded mutations, pinned storage, token mapping, and saved runtime snapshots.
use std::time::Duration;
@@ -28,6 +30,7 @@ use crate::oauth::save_oauth_tokens_to_file;
use crate::oauth::save_oauth_tokens_to_file_with_lock_held;
use crate::oauth::store_lock::OAuthStore;
use crate::oauth::store_lock::OAuthStoreLock;
use crate::oauth::store_lock::OAuthStoreWriteOperation;
#[tokio::test(flavor = "current_thread")]
async fn mutations_require_and_retain_the_transaction_guard() -> Result<()> {
@@ -68,7 +71,8 @@ async fn mutations_require_and_retain_the_transaction_guard() -> Result<()> {
if !clear {
store.clear().await?;
}
let aggregate_lock = OAuthStoreLock::acquire_for_write(OAuthStore::File)?;
let aggregate_lock =
OAuthStoreLock::acquire_for_write(OAuthStore::File, OAuthStoreWriteOperation::Save)?;
let mut mutation = if clear {
store.clear()
} else {
@@ -308,7 +312,8 @@ async fn cancelled_rejection_retains_guard_and_preserves_replacement() -> Result
Err(AuthError::CredentialStoreError(error)) if error.contains("active refresh guard")
));
let guard = store.acquire_transaction_guard().await?;
let aggregate_lock = OAuthStoreLock::acquire_for_write(OAuthStore::File)?;
let aggregate_lock =
OAuthStoreLock::acquire_for_write(OAuthStore::File, OAuthStoreWriteOperation::Save)?;
let mut mutation = store.on_refresh_token_rejected(&credentials);
assert!(futures::poll!(&mut mutation).is_pending());
drop(mutation);
codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs
+229 -10
View File
@@ -1,3 +1,6 @@
#![cfg(test)]
use std::collections::BTreeMap;
use std::process::Command;
use std::sync::mpsc;
use std::time::Duration;
@@ -13,6 +16,8 @@ use oauth2::RefreshToken;
use oauth2::Scope;
use oauth2::TokenResponse;
use oauth2::basic::BasicTokenType;
use opentelemetry_sdk::metrics::data::AggregatedMetrics;
use opentelemetry_sdk::metrics::data::MetricData;
use pretty_assertions::assert_eq;
use rmcp::transport::auth::OAuthTokenResponse;
use rmcp::transport::auth::VendorExtraTokenFields;
@@ -28,6 +33,7 @@ use super::OAuthStore;
use super::OAuthStoreLock;
use super::OAuthStoreLockFailure;
use super::OAuthStoreLockMode;
use super::OAuthStoreWriteOperation;
use crate::oauth::ResolvedOAuthCredentialStore;
use crate::oauth::StoredOAuthCredentialSnapshot;
use crate::oauth::StoredOAuthTokens;
@@ -46,6 +52,7 @@ use crate::oauth::test_support::GatedKeyringStore;
use crate::oauth::test_support::ReadOutcome;
use crate::oauth::test_support::TempCodexHome;
use codex_config::types::OAuthCredentialsStoreMode;
use codex_otel::auth_storage::AuthStorageOriginator;
const STORE_LOCK_CONTENTION_EVENT_TARGET: &str = "codex_rmcp_client::oauth::store_lock::contention";
// Contention is proven by the tracing event emitted after a real WouldBlock. Keep the timeout
@@ -216,7 +223,7 @@ fn store_lock_is_released_when_holder_process_exits() -> Result<()> {
env.path(),
OAuthStore::File,
Duration::from_millis(/*millis*/ 100),
OAuthStoreLockMode::Exclusive,
OAuthStoreLockMode::Exclusive(OAuthStoreWriteOperation::Save),
) {
Ok(_) => {
anyhow::bail!("live holder process should keep the OAuth store lock unavailable")
@@ -236,7 +243,7 @@ fn store_lock_is_released_when_holder_process_exits() -> Result<()> {
env.path(),
OAuthStore::File,
Duration::from_secs(/*secs*/ 1),
OAuthStoreLockMode::Exclusive,
OAuthStoreLockMode::Exclusive(OAuthStoreWriteOperation::Save),
)?;
Ok(())
})();
@@ -256,7 +263,8 @@ fn store_lock_is_released_when_holder_process_exits_child() -> Result<()> {
Some(path) => std::path::PathBuf::from(path),
None => return Ok(()),
};
let _lock = OAuthStoreLock::acquire_for_write(OAuthStore::File)?;
let _lock =
OAuthStoreLock::acquire_for_write(OAuthStore::File, OAuthStoreWriteOperation::Save)?;
std::fs::write(ready_file, b"ready")?;
loop {
std::thread::sleep(Duration::from_secs(/*secs*/ 60));
@@ -357,7 +365,8 @@ fn oauth_credential_probes_skip_contended_file_and_secrets_stores() -> Result<()
};
{
let _lock = OAuthStoreLock::acquire_for_write(OAuthStore::File)?;
let _lock =
OAuthStoreLock::acquire_for_write(OAuthStore::File, OAuthStoreWriteOperation::Save)?;
assert_eq!(reload(file, direct)?, None);
assert_eq!(reload(auto, secrets)?, None);
assert_eq!(probe(file, direct)?, None);
@@ -375,7 +384,8 @@ fn oauth_credential_probes_skip_contended_file_and_secrets_stores() -> Result<()
let refreshed = refresh(&tokens.url)?.expect("the unlocked store should still contain tokens");
assert!(!refreshed.store_was_contended());
let _lock = OAuthStoreLock::acquire_for_write(OAuthStore::Secrets)?;
let _lock =
OAuthStoreLock::acquire_for_write(OAuthStore::Secrets, OAuthStoreWriteOperation::Save)?;
assert_eq!(
probe(auto, secrets)?,
None,
@@ -440,7 +450,7 @@ where
codex_home,
store,
Duration::from_millis(/*millis*/ 100),
OAuthStoreLockMode::Exclusive,
OAuthStoreLockMode::Exclusive(OAuthStoreWriteOperation::Save),
)?;
let (contended_tx, contended_rx) = mpsc::channel();
let worker = scope.spawn(move || {
@@ -462,7 +472,29 @@ where
#[test]
fn aggregate_store_readers_share_access_while_writers_remain_exclusive() -> Result<()> {
// Isolate process-global metrics from other libtest cases, including under Bazel.
if std::env::var_os("CODEX_AUTH_STORAGE_TEST_CHILD").is_none() {
let output = Command::new(std::env::current_exe()?)
.args([
"--exact",
std::thread::current().name().unwrap(),
"--nocapture",
])
.env("CODEX_AUTH_STORAGE_TEST_CHILD", "1")
.output()?;
assert!(
output.status.success(),
"{}{}",
String::from_utf8_lossy(&output.stdout),
String::from_utf8_lossy(&output.stderr)
);
return Ok(());
}
let env = TempCodexHome::new();
let metrics = codex_otel::install_global_metrics(codex_otel::MetricsClient::new(
codex_otel::MetricsConfig::in_memory("test", "test", "1", Default::default())
.with_runtime_reader(),
)?);
for store in [OAuthStore::File, OAuthStore::Secrets] {
let readers = (0..2)
@@ -480,7 +512,7 @@ fn aggregate_store_readers_share_access_while_writers_remain_exclusive() -> Resu
env.path(),
store,
Duration::from_millis(/*millis*/ 100),
OAuthStoreLockMode::Exclusive,
OAuthStoreLockMode::Exclusive(OAuthStoreWriteOperation::Save),
) {
Ok(_) => anyhow::bail!("an active OAuth store reader must exclude writers"),
Err(error) => error,
@@ -490,14 +522,201 @@ fn aggregate_store_readers_share_access_while_writers_remain_exclusive() -> Resu
OAuthStoreLockFailure::Timeout { .. }
));
drop(readers);
let _writer = OAuthStoreLock::acquire_in_with_mode(
let writer = OAuthStoreLock::acquire_in_with_mode(
env.path(),
store,
Duration::from_millis(/*millis*/ 100),
OAuthStoreLockMode::Exclusive,
OAuthStoreLockMode::Exclusive(OAuthStoreWriteOperation::Save),
)?;
// Busy and successful zero-deadline probes must not enter the waiting population.
assert!(matches!(
OAuthStoreLock::try_acquire_for_read(store),
Err(OAuthStoreLockFailure::Timeout { .. })
));
drop(writer);
drop(OAuthStoreLock::try_acquire_for_read(store)?);
for mode in [
OAuthStoreLockMode::Shared,
OAuthStoreLockMode::Exclusive(OAuthStoreWriteOperation::Save),
] {
let path = env.path().to_path_buf();
drop(complete_after_store_lock_contention(
env.path(),
store,
|| Ok(()),
move || {
Ok(OAuthStoreLock::acquire_in_with_mode(
&path,
store,
Duration::from_secs(/*secs*/ 1),
mode,
)?)
},
)?);
}
}
// Pinned stores must retain their client even on a worker with another ambient originator.
let originator = AuthStorageOriginator::from_client_name("codex_vscode");
for source in [
originator.sync_scope(ResolvedOAuthCredentialStore::file),
originator
.sync_scope(|| ResolvedOAuthCredentialStore::keyring(AuthKeyringBackendKind::Secrets)),
] {
std::thread::spawn(move || {
AuthStorageOriginator::from_client_name("codex_cli").sync_scope(|| -> Result<()> {
let keyring = MockKeyringStore::default();
let tokens = sample_tokens();
source.save(&keyring, &tokens.server_name, &tokens)?;
source
.load(&keyring, &tokens.server_name, &tokens.url)?
.context("saved credential")?;
source.save_with_refresh_telemetry(
&keyring,
&tokens.server_name,
Some(&tokens),
&tokens,
)?;
source.retire_refresh_token_if_matches(
&keyring,
&tokens.server_name,
&tokens.url,
&tokens,
)?;
assert!(source.delete(&keyring, &tokens.server_name, &tokens.url)?);
Ok(())
})
})
.join()
.expect("storage worker should finish")?;
}
let snapshot = metrics.snapshot()?;
let mut actual = Vec::new();
for metric in snapshot
.scope_metrics()
.flat_map(opentelemetry_sdk::metrics::data::ScopeMetrics::metrics)
{
if !metric.name().starts_with("codex.auth_storage.lock_") {
continue;
}
match metric.data() {
AggregatedMetrics::U64(MetricData::Sum(sum)) => {
for point in sum.data_points() {
actual.push((
metric.name().to_owned(),
point.value(),
point
.attributes()
.map(|a| (a.key.to_string(), a.value.to_string()))
.collect::<BTreeMap<_, _>>(),
));
}
}
AggregatedMetrics::F64(MetricData::Histogram(histogram)) => {
for point in histogram.data_points() {
let tags = point
.attributes()
.map(|a| (a.key.to_string(), a.value.to_string()))
.collect::<BTreeMap<_, _>>();
if tags
.get("outcome")
.is_some_and(|outcome| outcome == "timeout")
{
assert!(
point.sum() >= 100.0,
"timeout wait must include the acquisition deadline"
);
}
actual.push((metric.name().to_owned(), point.count(), tags));
}
}
other => panic!("unexpected metric: {other:?}"),
}
}
let mut expected = Vec::new();
for store in ["file", "secrets"] {
for name in [
"codex.auth_storage.lock_acquire",
"codex.auth_storage.lock_wait.duration",
] {
for (mode, outcome, contended, count) in [
("shared", "acquired", "false", 2),
("exclusive", "timeout", "true", 1),
("exclusive", "acquired", "false", 3),
("shared", "acquired", "true", 1),
("exclusive", "acquired", "true", 1),
] {
expected.push((
name.to_owned(),
count,
BTreeMap::from([
("store".into(), store.into()),
("lock_mode".into(), mode.into()),
(
"operation".into(),
if mode == "shared" { "load" } else { "save" }.into(),
),
("originator".into(), "none".into()),
("outcome".into(), outcome.into()),
("contended".into(), contended.into()),
]),
));
}
}
for (mode, count) in [("shared", 3), ("exclusive", 4)] {
expected.push((
"codex.auth_storage.lock_hold.duration".to_owned(),
count,
BTreeMap::from([
("store".into(), store.into()),
("lock_mode".into(), mode.into()),
(
"operation".into(),
if mode == "shared" { "load" } else { "save" }.into(),
),
("originator".into(), "none".into()),
]),
));
}
for operation in ["load", "save", "delete", "retire_refresh_token"] {
let count = if operation == "save" { 2 } else { 1 };
let mut tags = BTreeMap::from([
("store".into(), store.into()),
(
"lock_mode".into(),
if operation == "load" {
"shared"
} else {
"exclusive"
}
.into(),
),
("operation".into(), operation.into()),
("originator".into(), "codex_vscode".into()),
]);
expected.push((
"codex.auth_storage.lock_hold.duration".to_owned(),
count,
tags.clone(),
));
tags.extend([
("outcome".into(), "acquired".into()),
("contended".into(), "false".into()),
]);
for name in [
"codex.auth_storage.lock_acquire",
"codex.auth_storage.lock_wait.duration",
] {
expected.push((name.to_owned(), count, tags.clone()));
}
}
}
actual.sort();
expected.sort();
assert_eq!(actual, expected);
metrics.shutdown()?;
Ok(())
}
@@ -622,7 +841,7 @@ async fn secrets_reads_release_the_file_lock_before_keyring_access() -> Result<(
home.path(),
OAuthStore::Secrets,
Duration::ZERO,
OAuthStoreLockMode::Exclusive,
OAuthStoreLockMode::Exclusive(OAuthStoreWriteOperation::Save),
)?;
release.send(ReadOutcome::Stored)?;
let result = reader.await??.expect("read captured the stored credential");