diff --git a/codex-rs/otel/src/metrics/buffered.rs b/codex-rs/otel/src/metrics/buffered.rs index 96d6e31bc6..4c0136fc04 100644 --- a/codex-rs/otel/src/metrics/buffered.rs +++ b/codex-rs/otel/src/metrics/buffered.rs @@ -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, + count_name: Option>, duration_name: Box, duration: Duration, tags: Vec<(Box, Box)>, @@ -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) } diff --git a/codex-rs/otel/src/metrics/buffered_tests.rs b/codex-rs/otel/src/metrics/buffered_tests.rs index 3fc1c2c630..fad5295b11 100644 --- a/codex-rs/otel/src/metrics/buffered_tests.rs +++ b/codex-rs/otel/src/metrics/buffered_tests.rs @@ -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(); diff --git a/codex-rs/otel/src/metrics/mod.rs b/codex-rs/otel/src/metrics/mod.rs index c4f9d3155d..314540314a 100644 --- a/codex-rs/otel/src/metrics/mod.rs +++ b/codex-rs/otel/src/metrics/mod.rs @@ -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; diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index 910811b912..defbfb00f3 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -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( 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( server_name: &str, url: &str, ) -> Result { - 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 { - 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), diff --git a/codex-rs/rmcp-client/src/oauth/resolved_store.rs b/codex-rs/rmcp-client/src/oauth/resolved_store.rs index 95c473b588..1c4ee65913 100644 --- a/codex-rs/rmcp-client/src/oauth/resolved_store.rs +++ b/codex-rs/rmcp-client/src/oauth/resolved_store.rs @@ -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> { 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), ¤t); if matches!(self.backend, Backend::Keyring(_)) && let Err(error) = ¤t @@ -339,7 +348,7 @@ impl ResolvedOAuthCredentialStore { url: &str, ) -> Result { 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 diff --git a/codex-rs/rmcp-client/src/oauth/store_lock.rs b/codex-rs/rmcp-client/src/oauth/store_lock.rs index f06b9b9adc..39d6d30f73 100644 --- a/codex-rs/rmcp-client/src/oauth/store_lock.rs +++ b/codex-rs/rmcp-client/src/oauth/store_lock.rs @@ -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, +} + +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 { + pub(super) fn acquire_for_write( + store: OAuthStore, + operation: OAuthStoreWriteOperation, + ) -> Result { 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 } } diff --git a/codex-rs/rmcp-client/src/oauth/tests/credential_store_tests.rs b/codex-rs/rmcp-client/src/oauth/tests/credential_store_tests.rs index f9e2ba2907..3e439b9ccf 100644 --- a/codex-rs/rmcp-client/src/oauth/tests/credential_store_tests.rs +++ b/codex-rs/rmcp-client/src/oauth/tests/credential_store_tests.rs @@ -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); diff --git a/codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs b/codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs index f9f6b0f185..fb79b22e87 100644 --- a/codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs +++ b/codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs @@ -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::>(), + )); + } + } + 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::>(); + 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");