mirror of
https://github.com/openai/codex.git
synced 2026-10-22 17:58:46 +00:00
Matthew Zeng
88912c04cd
## What changed - Honor `McpOAuthRefreshMode::Coordinated` for streamable HTTP MCP connections by letting RMCP refresh and persist credentials through the pinned credential store. - Reread credentials while holding the refresh lock, persist refreshed tokens before retrying requests, and serialize login and logout writes with in-flight refreshes. - Keep OAuth preparation outside MCP operation and handshake deadlines, and allow refresh persistence to finish if the initiating caller is cancelled. ## Testing - Cover coordinated startup refresh, concurrent refreshes, `401` retries, credential mutation races, provider timeouts, and caller cancellation. GitOrigin-RevId: 100ae788048087be23e26bb43f868b1fcebb007a
65 lines
2.2 KiB
Rust
65 lines
2.2 KiB
Rust
//! Keeps one OAuth refresh and persistence owner for the lifetime of a connection.
|
|
//!
|
|
//! Both modes use cached expiry to prepare credentials before the MCP request budget.
|
|
//! Coordinated preparation owns a task so caller cancellation cannot interrupt persistence.
|
|
//! It locks the manager before the credential store, matching RMCP's lock order.
|
|
//! RMCP commits through the pinned store while Legacy retains Codex's persistor.
|
|
|
|
use std::sync::Arc;
|
|
|
|
use anyhow::Context;
|
|
use anyhow::Result;
|
|
use rmcp::transport::auth::AuthorizationManager;
|
|
use tokio::sync::Mutex;
|
|
use tracing::warn;
|
|
|
|
use super::OAuthPersistor;
|
|
use super::credential_store::OAuthCredentialStore;
|
|
use super::token_needs_refresh;
|
|
|
|
#[derive(Clone)]
|
|
pub(crate) enum OAuthRuntime {
|
|
Legacy(OAuthPersistor),
|
|
Coordinated {
|
|
auth_manager: Arc<Mutex<AuthorizationManager>>,
|
|
store: OAuthCredentialStore,
|
|
},
|
|
}
|
|
|
|
impl OAuthRuntime {
|
|
#[expect(
|
|
clippy::await_holding_invalid_type,
|
|
reason = "AuthorizationManager access follows RMCP's manager-before-credential lock order"
|
|
)]
|
|
pub(crate) async fn refresh_if_needed(&self) -> Result<()> {
|
|
match self {
|
|
Self::Legacy(persistor) => persistor.refresh_if_needed().await,
|
|
Self::Coordinated {
|
|
auth_manager,
|
|
store,
|
|
} => {
|
|
let expires_at = store
|
|
.stored_credentials()
|
|
.await
|
|
.and_then(|tokens| tokens.expires_at);
|
|
if !token_needs_refresh(expires_at) {
|
|
return Ok(());
|
|
}
|
|
let auth_manager = Arc::clone(auth_manager);
|
|
let store = store.clone();
|
|
tokio::spawn(async move {
|
|
let mut manager = auth_manager.lock().await;
|
|
let result = store.refresh_if_needed(&mut manager).await;
|
|
if result.is_err() {
|
|
// Keep the summary in the owned task without logging credential/provider data.
|
|
warn!("MCP OAuth preparation failed");
|
|
}
|
|
result
|
|
})
|
|
.await
|
|
.context("OAuth refresh task failed")?
|
|
}
|
|
}
|
|
}
|
|
}
|