openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-04 15:08:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

Split MCP OAuth lifecycle modules

Browse Source
This commit is contained in:
Steven Lee
2026-06-25 02:24:38 +00:00
parent 04cfc03825
commit e563b66000
3 changed files with 490 additions and 431 deletions
Download Patch File Download Diff File Expand all files Collapse all files

445
codex-rs/rmcp-client/src/oauth.rs
View File

@@ -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<Option<LoadedOAuthTokens>> {
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<K: KeyringStore + Clone + 'static>(
keyring_store: &K,
server_name: &str,
url: &str,
store_mode: OAuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
) -> Result<Option<LoadedOAuthTokens>> {
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<Option<StoredOAuthTokens>> {
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<K: KeyringStore + Clone + 'static>(
keyring_store: &K,
keyring_backend_kind: AuthKeyringBackendKind,
server_name: &str,
url: &str,
) -> Result<Option<LoadedOAuthTokens>> {
// 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<K: KeyringStore + Clone + 'static>(
keyring_store: &K,
keyring_backend_kind: AuthKeyringBackendKind,
@@ -578,303 +458,6 @@ fn delete_oauth_tokens_from_secrets_keyring<K: KeyringStore + Clone + 'static>(
Ok(secrets_removed)
}
#[derive(Clone)]
pub(crate) struct OAuthPersistor {
inner: Arc<OAuthPersistorInner>,
}
struct OAuthPersistorInner {
server_name: String,
url: String,
authorization_manager: Arc<Mutex<AuthorizationManager>>,
credential_store: ResolvedOAuthCredentialStore,
last_credentials: Mutex<Option<StoredOAuthTokens>>,
}
impl OAuthPersistor {
pub(crate) fn new(
server_name: String,
url: String,
authorization_manager: Arc<Mutex<AuthorizationManager>>,
credential_store: ResolvedOAuthCredentialStore,
initial_credentials: Option<StoredOAuthTokens>,
) -> 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<K: KeyringStore + Clone + 'static>(
&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<K: KeyringStore + Clone + 'static>(
&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<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<Mutex<AuthorizationManager>>,
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";

336
codex-rs/rmcp-client/src/oauth/persistor.rs Normal file
View File

@@ -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<OAuthPersistorInner>,
}
struct OAuthPersistorInner {
server_name: String,
url: String,
authorization_manager: Arc<Mutex<AuthorizationManager>>,
credential_store: ResolvedOAuthCredentialStore,
last_credentials: Mutex<Option<StoredOAuthTokens>>,
}
impl OAuthPersistor {
pub(crate) fn new(
server_name: String,
url: String,
authorization_manager: Arc<Mutex<AuthorizationManager>>,
credential_store: ResolvedOAuthCredentialStore,
initial_credentials: Option<StoredOAuthTokens>,
) -> 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<K: KeyringStore + Clone + 'static>(
&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<K: KeyringStore + Clone + 'static>(
&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<Mutex<AuthorizationManager>>,
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,
)
}

140
codex-rs/rmcp-client/src/oauth/resolved_store.rs Normal file
View File

@@ -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<Option<LoadedOAuthTokens>> {
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<K: KeyringStore + Clone + 'static>(
keyring_store: &K,
server_name: &str,
url: &str,
store_mode: OAuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
) -> Result<Option<LoadedOAuthTokens>> {
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<Option<StoredOAuthTokens>> {
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<Option<LoadedOAuthTokens>> {
// 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,
}))
}
}
}