From e563b6600090c997efdcfd28374fbf00525517d2 Mon Sep 17 00:00:00 2001 From: Steven Lee Date: Thu, 25 Jun 2026 02:24:38 +0000 Subject: [PATCH] Split MCP OAuth lifecycle modules --- codex-rs/rmcp-client/src/oauth.rs | 445 +----------------- codex-rs/rmcp-client/src/oauth/persistor.rs | 336 +++++++++++++ .../rmcp-client/src/oauth/resolved_store.rs | 140 ++++++ 3 files changed, 490 insertions(+), 431 deletions(-) create mode 100644 codex-rs/rmcp-client/src/oauth/persistor.rs create mode 100644 codex-rs/rmcp-client/src/oauth/resolved_store.rs diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index 1799f2d2e6..369d0939d5 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -16,7 +16,9 @@ //! //! 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 persistor; mod refresh_lock; +mod resolved_store; use anyhow::Context; use anyhow::Error; @@ -53,20 +55,23 @@ use tracing::warn; use codex_keyring_store::DefaultKeyringStore; use codex_keyring_store::KeyringStore; -use rmcp::transport::auth::AuthorizationManager; -use rmcp::transport::auth::CredentialStore as _; -use rmcp::transport::auth::InMemoryCredentialStore; -use rmcp::transport::auth::StoredCredentials; -use tokio::sync::Mutex; -use tokio::time::timeout; - -use self::refresh_lock::RefreshCredentialLock; use codex_utils_home_dir::find_codex_home; +pub(crate) use self::persistor::OAuthPersistor; +pub(crate) use self::resolved_store::LoadedOAuthTokens; +pub(crate) use self::resolved_store::ResolvedOAuthCredentialStore; +#[cfg(test)] +use self::resolved_store::load_oauth_tokens_from_keyring_with_fallback_to_file; +pub(crate) use self::resolved_store::load_oauth_tokens_from_resolved_store; +pub(crate) use self::resolved_store::load_oauth_tokens_with_source; +#[cfg(test)] +use self::resolved_store::load_oauth_tokens_with_source_and_keyring_store; +#[cfg(test)] +use rmcp::transport::auth::AuthorizationManager; + const KEYRING_SERVICE: &str = "Codex MCP Credentials"; const MCP_OAUTH_SECRET_PREFIX: &str = "MCP_OAUTH"; const REFRESH_SKEW_MILLIS: u64 = 30_000; -const REFRESH_REQUEST_TIMEOUT: Duration = Duration::from_secs(45); #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct StoredOAuthTokens { @@ -91,24 +96,6 @@ impl PartialEq for WrappedOAuthTokenResponse { } } -/// Concrete credential store resolved for one MCP OAuth client lifecycle. -/// -/// This is intentionally not durable. `Auto` may resolve differently in a later process, but a -/// client that loaded credentials from one store must reread, refresh, persist, and remove only -/// through that store. A mid-lifecycle backend failure is unexpected and must return an error -/// rather than falling back to another possibly stale refresh token. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum ResolvedOAuthCredentialStore { - File, - Keyring(AuthKeyringBackendKind), -} - -#[derive(Debug)] -pub(crate) struct LoadedOAuthTokens { - pub(crate) tokens: StoredOAuthTokens, - pub(crate) store: ResolvedOAuthCredentialStore, -} - #[derive(Debug, PartialEq, Eq)] pub(crate) enum StoredOAuthTokenStatus { Missing, @@ -128,78 +115,6 @@ pub(crate) fn load_oauth_tokens( ) } -pub(crate) fn load_oauth_tokens_with_source( - server_name: &str, - url: &str, - store_mode: OAuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, -) -> Result> { - let keyring_store = DefaultKeyringStore; - load_oauth_tokens_with_source_and_keyring_store( - &keyring_store, - server_name, - url, - store_mode, - keyring_backend_kind, - ) -} - -fn load_oauth_tokens_with_source_and_keyring_store( - keyring_store: &K, - server_name: &str, - url: &str, - store_mode: OAuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, -) -> Result> { - match store_mode { - OAuthCredentialsStoreMode::Auto => load_oauth_tokens_from_keyring_with_fallback_to_file( - keyring_store, - keyring_backend_kind, - server_name, - url, - ), - OAuthCredentialsStoreMode::File => Ok(load_oauth_tokens_from_file(server_name, url)?.map( - |tokens| LoadedOAuthTokens { - tokens, - store: ResolvedOAuthCredentialStore::File, - }, - )), - OAuthCredentialsStoreMode::Keyring => Ok(load_oauth_tokens_from_keyring( - keyring_store, - keyring_backend_kind, - server_name, - url, - ) - .with_context(|| "failed to read OAuth tokens from keyring".to_string())? - .map(|tokens| LoadedOAuthTokens { - tokens, - store: ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind), - })), - } -} - -pub(crate) fn load_oauth_tokens_from_resolved_store( - server_name: &str, - url: &str, - store: ResolvedOAuthCredentialStore, -) -> Result> { - match store { - ResolvedOAuthCredentialStore::File => load_oauth_tokens_from_file(server_name, url) - .context("failed to read OAuth tokens from resolved file storage"), - ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind) => { - load_oauth_tokens_from_keyring( - &DefaultKeyringStore, - keyring_backend_kind, - server_name, - url, - ) - .context( - "failed to read OAuth tokens from resolved keyring storage; refusing file fallback", - ) - } - } -} - pub(crate) fn oauth_token_status( server_name: &str, url: &str, @@ -252,41 +167,6 @@ fn refresh_expires_in_from_timestamp(tokens: &mut StoredOAuthTokens) { } } -fn load_oauth_tokens_from_keyring_with_fallback_to_file( - keyring_store: &K, - keyring_backend_kind: AuthKeyringBackendKind, - server_name: &str, - url: &str, -) -> Result> { - // Auto remains keyring-first at lifecycle startup. The returned source is then pinned by the - // client transport recipe and OAuth persistor so retries, recovery, and refresh work cannot - // hot-switch stores. - // TODO(stevenlee): Different processes can still resolve Auto to different stores when - // keyring availability differs. Solving that safely requires durable backend selection or - // reconciliation of legacy entries and is intentionally outside this stack. - match load_oauth_tokens_from_keyring(keyring_store, keyring_backend_kind, server_name, url) { - Ok(Some(tokens)) => Ok(Some(LoadedOAuthTokens { - tokens, - store: ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind), - })), - Ok(None) => Ok( - load_oauth_tokens_from_file(server_name, url)?.map(|tokens| LoadedOAuthTokens { - tokens, - store: ResolvedOAuthCredentialStore::File, - }), - ), - Err(error) => { - warn!("failed to read OAuth tokens from keyring: {error}"); - Ok(load_oauth_tokens_from_file(server_name, url) - .with_context(|| format!("failed to read OAuth tokens from keyring: {error}"))? - .map(|tokens| LoadedOAuthTokens { - tokens, - store: ResolvedOAuthCredentialStore::File, - })) - } - } -} - fn load_oauth_tokens_from_keyring( keyring_store: &K, keyring_backend_kind: AuthKeyringBackendKind, @@ -578,303 +458,6 @@ fn delete_oauth_tokens_from_secrets_keyring( Ok(secrets_removed) } -#[derive(Clone)] -pub(crate) struct OAuthPersistor { - inner: Arc, -} - -struct OAuthPersistorInner { - server_name: String, - url: String, - authorization_manager: Arc>, - credential_store: ResolvedOAuthCredentialStore, - last_credentials: Mutex>, -} - -impl OAuthPersistor { - pub(crate) fn new( - server_name: String, - url: String, - authorization_manager: Arc>, - credential_store: ResolvedOAuthCredentialStore, - initial_credentials: Option, - ) -> Self { - Self { - inner: Arc::new(OAuthPersistorInner { - server_name, - url, - authorization_manager, - credential_store, - last_credentials: Mutex::new(initial_credentials), - }), - } - } - - /// Persists the latest stored credentials if they have changed. - /// Deletes the credentials if they are no longer present. - pub(crate) async fn persist_if_needed(&self) -> Result<()> { - self.persist_if_needed_with_keyring_store(&DefaultKeyringStore) - .await - } - - #[expect( - clippy::await_holding_invalid_type, - reason = "AuthorizationManager async access must be serialized through its mutex" - )] - async fn persist_if_needed_with_keyring_store( - &self, - keyring_store: &K, - ) -> Result<()> { - let (client_id, maybe_credentials) = { - let manager = self.inner.authorization_manager.clone(); - let guard = manager.lock().await; - guard.get_credentials().await - }?; - - match maybe_credentials { - Some(credentials) => { - let mut last_credentials = self.inner.last_credentials.lock().await; - let new_token_response = WrappedOAuthTokenResponse(credentials.clone()); - let same_token = last_credentials - .as_ref() - .map(|prev| prev.token_response == new_token_response) - .unwrap_or(false); - let expires_at = if same_token { - last_credentials.as_ref().and_then(|prev| prev.expires_at) - } else { - compute_expires_at_millis(&credentials) - }; - let stored = StoredOAuthTokens { - server_name: self.inner.server_name.clone(), - url: self.inner.url.clone(), - client_id, - token_response: new_token_response, - expires_at, - }; - if last_credentials.as_ref() != Some(&stored) { - match self.inner.credential_store { - ResolvedOAuthCredentialStore::File => save_oauth_tokens_to_file(&stored)?, - ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind) => { - save_oauth_tokens_with_keyring( - keyring_store, - keyring_backend_kind, - &self.inner.server_name, - &stored, - )?; - } - } - *last_credentials = Some(stored); - } - } - None => { - let mut last_serialized = self.inner.last_credentials.lock().await; - if last_serialized.take().is_some() - && let Err(error) = match self.inner.credential_store { - ResolvedOAuthCredentialStore::File => { - let key = compute_store_key(&self.inner.server_name, &self.inner.url)?; - delete_oauth_tokens_from_file(&key).map(|_| ()) - } - ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct) => { - delete_oauth_tokens_from_direct_keyring( - keyring_store, - &self.inner.server_name, - &self.inner.url, - ) - .map(|_| ()) - } - ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Secrets) => { - delete_oauth_tokens_from_secrets_keyring( - keyring_store, - &self.inner.server_name, - &self.inner.url, - ) - .map(|_| ()) - } - } - { - warn!( - "failed to remove OAuth tokens for server {}: {error}", - self.inner.server_name - ); - } - } - } - - Ok(()) - } - - pub(crate) async fn refresh_if_needed(&self) -> Result<()> { - self.refresh_if_needed_with_keyring_store(&DefaultKeyringStore) - .await - } - - async fn refresh_if_needed_with_keyring_store( - &self, - keyring_store: &K, - ) -> Result<()> { - self.refresh_if_needed_with_keyring_store_and_timeout( - keyring_store, - REFRESH_REQUEST_TIMEOUT, - ) - .await - } - - #[expect( - clippy::await_holding_invalid_type, - reason = "AuthorizationManager async access must be serialized through its mutex" - )] - async fn refresh_if_needed_with_keyring_store_and_timeout( - &self, - keyring_store: &K, - refresh_request_timeout: Duration, - ) -> Result<()> { - let expires_at = { - let guard = self.inner.last_credentials.lock().await; - guard.as_ref().and_then(|tokens| tokens.expires_at) - }; - - if !token_needs_refresh(expires_at) { - return Ok(()); - } - - let snapshot = { - let guard = self.inner.last_credentials.lock().await; - guard.clone() - }; - let key = compute_store_key(&self.inner.server_name, &self.inner.url)?; - let _lock = RefreshCredentialLock::acquire(&key).await?; - // The refresh transaction must stay on the store that supplied its snapshot. Falling back - // here could replay an older rotating refresh token from the other store. We assume store - // availability is stable for this client lifecycle and surface violations of that - // assumption instead of switching stores. - let latest = match self.inner.credential_store { - ResolvedOAuthCredentialStore::File => { - load_oauth_tokens_from_file(&self.inner.server_name, &self.inner.url) - .context("failed to reread OAuth tokens from resolved file storage")? - } - ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind) => { - load_oauth_tokens_from_keyring( - keyring_store, - keyring_backend_kind, - &self.inner.server_name, - &self.inner.url, - ) - .context( - "failed to reread OAuth tokens from resolved keyring storage; refusing file fallback", - )? - } - }; - - if latest.is_none() && snapshot.is_some() { - self.clear_manager_credentials().await; - let mut last_credentials = self.inner.last_credentials.lock().await; - *last_credentials = None; - anyhow::bail!( - "OAuth tokens for server {} were removed before refresh; authorization required", - self.inner.server_name - ); - } - - if latest != snapshot { - if let Some(latest) = latest { - let needs_refresh = token_needs_refresh(latest.expires_at); - self.adopt_credentials(latest).await?; - // `expires_in` is derived from `expires_at` on each load and can drift without a - // persisted change. Even for a real concurrent update, keep going when the - // authoritative token is still inside the refresh window. - if !needs_refresh { - return Ok(()); - } - } else { - return Ok(()); - } - } - - { - let manager = self.inner.authorization_manager.clone(); - let guard = manager.lock().await; - match timeout(refresh_request_timeout, guard.refresh_token()).await { - Ok(result) => { - result.with_context(|| { - format!( - "failed to refresh OAuth tokens for server {}", - self.inner.server_name - ) - })?; - } - Err(_) => anyhow::bail!( - "timed out after {refresh_request_timeout:?} refreshing OAuth tokens for server {}", - self.inner.server_name - ), - } - } - - // Once the provider returns a rotated token, persistence must finish before the credential - // lock is released. In particular, caller startup deadlines must not cancel this step. - self.persist_if_needed_with_keyring_store(keyring_store) - .await - } - - async fn adopt_credentials(&self, tokens: StoredOAuthTokens) -> Result<()> { - install_tokens_in_manager(&self.inner.authorization_manager, &tokens).await?; - let mut last_credentials = self.inner.last_credentials.lock().await; - *last_credentials = Some(tokens); - Ok(()) - } - - async fn clear_manager_credentials(&self) { - let manager = self.inner.authorization_manager.clone(); - let mut guard = manager.lock().await; - guard.set_credential_store(InMemoryCredentialStore::new()); - } -} - -#[expect( - clippy::await_holding_invalid_type, - reason = "AuthorizationManager async access must be serialized through its mutex" -)] -async fn install_tokens_in_manager( - authorization_manager: &Arc>, - tokens: &StoredOAuthTokens, -) -> Result<()> { - let store = InMemoryCredentialStore::new(); - store - .save(stored_credentials_from_tokens(tokens)) - .await - .context("failed to stage OAuth tokens for authorization manager")?; - - let manager = authorization_manager.clone(); - let mut guard = manager.lock().await; - guard.set_credential_store(store); - // TODO(stevenlee): RMCP's `initialize_from_store` updates the credential store and client ID - // but not its private `current_scopes`. Credential adoption can therefore leave scope-upgrade - // state stale until RMCP exposes an adoption API that synchronizes both. - guard - .initialize_from_store() - .await - .context("failed to adopt refreshed OAuth tokens")?; - Ok(()) -} - -fn stored_credentials_from_tokens(tokens: &StoredOAuthTokens) -> StoredCredentials { - let token_response = tokens.token_response.0.clone(); - let granted_scopes = token_response - .scopes() - .map(|scopes| scopes.iter().map(|scope| scope.to_string()).collect()) - .unwrap_or_default(); - let token_received_at = SystemTime::now() - .duration_since(UNIX_EPOCH) - .ok() - .map(|duration| duration.as_secs()); - - StoredCredentials::new( - tokens.client_id.clone(), - Some(token_response), - granted_scopes, - token_received_at, - ) -} - const FALLBACK_FILENAME: &str = ".credentials.json"; const MCP_SERVER_TYPE: &str = "http"; diff --git a/codex-rs/rmcp-client/src/oauth/persistor.rs b/codex-rs/rmcp-client/src/oauth/persistor.rs new file mode 100644 index 0000000000..303fdb0d97 --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/persistor.rs @@ -0,0 +1,336 @@ +//! Lifecycle-local persistence and serialized refresh transactions for MCP OAuth credentials. + +use std::sync::Arc; +use std::time::Duration; +use std::time::SystemTime; +use std::time::UNIX_EPOCH; + +use anyhow::Context; +use anyhow::Result; +use codex_config::types::AuthKeyringBackendKind; +use codex_keyring_store::DefaultKeyringStore; +use codex_keyring_store::KeyringStore; +use oauth2::TokenResponse; +use rmcp::transport::auth::AuthorizationManager; +use rmcp::transport::auth::CredentialStore as _; +use rmcp::transport::auth::InMemoryCredentialStore; +use rmcp::transport::auth::StoredCredentials; +use tokio::sync::Mutex; +use tokio::time::timeout; +use tracing::warn; + +use super::ResolvedOAuthCredentialStore; +use super::StoredOAuthTokens; +use super::WrappedOAuthTokenResponse; +use super::compute_expires_at_millis; +use super::compute_store_key; +use super::delete_oauth_tokens_from_direct_keyring; +use super::delete_oauth_tokens_from_file; +use super::delete_oauth_tokens_from_secrets_keyring; +use super::load_oauth_tokens_from_file; +use super::load_oauth_tokens_from_keyring; +use super::refresh_lock::RefreshCredentialLock; +use super::save_oauth_tokens_to_file; +use super::save_oauth_tokens_with_keyring; +use super::token_needs_refresh; + +const REFRESH_REQUEST_TIMEOUT: Duration = Duration::from_secs(45); + +#[derive(Clone)] +pub(crate) struct OAuthPersistor { + inner: Arc, +} + +struct OAuthPersistorInner { + server_name: String, + url: String, + authorization_manager: Arc>, + credential_store: ResolvedOAuthCredentialStore, + last_credentials: Mutex>, +} + +impl OAuthPersistor { + pub(crate) fn new( + server_name: String, + url: String, + authorization_manager: Arc>, + credential_store: ResolvedOAuthCredentialStore, + initial_credentials: Option, + ) -> Self { + Self { + inner: Arc::new(OAuthPersistorInner { + server_name, + url, + authorization_manager, + credential_store, + last_credentials: Mutex::new(initial_credentials), + }), + } + } + + /// Persists the latest stored credentials if they have changed. + /// Deletes the credentials if they are no longer present. + pub(crate) async fn persist_if_needed(&self) -> Result<()> { + self.persist_if_needed_with_keyring_store(&DefaultKeyringStore) + .await + } + + #[expect( + clippy::await_holding_invalid_type, + reason = "AuthorizationManager async access must be serialized through its mutex" + )] + pub(super) async fn persist_if_needed_with_keyring_store( + &self, + keyring_store: &K, + ) -> Result<()> { + let (client_id, maybe_credentials) = { + let manager = self.inner.authorization_manager.clone(); + let guard = manager.lock().await; + guard.get_credentials().await + }?; + + match maybe_credentials { + Some(credentials) => { + let mut last_credentials = self.inner.last_credentials.lock().await; + let new_token_response = WrappedOAuthTokenResponse(credentials.clone()); + let same_token = last_credentials + .as_ref() + .map(|prev| prev.token_response == new_token_response) + .unwrap_or(false); + let expires_at = if same_token { + last_credentials.as_ref().and_then(|prev| prev.expires_at) + } else { + compute_expires_at_millis(&credentials) + }; + let stored = StoredOAuthTokens { + server_name: self.inner.server_name.clone(), + url: self.inner.url.clone(), + client_id, + token_response: new_token_response, + expires_at, + }; + if last_credentials.as_ref() != Some(&stored) { + match self.inner.credential_store { + ResolvedOAuthCredentialStore::File => save_oauth_tokens_to_file(&stored)?, + ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind) => { + save_oauth_tokens_with_keyring( + keyring_store, + keyring_backend_kind, + &self.inner.server_name, + &stored, + )?; + } + } + *last_credentials = Some(stored); + } + } + None => { + let mut last_serialized = self.inner.last_credentials.lock().await; + if last_serialized.take().is_some() + && let Err(error) = match self.inner.credential_store { + ResolvedOAuthCredentialStore::File => { + let key = compute_store_key(&self.inner.server_name, &self.inner.url)?; + delete_oauth_tokens_from_file(&key).map(|_| ()) + } + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct) => { + delete_oauth_tokens_from_direct_keyring( + keyring_store, + &self.inner.server_name, + &self.inner.url, + ) + .map(|_| ()) + } + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Secrets) => { + delete_oauth_tokens_from_secrets_keyring( + keyring_store, + &self.inner.server_name, + &self.inner.url, + ) + .map(|_| ()) + } + } + { + warn!( + "failed to remove OAuth tokens for server {}: {error}", + self.inner.server_name + ); + } + } + } + + Ok(()) + } + + pub(crate) async fn refresh_if_needed(&self) -> Result<()> { + self.refresh_if_needed_with_keyring_store(&DefaultKeyringStore) + .await + } + + pub(super) async fn refresh_if_needed_with_keyring_store( + &self, + keyring_store: &K, + ) -> Result<()> { + self.refresh_if_needed_with_keyring_store_and_timeout( + keyring_store, + REFRESH_REQUEST_TIMEOUT, + ) + .await + } + + #[expect( + clippy::await_holding_invalid_type, + reason = "AuthorizationManager async access must be serialized through its mutex" + )] + pub(super) async fn refresh_if_needed_with_keyring_store_and_timeout< + K: KeyringStore + Clone + 'static, + >( + &self, + keyring_store: &K, + refresh_request_timeout: Duration, + ) -> Result<()> { + let expires_at = { + let guard = self.inner.last_credentials.lock().await; + guard.as_ref().and_then(|tokens| tokens.expires_at) + }; + + if !token_needs_refresh(expires_at) { + return Ok(()); + } + + let snapshot = { + let guard = self.inner.last_credentials.lock().await; + guard.clone() + }; + let key = compute_store_key(&self.inner.server_name, &self.inner.url)?; + let _lock = RefreshCredentialLock::acquire(&key).await?; + // The refresh transaction must stay on the store that supplied its snapshot. Falling back + // here could replay an older rotating refresh token from the other store. We assume store + // availability is stable for this client lifecycle and surface violations of that + // assumption instead of switching stores. + let latest = match self.inner.credential_store { + ResolvedOAuthCredentialStore::File => { + load_oauth_tokens_from_file(&self.inner.server_name, &self.inner.url) + .context("failed to reread OAuth tokens from resolved file storage")? + } + ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind) => { + load_oauth_tokens_from_keyring( + keyring_store, + keyring_backend_kind, + &self.inner.server_name, + &self.inner.url, + ) + .context( + "failed to reread OAuth tokens from resolved keyring storage; refusing file fallback", + )? + } + }; + + if latest.is_none() && snapshot.is_some() { + self.clear_manager_credentials().await; + let mut last_credentials = self.inner.last_credentials.lock().await; + *last_credentials = None; + anyhow::bail!( + "OAuth tokens for server {} were removed before refresh; authorization required", + self.inner.server_name + ); + } + + if latest != snapshot { + if let Some(latest) = latest { + let needs_refresh = token_needs_refresh(latest.expires_at); + self.adopt_credentials(latest).await?; + // `expires_in` is derived from `expires_at` on each load and can drift without a + // persisted change. Even for a real concurrent update, keep going when the + // authoritative token is still inside the refresh window. + if !needs_refresh { + return Ok(()); + } + } else { + return Ok(()); + } + } + + { + let manager = self.inner.authorization_manager.clone(); + let guard = manager.lock().await; + match timeout(refresh_request_timeout, guard.refresh_token()).await { + Ok(result) => { + result.with_context(|| { + format!( + "failed to refresh OAuth tokens for server {}", + self.inner.server_name + ) + })?; + } + Err(_) => anyhow::bail!( + "timed out after {refresh_request_timeout:?} refreshing OAuth tokens for server {}", + self.inner.server_name + ), + } + } + + // Once the provider returns a rotated token, persistence must finish before the credential + // lock is released. In particular, caller startup deadlines must not cancel this step. + self.persist_if_needed_with_keyring_store(keyring_store) + .await + } + + async fn adopt_credentials(&self, tokens: StoredOAuthTokens) -> Result<()> { + install_tokens_in_manager(&self.inner.authorization_manager, &tokens).await?; + let mut last_credentials = self.inner.last_credentials.lock().await; + *last_credentials = Some(tokens); + Ok(()) + } + + async fn clear_manager_credentials(&self) { + let manager = self.inner.authorization_manager.clone(); + let mut guard = manager.lock().await; + guard.set_credential_store(InMemoryCredentialStore::new()); + } +} + +#[expect( + clippy::await_holding_invalid_type, + reason = "AuthorizationManager async access must be serialized through its mutex" +)] +async fn install_tokens_in_manager( + authorization_manager: &Arc>, + tokens: &StoredOAuthTokens, +) -> Result<()> { + let store = InMemoryCredentialStore::new(); + store + .save(stored_credentials_from_tokens(tokens)) + .await + .context("failed to stage OAuth tokens for authorization manager")?; + + let manager = authorization_manager.clone(); + let mut guard = manager.lock().await; + guard.set_credential_store(store); + // TODO(stevenlee): RMCP's `initialize_from_store` updates the credential store and client ID + // but not its private `current_scopes`. Credential adoption can therefore leave scope-upgrade + // state stale until RMCP exposes an adoption API that synchronizes both. + guard + .initialize_from_store() + .await + .context("failed to adopt refreshed OAuth tokens")?; + Ok(()) +} + +fn stored_credentials_from_tokens(tokens: &StoredOAuthTokens) -> StoredCredentials { + let token_response = tokens.token_response.0.clone(); + let granted_scopes = token_response + .scopes() + .map(|scopes| scopes.iter().map(|scope| scope.to_string()).collect()) + .unwrap_or_default(); + let token_received_at = SystemTime::now() + .duration_since(UNIX_EPOCH) + .ok() + .map(|duration| duration.as_secs()); + + StoredCredentials::new( + tokens.client_id.clone(), + Some(token_response), + granted_scopes, + token_received_at, + ) +} diff --git a/codex-rs/rmcp-client/src/oauth/resolved_store.rs b/codex-rs/rmcp-client/src/oauth/resolved_store.rs new file mode 100644 index 0000000000..3919521883 --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/resolved_store.rs @@ -0,0 +1,140 @@ +//! Resolves the configured MCP OAuth store and pins that concrete source for one client lifecycle. + +use anyhow::Context; +use anyhow::Result; +use codex_config::types::AuthKeyringBackendKind; +use codex_config::types::OAuthCredentialsStoreMode; +use codex_keyring_store::DefaultKeyringStore; +use codex_keyring_store::KeyringStore; +use tracing::warn; + +use super::StoredOAuthTokens; +use super::load_oauth_tokens_from_file; +use super::load_oauth_tokens_from_keyring; + +/// Concrete credential store resolved for one MCP OAuth client lifecycle. +/// +/// This is intentionally not durable. `Auto` may resolve differently in a later process, but a +/// client that loaded credentials from one store must reread, refresh, persist, and remove only +/// through that store. A mid-lifecycle backend failure is unexpected and must return an error +/// rather than falling back to another possibly stale refresh token. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ResolvedOAuthCredentialStore { + File, + Keyring(AuthKeyringBackendKind), +} + +#[derive(Debug)] +pub(crate) struct LoadedOAuthTokens { + pub(crate) tokens: StoredOAuthTokens, + pub(crate) store: ResolvedOAuthCredentialStore, +} + +pub(crate) fn load_oauth_tokens_with_source( + server_name: &str, + url: &str, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> Result> { + let keyring_store = DefaultKeyringStore; + load_oauth_tokens_with_source_and_keyring_store( + &keyring_store, + server_name, + url, + store_mode, + keyring_backend_kind, + ) +} + +pub(super) fn load_oauth_tokens_with_source_and_keyring_store( + keyring_store: &K, + server_name: &str, + url: &str, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> Result> { + match store_mode { + OAuthCredentialsStoreMode::Auto => load_oauth_tokens_from_keyring_with_fallback_to_file( + keyring_store, + keyring_backend_kind, + server_name, + url, + ), + OAuthCredentialsStoreMode::File => Ok(load_oauth_tokens_from_file(server_name, url)?.map( + |tokens| LoadedOAuthTokens { + tokens, + store: ResolvedOAuthCredentialStore::File, + }, + )), + OAuthCredentialsStoreMode::Keyring => Ok(load_oauth_tokens_from_keyring( + keyring_store, + keyring_backend_kind, + server_name, + url, + ) + .with_context(|| "failed to read OAuth tokens from keyring".to_string())? + .map(|tokens| LoadedOAuthTokens { + tokens, + store: ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind), + })), + } +} + +pub(crate) fn load_oauth_tokens_from_resolved_store( + server_name: &str, + url: &str, + store: ResolvedOAuthCredentialStore, +) -> Result> { + match store { + ResolvedOAuthCredentialStore::File => load_oauth_tokens_from_file(server_name, url) + .context("failed to read OAuth tokens from resolved file storage"), + ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind) => { + load_oauth_tokens_from_keyring( + &DefaultKeyringStore, + keyring_backend_kind, + server_name, + url, + ) + .context( + "failed to read OAuth tokens from resolved keyring storage; refusing file fallback", + ) + } + } +} + +pub(super) fn load_oauth_tokens_from_keyring_with_fallback_to_file< + K: KeyringStore + Clone + 'static, +>( + keyring_store: &K, + keyring_backend_kind: AuthKeyringBackendKind, + server_name: &str, + url: &str, +) -> Result> { + // Auto remains keyring-first at lifecycle startup. The returned source is then pinned by the + // client transport recipe and OAuth persistor so retries, recovery, and refresh work cannot + // hot-switch stores. + // TODO(stevenlee): Different processes can still resolve Auto to different stores when + // keyring availability differs. Solving that safely requires durable backend selection or + // reconciliation of legacy entries and is intentionally outside this stack. + match load_oauth_tokens_from_keyring(keyring_store, keyring_backend_kind, server_name, url) { + Ok(Some(tokens)) => Ok(Some(LoadedOAuthTokens { + tokens, + store: ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind), + })), + Ok(None) => Ok( + load_oauth_tokens_from_file(server_name, url)?.map(|tokens| LoadedOAuthTokens { + tokens, + store: ResolvedOAuthCredentialStore::File, + }), + ), + Err(error) => { + warn!("failed to read OAuth tokens from keyring: {error}"); + Ok(load_oauth_tokens_from_file(server_name, url) + .with_context(|| format!("failed to read OAuth tokens from keyring: {error}"))? + .map(|tokens| LoadedOAuthTokens { + tokens, + store: ResolvedOAuthCredentialStore::File, + })) + } + } +}