From 8b1b065719103f5c5a1bdd8b6a95354a72b06d97 Mon Sep 17 00:00:00 2001 From: iceweasel-oai Date: Mon, 10 Aug 2026 15:34:35 +0000 Subject: [PATCH] Speed up MCP OAuth credential reads (#37842) ## Why Concurrent MCP startup and status checks should not serialize when they only read the shared credential store. Repeated reads of the encrypted MCP OAuth store also needlessly decrypted unchanged contents. ## What changed - Use shared locks for `File` and `Secrets` credential reads while keeping saves and deletes exclusive. - Cache decrypted MCP OAuth secrets by store path, ciphertext, and passphrase, and invalidate the cache after writes. ## Testing - Cover concurrent readers, reader/writer exclusion, shared credential loads, and cache invalidation after updates and deletes. GitOrigin-RevId: f13512e6404d4919879ba5ba77a3e34e52b35640 --- codex-rs/rmcp-client/src/oauth.rs | 12 +- codex-rs/rmcp-client/src/oauth/store_lock.rs | 35 ++++- .../src/oauth/tests/store_lock_tests.rs | 101 ++++++++++++- codex-rs/secrets/src/local.rs | 138 ++++++++++++++++++ 4 files changed, 269 insertions(+), 17 deletions(-) diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index 913108b8fe..6af8d95f2f 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -282,7 +282,7 @@ fn load_oauth_tokens_from_secrets_keyring( server_name: &str, url: &str, ) -> std::result::Result, OAuthKeyringLoadError> { - let _store_lock = OAuthStoreLock::acquire(OAuthStore::Secrets)?; + let _store_lock = OAuthStoreLock::acquire_for_read(OAuthStore::Secrets)?; let codex_home = find_codex_home().map_err(anyhow::Error::from)?; let manager = SecretsManager::new_with_keyring_store_and_namespace( codex_home.to_path_buf(), @@ -386,7 +386,7 @@ 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(OAuthStore::Secrets)?; + let _store_lock = OAuthStoreLock::acquire_for_write(OAuthStore::Secrets)?; save_oauth_tokens_to_secrets_keyring_with_lock_held( keyring_store, server_name, @@ -541,7 +541,7 @@ fn delete_oauth_tokens_from_secrets_keyring( server_name: &str, url: &str, ) -> Result { - let _store_lock = OAuthStoreLock::acquire(OAuthStore::Secrets)?; + let _store_lock = OAuthStoreLock::acquire_for_write(OAuthStore::Secrets)?; let codex_home = find_codex_home()?; let manager = SecretsManager::new_with_keyring_store_and_namespace( codex_home.to_path_buf(), @@ -681,7 +681,7 @@ struct FallbackTokenEntry { } fn load_oauth_tokens_from_file(server_name: &str, url: &str) -> Result> { - let _store_lock = OAuthStoreLock::acquire(OAuthStore::File)?; + let _store_lock = OAuthStoreLock::acquire_for_read(OAuthStore::File)?; let Some(store) = read_fallback_file_unlocked()? else { return Ok(None); }; @@ -739,7 +739,7 @@ fn load_oauth_tokens_from_file(server_name: &str, url: &str) -> Result Result<()> { - let _store_lock = OAuthStoreLock::acquire(OAuthStore::File)?; + let _store_lock = OAuthStoreLock::acquire_for_write(OAuthStore::File)?; save_oauth_tokens_to_file_with_lock_held(tokens) } @@ -779,7 +779,7 @@ fn save_oauth_tokens_to_file_with_lock_held(tokens: &StoredOAuthTokens) -> Resul } fn delete_oauth_tokens_from_file(key: &str) -> Result { - let _store_lock = OAuthStoreLock::acquire(OAuthStore::File)?; + let _store_lock = OAuthStoreLock::acquire_for_write(OAuthStore::File)?; let mut store = match read_fallback_file_unlocked()? { Some(store) => store, None => return Ok(false), diff --git a/codex-rs/rmcp-client/src/oauth/store_lock.rs b/codex-rs/rmcp-client/src/oauth/store_lock.rs index b053f4e5ae..0a6f3c6494 100644 --- a/codex-rs/rmcp-client/src/oauth/store_lock.rs +++ b/codex-rs/rmcp-client/src/oauth/store_lock.rs @@ -1,8 +1,9 @@ //! Cross-process serialization for MCP OAuth stores shared by multiple credentials. //! //! File and Secrets each keep credentials for multiple MCP servers in one aggregate document. -//! Their lock therefore protects the complete read-modify-write operation. Direct keyring entries -//! are already stored independently per credential and do not use this lock. +//! 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. use std::fs; use std::fs::File; @@ -50,8 +51,25 @@ pub(super) struct OAuthStoreLock { _file: File, } +#[derive(Clone, Copy, Debug)] +enum OAuthStoreLockMode { + Shared, + Exclusive, +} + impl OAuthStoreLock { - pub(super) fn acquire(store: OAuthStore) -> Result { + pub(super) fn acquire_for_write(store: OAuthStore) -> Result { + Self::acquire_with_mode(store, OAuthStoreLockMode::Exclusive) + } + + pub(super) fn acquire_for_read(store: OAuthStore) -> Result { + Self::acquire_with_mode(store, OAuthStoreLockMode::Shared) + } + + fn acquire_with_mode( + store: OAuthStore, + mode: OAuthStoreLockMode, + ) -> Result { // This lock intentionally follows the existing local File/Secrets credential-store // authority. Those stores are CODEX_HOME-backed today: if CODEX_HOME is unset they use // the default home (`~/.codex`), and if an embedder has no local home/filesystem authority @@ -59,13 +77,14 @@ impl OAuthStoreLock { // provide its own matching lock authority instead of using this local path. let codex_home = find_codex_home() .map_err(|source| OAuthStoreLockFailure::CodexHome { store, source })?; - Self::acquire_in(&codex_home, store, STORE_LOCK_ACQUIRE_TIMEOUT) + Self::acquire_in_with_mode(&codex_home, store, STORE_LOCK_ACQUIRE_TIMEOUT, mode) } - pub(super) fn acquire_in( + fn acquire_in_with_mode( codex_home: &Path, store: OAuthStore, acquire_timeout: Duration, + mode: OAuthStoreLockMode, ) -> Result { let path = oauth_store_lock_path(codex_home, store); if let Some(parent) = path.parent() { @@ -91,7 +110,11 @@ impl OAuthStoreLock { let mut reported_contention = false; loop { - match file.try_lock() { + let result = match mode { + OAuthStoreLockMode::Shared => file.try_lock_shared(), + OAuthStoreLockMode::Exclusive => file.try_lock(), + }; + match result { Ok(()) => return Ok(Self { _file: file }), Err(std::fs::TryLockError::WouldBlock) if started.elapsed() >= acquire_timeout => { return Err(OAuthStoreLockFailure::Timeout { 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 3278365e01..037f05cff4 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 @@ -27,6 +27,7 @@ use tracing::subscriber::Interest; use super::OAuthStore; use super::OAuthStoreLock; use super::OAuthStoreLockFailure; +use super::OAuthStoreLockMode; use crate::oauth::StoredOAuthTokens; use crate::oauth::WrappedOAuthTokenResponse; use crate::oauth::fallback_file_path; @@ -173,10 +174,11 @@ fn store_lock_is_released_when_holder_process_exits() -> Result<()> { std::thread::sleep(Duration::from_millis(/*millis*/ 20)); } - let error = match OAuthStoreLock::acquire_in( + let error = match OAuthStoreLock::acquire_in_with_mode( env.path(), OAuthStore::File, Duration::from_millis(/*millis*/ 100), + OAuthStoreLockMode::Exclusive, ) { Ok(_) => { anyhow::bail!("live holder process should keep the OAuth store lock unavailable") @@ -192,10 +194,11 @@ fn store_lock_is_released_when_holder_process_exits() -> Result<()> { .wait() .context("wait for killed OAuth store lock holder process")?; assert!(!status.success()); - let _lock = OAuthStoreLock::acquire_in( + let _lock = OAuthStoreLock::acquire_in_with_mode( env.path(), OAuthStore::File, Duration::from_secs(/*secs*/ 1), + OAuthStoreLockMode::Exclusive, )?; Ok(()) })(); @@ -215,7 +218,7 @@ 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(OAuthStore::File)?; + let _lock = OAuthStoreLock::acquire_for_write(OAuthStore::File)?; std::fs::write(ready_file, b"ready")?; loop { std::thread::sleep(Duration::from_secs(/*secs*/ 60)); @@ -328,8 +331,12 @@ where T: Send + 'static, { std::thread::scope(|scope| { - let held_lock = - OAuthStoreLock::acquire_in(codex_home, store, Duration::from_millis(/*millis*/ 100))?; + let held_lock = OAuthStoreLock::acquire_in_with_mode( + codex_home, + store, + Duration::from_millis(/*millis*/ 100), + OAuthStoreLockMode::Exclusive, + )?; let (contended_tx, contended_rx) = mpsc::channel(); let worker = scope.spawn(move || { tracing::subscriber::with_default(LockContentionSubscriber { contended_tx }, operation) @@ -348,6 +355,90 @@ where }) } +#[test] +fn aggregate_store_readers_share_access_while_writers_remain_exclusive() -> Result<()> { + let env = TempCodexHome::new(); + + for store in [OAuthStore::File, OAuthStore::Secrets] { + let readers = (0..2) + .map(|_| { + OAuthStoreLock::acquire_in_with_mode( + env.path(), + store, + Duration::from_millis(/*millis*/ 100), + OAuthStoreLockMode::Shared, + ) + }) + .collect::, _>>()?; + + let writer_error = match OAuthStoreLock::acquire_in_with_mode( + env.path(), + store, + Duration::from_millis(/*millis*/ 100), + OAuthStoreLockMode::Exclusive, + ) { + Ok(_) => anyhow::bail!("an active OAuth store reader must exclude writers"), + Err(error) => error, + }; + assert!(matches!( + writer_error, + OAuthStoreLockFailure::Timeout { .. } + )); + drop(readers); + let _writer = OAuthStoreLock::acquire_in_with_mode( + env.path(), + store, + Duration::from_millis(/*millis*/ 100), + OAuthStoreLockMode::Exclusive, + )?; + } + + Ok(()) +} + +#[test] +fn aggregate_store_credential_loads_can_share_an_existing_reader() -> Result<()> { + let env = TempCodexHome::new(); + let keyring_store = MockKeyringStore::default(); + let tokens = sample_tokens(); + save_oauth_tokens_to_file(&tokens)?; + save_oauth_tokens_with_keyring( + &keyring_store, + AuthKeyringBackendKind::Secrets, + &tokens.server_name, + &tokens, + )?; + + let file_reader = OAuthStoreLock::acquire_in_with_mode( + env.path(), + OAuthStore::File, + Duration::from_millis(/*millis*/ 100), + OAuthStoreLockMode::Shared, + )?; + let file_tokens = load_oauth_tokens_from_file(&tokens.server_name, &tokens.url)? + .expect("a File credential read should coexist with another reader"); + assert_tokens_match_without_expiry(&file_tokens, &tokens); + drop(file_reader); + + let secrets_reader = OAuthStoreLock::acquire_in_with_mode( + env.path(), + OAuthStore::Secrets, + Duration::from_millis(/*millis*/ 100), + OAuthStoreLockMode::Shared, + )?; + let secrets_tokens = load_oauth_tokens_from_keyring( + &keyring_store, + AuthKeyringBackendKind::Secrets, + &tokens.server_name, + &tokens.url, + )? + .expect("a Secrets credential read should coexist with another reader"); + assert_tokens_match_without_expiry(&secrets_tokens, &tokens); + drop(secrets_reader); + + Ok(()) +} + #[test] fn file_store_lock_preserves_updates_for_different_servers() -> Result<()> { let env = TempCodexHome::new(); diff --git a/codex-rs/secrets/src/local.rs b/codex-rs/secrets/src/local.rs index 366be386b9..e76e769dd3 100644 --- a/codex-rs/secrets/src/local.rs +++ b/codex-rs/secrets/src/local.rs @@ -4,6 +4,8 @@ use std::io::Write; use std::path::Path; use std::path::PathBuf; use std::sync::Arc; +use std::sync::Mutex; +use std::sync::PoisonError; use std::sync::atomic::Ordering; use std::sync::atomic::compiler_fence; use std::time::SystemTime; @@ -24,6 +26,8 @@ use rand::TryRngCore; use rand::rngs::OsRng; use serde::Deserialize; use serde::Serialize; +use sha2::Digest; +use sha2::Sha256; use tracing::warn; use super::SecretListEntry; @@ -37,6 +41,7 @@ const SECRETS_VERSION: u8 = 1; const LOCAL_SECRETS_FILENAME: &str = "local.age"; const CODEX_AUTH_SECRETS_FILENAME: &str = "codex_auth.age"; const MCP_OAUTH_SECRETS_FILENAME: &str = "mcp_oauth.age"; +static MCP_OAUTH_CACHE: Mutex> = Mutex::new(None); /// Selects the local encrypted file used by a `LocalSecretsBackend`. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] @@ -56,6 +61,13 @@ struct SecretsFile { secrets: BTreeMap, } +struct CachedMcpSecrets { + path: PathBuf, + ciphertext_hash: [u8; 32], + passphrase_hash: [u8; 32], + file: Arc, +} + impl SecretsFile { fn new_empty() -> Self { Self { @@ -157,6 +169,23 @@ impl LocalSecretsBackend { let ciphertext = fs::read(&path) .with_context(|| format!("failed to read secrets file at {}", path.display()))?; let passphrase = self.load_or_create_passphrase()?; + let cache = (self.namespace == LocalSecretsNamespace::McpOAuth).then(|| { + let ciphertext_hash: [u8; 32] = Sha256::digest(&ciphertext).into(); + let passphrase_hash: [u8; 32] = + Sha256::digest(passphrase.expose_secret().as_bytes()).into(); + let cache = MCP_OAUTH_CACHE + .lock() + .unwrap_or_else(PoisonError::into_inner); + (cache, ciphertext_hash, passphrase_hash) + }); + if let Some((cache, ciphertext_hash, passphrase_hash)) = cache.as_ref() + && let Some(cached) = cache.as_ref() + && cached.path == path + && cached.ciphertext_hash == *ciphertext_hash + && cached.passphrase_hash == *passphrase_hash + { + return Ok(cached.file.as_ref().clone()); + } let plaintext = decrypt_with_passphrase(&ciphertext, &passphrase)?; let mut parsed: SecretsFile = serde_json::from_slice(&plaintext).with_context(|| { format!( @@ -173,6 +202,14 @@ impl LocalSecretsBackend { parsed.version, SECRETS_VERSION ); + if let Some((mut cache, ciphertext_hash, passphrase_hash)) = cache { + *cache = Some(CachedMcpSecrets { + path, + ciphertext_hash, + passphrase_hash, + file: Arc::new(parsed.clone()), + }); + } Ok(parsed) } @@ -186,6 +223,14 @@ impl LocalSecretsBackend { let ciphertext = encrypt_with_passphrase(&plaintext, &passphrase)?; let path = self.secrets_path(); write_file_atomically(&path, &ciphertext)?; + if self.namespace == LocalSecretsNamespace::McpOAuth { + let mut cache = MCP_OAUTH_CACHE + .lock() + .unwrap_or_else(PoisonError::into_inner); + if cache.as_ref().is_some_and(|cached| cached.path == path) { + *cache = None; + } + } Ok(()) } @@ -376,6 +421,8 @@ mod tests { use keyring::Error as KeyringError; use pretty_assertions::assert_eq; + static MCP_OAUTH_CACHE_TEST_LOCK: Mutex<()> = Mutex::new(()); + #[test] fn load_file_rejects_newer_schema_versions() -> Result<()> { let codex_home = tempfile::tempdir().expect("tempdir"); @@ -446,11 +493,21 @@ mod tests { .collect(); assert_eq!(filenames, vec![LOCAL_SECRETS_FILENAME.to_string()]); assert_eq!(backend.get(&scope, &name)?, Some("two".to_string())); + assert!( + MCP_OAUTH_CACHE + .lock() + .unwrap_or_else(PoisonError::into_inner) + .as_ref() + .is_none_or(|cached| cached.path != backend.secrets_path()) + ); Ok(()) } #[test] fn local_namespaces_write_separate_files() -> Result<()> { + let _cache_lock = MCP_OAUTH_CACHE_TEST_LOCK + .lock() + .unwrap_or_else(PoisonError::into_inner); let codex_home = tempfile::tempdir().expect("tempdir"); let keyring = Arc::new(MockKeyringStore::default()); let codex_auth_backend = LocalSecretsBackend::new_with_namespace( @@ -473,6 +530,13 @@ mod tests { codex_auth_backend.get(&scope, &name)?, Some("codex-auth-value".to_string()) ); + assert!( + MCP_OAUTH_CACHE + .lock() + .unwrap_or_else(PoisonError::into_inner) + .as_ref() + .is_none_or(|cached| cached.path != codex_auth_backend.secrets_path()) + ); assert_eq!( mcp_backend.get(&scope, &name)?, Some("mcp-value".to_string()) @@ -494,4 +558,78 @@ mod tests { assert!(!codex_home.path().join("secrets").join("local.age").exists()); Ok(()) } + + #[test] + fn mcp_oauth_cache_reuses_plaintext_and_invalidates_when_ciphertext_changes() -> Result<()> { + let _cache_lock = MCP_OAUTH_CACHE_TEST_LOCK + .lock() + .unwrap_or_else(PoisonError::into_inner); + let codex_home = tempfile::tempdir().expect("tempdir"); + let keyring = Arc::new(MockKeyringStore::default()); + let first = LocalSecretsBackend::new_with_namespace( + codex_home.path().to_path_buf(), + keyring.clone(), + LocalSecretsNamespace::McpOAuth, + ); + let second = LocalSecretsBackend::new_with_namespace( + codex_home.path().to_path_buf(), + keyring, + LocalSecretsNamespace::McpOAuth, + ); + let scope = SecretScope::Global; + let name = SecretName::new("TEST_SECRET")?; + let cached_file = || { + Arc::clone( + &MCP_OAUTH_CACHE + .lock() + .unwrap_or_else(PoisonError::into_inner) + .as_ref() + .expect("MCP OAuth credentials should be cached") + .file, + ) + }; + + first.set(&scope, &name, "one")?; + let (first_cached, second_cached) = std::thread::scope(|threads| { + let first_reader = threads.spawn(|| { + assert_eq!(first.get(&scope, &name)?, Some("one".to_string())); + Ok::<_, anyhow::Error>(cached_file()) + }); + let second_reader = threads.spawn(|| { + assert_eq!(second.get(&scope, &name)?, Some("one".to_string())); + Ok::<_, anyhow::Error>(cached_file()) + }); + Ok::<_, anyhow::Error>(( + first_reader.join().expect("first credential reader")?, + second_reader.join().expect("second credential reader")?, + )) + })?; + assert!(Arc::ptr_eq(&first_cached, &second_cached)); + + assert_eq!(second.get(&scope, &name)?, Some("one".to_string())); + assert!(Arc::ptr_eq(&first_cached, &cached_file())); + + first.set(&scope, &name, "two")?; + assert!( + MCP_OAUTH_CACHE + .lock() + .unwrap_or_else(PoisonError::into_inner) + .as_ref() + .is_none_or(|cached| cached.path != first.secrets_path()) + ); + assert_eq!(second.get(&scope, &name)?, Some("two".to_string())); + assert!(!Arc::ptr_eq(&first_cached, &cached_file())); + + assert!(first.delete(&scope, &name)?); + assert!( + MCP_OAUTH_CACHE + .lock() + .unwrap_or_else(PoisonError::into_inner) + .as_ref() + .is_none_or(|cached| cached.path != first.secrets_path()) + ); + assert_eq!(second.get(&scope, &name)?, None); + + Ok(()) + } }