From f089e6ef96744693e8f7375a7c2962fc9c837709 Mon Sep 17 00:00:00 2001 From: Steven Lee Date: Thu, 18 Jun 2026 20:10:26 +0000 Subject: [PATCH] Keep MCP OAuth refreshes in Codex --- codex-rs/rmcp-client/src/oauth.rs | 356 ++++++++++++------ codex-rs/rmcp-client/src/rmcp_client.rs | 91 +++-- .../rmcp-client/src/streamable_http_retry.rs | 77 ++++ .../tests/streamable_http_oauth_startup.rs | 196 ++++++++++ 4 files changed, 555 insertions(+), 165 deletions(-) diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index 95a5688e6b..769429637e 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -66,8 +66,7 @@ use codex_utils_home_dir::find_codex_home; const KEYRING_SERVICE: &str = "Codex MCP Credentials"; const MCP_OAUTH_SECRET_PREFIX: &str = "MCP_OAUTH"; -// Keep a 30-second guard band ahead of RMCP's 30-second automatic refresh -// buffer so preflight normally refreshes before the request path can do so. +// Refresh proactively so ordinary requests do not race token expiry. const REFRESH_SKEW_MILLIS: u64 = 60_000; const REFRESH_LOCK_DIR: &str = "mcp-oauth-refresh-locks"; const LOCK_ACQUIRE_TIMEOUT: Duration = Duration::from_secs(60); @@ -610,81 +609,13 @@ impl OAuthPersistor { } } - /// 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) + pub(crate) async fn refresh_if_needed(&self) -> Result<()> { + self.refresh_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) { - save_oauth_tokens_with_keyring_store( - keyring_store, - &self.inner.server_name, - &stored, - self.inner.store_mode, - self.inner.keyring_backend_kind, - )?; - *last_credentials = Some(stored); - } - } - None => { - let mut last_serialized = self.inner.last_credentials.lock().await; - if last_serialized.take().is_some() - && let Err(error) = delete_oauth_tokens( - &self.inner.server_name, - &self.inner.url, - self.inner.store_mode, - self.inner.keyring_backend_kind, - ) - { - 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) + pub(crate) async fn refresh_after_unauthorized(&self) -> Result<()> { + self.refresh_after_unauthorized_with_keyring_store(&DefaultKeyringStore) .await } @@ -699,10 +630,6 @@ impl OAuthPersistor { .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, @@ -717,6 +644,43 @@ impl OAuthPersistor { return Ok(()); } + self.refresh_transaction( + keyring_store, + RefreshReason::Expiry, + refresh_request_timeout, + ) + .await + } + + async fn refresh_after_unauthorized_with_keyring_store( + &self, + keyring_store: &K, + ) -> Result<()> { + self.refresh_transaction( + keyring_store, + RefreshReason::Unauthorized, + REFRESH_REQUEST_TIMEOUT, + ) + .await + } + + #[expect( + clippy::await_holding_invalid_type, + reason = "AuthorizationManager async access must be serialized through its mutex" + )] + async fn refresh_transaction( + &self, + keyring_store: &K, + reason: RefreshReason, + refresh_request_timeout: Duration, + ) -> Result<()> { + let local_access_token = { + let last_credentials = self.inner.last_credentials.lock().await; + last_credentials + .as_ref() + .map(|tokens| tokens.token_response.0.access_token().secret().to_string()) + }; + let _lock = RefreshCredentialLock::acquire_for_server(&self.inner.server_name, &self.inner.url) .await?; @@ -738,34 +702,62 @@ impl OAuthPersistor { ); }; - if !token_needs_refresh(latest.expires_at) { + let latest_access_token = latest.token_response.0.access_token().secret(); + let should_adopt = !token_needs_refresh(latest.expires_at) + && match reason { + RefreshReason::Expiry => true, + RefreshReason::Unauthorized => { + local_access_token.as_deref() != Some(latest_access_token) + } + }; + if should_adopt { self.adopt_credentials(latest).await?; return Ok(()); } - self.adopt_credentials(latest).await?; - + let manager = self.inner.authorization_manager.clone(); + let mut guard = manager.lock().await; + if let Err(error) = + install_tokens_in_manager_guard(&mut guard, &latest, CredentialExposure::Refresh).await { - 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 - ), - } + install_tokens_in_manager_guard(&mut guard, &latest, CredentialExposure::Request) + .await + .context("failed to restore request-only OAuth credentials")?; + return Err(error).context("failed to stage OAuth credentials for refresh"); } + let refresh_result = match timeout(refresh_request_timeout, guard.refresh_token()).await { + Ok(Ok(token_response)) => Ok(refreshed_tokens(token_response, &latest, &self.inner)), + Ok(Err(error)) => Err(error).with_context(|| { + format!( + "failed to refresh OAuth tokens for server {}", + self.inner.server_name + ) + }), + Err(_) => Err(anyhow::anyhow!( + "timed out after {refresh_request_timeout:?} refreshing OAuth tokens for server {}", + self.inner.server_name + )), + }; + let request_tokens = refresh_result.as_ref().unwrap_or(&latest); + if let Err(error) = + install_tokens_in_manager_guard(&mut guard, request_tokens, CredentialExposure::Request) + .await + { + return Err(error).context("failed to restore request-only OAuth credentials"); + } + drop(guard); - self.persist_if_needed_with_keyring_store(keyring_store) - .await + let refreshed = refresh_result?; + save_oauth_tokens_with_keyring_store( + keyring_store, + &self.inner.server_name, + &refreshed, + self.inner.store_mode, + self.inner.keyring_backend_kind, + )?; + let mut last_credentials = self.inner.last_credentials.lock().await; + *last_credentials = Some(refreshed); + Ok(()) } async fn adopt_credentials(&self, tokens: StoredOAuthTokens) -> Result<()> { @@ -782,6 +774,12 @@ impl OAuthPersistor { } } +#[derive(Clone, Copy)] +enum RefreshReason { + Expiry, + Unauthorized, +} + struct RefreshCredentialLock { _file: File, } @@ -927,33 +925,56 @@ impl OAuthStoreLock { async fn install_tokens_in_manager( authorization_manager: &Arc>, tokens: &StoredOAuthTokens, +) -> Result<()> { + let manager = authorization_manager.clone(); + let mut guard = manager.lock().await; + install_tokens_in_manager_guard(&mut guard, tokens, CredentialExposure::Request).await +} + +async fn install_tokens_in_manager_guard( + authorization_manager: &mut AuthorizationManager, + tokens: &StoredOAuthTokens, + exposure: CredentialExposure, ) -> Result<()> { let store = InMemoryCredentialStore::new(); store - .save(stored_credentials_from_tokens(tokens)) + .save(stored_credentials_from_tokens(tokens, exposure)) .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); - guard + authorization_manager.set_credential_store(store); + authorization_manager .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(); +#[derive(Clone, Copy)] +enum CredentialExposure { + Request, + Refresh, +} + +fn stored_credentials_from_tokens( + tokens: &StoredOAuthTokens, + exposure: CredentialExposure, +) -> StoredCredentials { + let token_response = match exposure { + CredentialExposure::Request => request_oauth_token_response(tokens), + CredentialExposure::Refresh => 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()); + let token_received_at = match exposure { + CredentialExposure::Request => None, + CredentialExposure::Refresh => SystemTime::now() + .duration_since(UNIX_EPOCH) + .ok() + .map(|duration| duration.as_secs()), + }; StoredCredentials::new( tokens.client_id.clone(), @@ -963,6 +984,34 @@ fn stored_credentials_from_tokens(tokens: &StoredOAuthTokens) -> StoredCredentia ) } +pub(crate) fn request_oauth_token_response(tokens: &StoredOAuthTokens) -> OAuthTokenResponse { + let mut token_response = tokens.token_response.0.clone(); + token_response.set_refresh_token(None); + token_response.set_expires_in(None); + token_response +} + +fn refreshed_tokens( + mut token_response: OAuthTokenResponse, + previous: &StoredOAuthTokens, + inner: &OAuthPersistorInner, +) -> StoredOAuthTokens { + if token_response.refresh_token().is_none() { + token_response.set_refresh_token(previous.token_response.0.refresh_token().cloned()); + } + if token_response.scopes().is_none() { + token_response.set_scopes(previous.token_response.0.scopes().cloned()); + } + let expires_at = compute_expires_at_millis(&token_response); + StoredOAuthTokens { + server_name: inner.server_name.clone(), + url: inner.url.clone(), + client_id: previous.client_id.clone(), + token_response: WrappedOAuthTokenResponse(token_response), + expires_at, + } +} + const FALLBACK_FILENAME: &str = ".credentials.json"; const MCP_SERVER_TYPE: &str = "http"; @@ -1646,11 +1695,84 @@ mod tests { ); let loser_tokens = tokens_from_manager(&loser_manager).await?; - assert_eq!(access_token(&loser_tokens), "refreshed-access-token"); assert_eq!( - refresh_token(&loser_tokens), + loser_tokens.token_response, + WrappedOAuthTokenResponse(request_oauth_token_response(&stored)) + ); + Ok(()) + } + + #[tokio::test] + async fn unauthorized_transaction_without_expiry_refreshes_once_across_clients() -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + mount_refresh_response( + &server, + "refresh-token", + "refreshed-access-token", + "rotated-refresh-token", + ) + .await; + + let store = MockKeyringStore::default(); + let mut initial_tokens = sample_tokens(); + initial_tokens.url = format!("{}/mcp", server.uri()); + initial_tokens.expires_at = None; + initial_tokens.token_response.0.set_expires_in(None); + super::save_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + + let first_manager = authorization_manager_for(&initial_tokens).await?; + let second_manager = authorization_manager_for(&initial_tokens).await?; + let first = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + first_manager, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + Some(initial_tokens.clone()), + ); + let second = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + second_manager.clone(), + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + Some(initial_tokens.clone()), + ); + + first + .refresh_after_unauthorized_with_keyring_store(&store) + .await?; + second + .refresh_after_unauthorized_with_keyring_store(&store) + .await?; + + server.verify().await; + let stored = super::load_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens.url, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )? + .expect("refreshed tokens should be persisted"); + assert_eq!(access_token(&stored), "refreshed-access-token"); + assert_eq!( + refresh_token(&stored), Some("rotated-refresh-token".to_string()) ); + let second_tokens = tokens_from_manager(&second_manager).await?; + assert_eq!( + second_tokens.token_response, + WrappedOAuthTokenResponse(request_oauth_token_response(&stored)) + ); Ok(()) } @@ -1705,12 +1827,8 @@ mod tests { let manager_tokens = tokens_from_manager(&manager).await?; assert_eq!( - access_token(&manager_tokens), - "already-refreshed-access-token" - ); - assert_eq!( - refresh_token(&manager_tokens), - Some("already-rotated-refresh-token".to_string()) + manager_tokens.token_response, + WrappedOAuthTokenResponse(request_oauth_token_response(&latest_tokens)) ); Ok(()) } @@ -2504,7 +2622,7 @@ mod tests { ) -> Result>> { let mut state = OAuthState::new(tokens.url.clone(), Some(reqwest::Client::new())).await?; state - .set_credentials(&tokens.client_id, tokens.token_response.0.clone()) + .set_credentials(&tokens.client_id, request_oauth_token_response(tokens)) .await?; let manager = match state { OAuthState::Authorized(manager) | OAuthState::Unauthorized(manager) => manager, diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index fad6fb0200..3bd83e63f7 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -68,6 +68,7 @@ use crate::in_process_transport::InProcessTransportFactory; use crate::load_oauth_tokens; use crate::oauth::OAuthPersistor; use crate::oauth::StoredOAuthTokens; +use crate::oauth::request_oauth_token_response; use crate::stdio_server_launcher::StdioServerCommand; use crate::stdio_server_launcher::StdioServerLauncher; use crate::stdio_server_launcher::StdioServerProcessHandle; @@ -445,7 +446,7 @@ impl RmcpClient { }; let (service, oauth_persistor) = self - .connect_pending_transport_with_initialize_retries( + .connect_pending_transport_with_oauth_recovery( pending_transport, client_service.clone(), timeout, @@ -477,12 +478,6 @@ impl RmcpClient { }; } - if let Some(runtime) = oauth_persistor - && let Err(error) = runtime.persist_if_needed().await - { - warn!("failed to persist OAuth tokens after initialize: {error}"); - } - Ok(initialize_result) } @@ -498,7 +493,6 @@ impl RmcpClient { async move { service.list_tools(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -532,7 +526,6 @@ impl RmcpClient { }) }) .collect::>>()?; - self.persist_oauth_tokens().await; Ok(ListToolsWithConnectorIdResult { next_cursor: result.next_cursor, tools, @@ -559,7 +552,6 @@ impl RmcpClient { async move { service.list_resources(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -575,7 +567,6 @@ impl RmcpClient { async move { service.list_resource_templates(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -591,7 +582,6 @@ impl RmcpClient { async move { service.read_resource(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -649,7 +639,6 @@ impl RmcpClient { .boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -679,7 +668,6 @@ impl RmcpClient { }, ) .await?; - self.persist_oauth_tokens().await; Ok(()) } @@ -702,7 +690,6 @@ impl RmcpClient { .boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(response) } @@ -742,16 +729,6 @@ impl RmcpClient { drop(previous_state); } - /// This should be called after every tool call so that if a given tool call triggered - /// a refresh of the OAuth tokens, they are persisted. - async fn persist_oauth_tokens(&self) { - if let Some(runtime) = self.oauth_persistor().await - && let Err(error) = runtime.persist_if_needed().await - { - warn!("failed to persist OAuth tokens: {error}"); - } - } - async fn refresh_oauth_if_needed(&self) -> Result<()> { if let Some(runtime) = self.oauth_persistor().await { runtime.refresh_if_needed().await?; @@ -914,19 +891,7 @@ impl RmcpClient { .await .map_err(|source| anyhow::Error::from(HandshakeError { source })), }; - let service = match service_result { - Ok(service) => service, - Err(error) => { - if let Some(runtime) = oauth_persistor.as_ref() - && let Err(persist_error) = runtime.persist_if_needed().await - { - warn!( - "failed to persist OAuth tokens after failed initialize: {persist_error}" - ); - } - return Err(error); - } - }; + let service = service_result?; Ok((Arc::new(service), oauth_persistor)) } @@ -952,6 +917,21 @@ impl RmcpClient { .await { Ok(result) => Ok(result), + Err(error) if Self::is_unauthorized_operation_error(&error) => { + let Some(oauth_persistor) = self.oauth_persistor().await else { + return Err(error.into()); + }; + oauth_persistor.refresh_after_unauthorized().await?; + Self::run_service_operation_with_transient_retries( + service, + label, + timeout, + self.elicitation_pause_state.clone(), + &operation, + ) + .await + .map_err(Into::into) + } Err(error) if Self::is_session_expired_404(&error) => { self.reinitialize_after_session_expiry(&service).await?; let recovered_service = self.service().await?; @@ -1086,6 +1066,31 @@ impl RmcpClient { }) } + fn is_unauthorized_operation_error(error: &ClientOperationError) -> bool { + let ClientOperationError::Service(rmcp::service::ServiceError::TransportSend(error)) = + error + else { + return false; + }; + + error + .error + .downcast_ref::>() + .is_some_and(Self::is_unauthorized_streamable_http_error) + } + + pub(super) fn is_unauthorized_streamable_http_error( + error: &StreamableHttpError, + ) -> bool { + match error { + StreamableHttpError::AuthRequired(_) => true, + StreamableHttpError::UnexpectedServerResponse(message) => { + message.starts_with("HTTP 401") + } + _ => false, + } + } + async fn reinitialize_after_session_expiry( &self, failed_service: &Arc>, @@ -1120,7 +1125,7 @@ impl RmcpClient { .ok_or_else(|| anyhow!("MCP client cannot recover before initialize succeeds"))?; let pending_transport = Self::create_pending_transport(&self.transport_recipe).await?; let (service, oauth_persistor) = self - .connect_pending_transport_with_initialize_retries( + .connect_pending_transport_with_oauth_recovery( pending_transport, initialize_context.client_service, initialize_context.timeout, @@ -1138,12 +1143,6 @@ impl RmcpClient { }; } - if let Some(runtime) = oauth_persistor - && let Err(error) = runtime.persist_if_needed().await - { - warn!("failed to persist OAuth tokens after session recovery: {error}"); - } - Ok(()) } } @@ -1174,7 +1173,7 @@ async fn create_oauth_transport_and_runtime( oauth_state .set_credentials( &initial_tokens.client_id, - initial_tokens.token_response.0.clone(), + request_oauth_token_response(&initial_tokens), ) .await?; diff --git a/codex-rs/rmcp-client/src/streamable_http_retry.rs b/codex-rs/rmcp-client/src/streamable_http_retry.rs index 73da95de58..24b73afab3 100644 --- a/codex-rs/rmcp-client/src/streamable_http_retry.rs +++ b/codex-rs/rmcp-client/src/streamable_http_retry.rs @@ -23,6 +23,60 @@ const JSON_RPC_INTERNAL_ERROR_CODE: i64 = -32603; pub(super) const STREAMABLE_HTTP_RETRY_DELAYS_MS: [u64; 2] = [250, 1_000]; impl RmcpClient { + pub(super) async fn connect_pending_transport_with_oauth_recovery( + &self, + initial_transport: PendingTransport, + client_service: ElicitationClientService, + timeout: Option, + ) -> Result<( + Arc>, + Option, + )> { + let oauth_persistor = match &initial_transport { + PendingTransport::StreamableHttpWithOAuth { + oauth_persistor, .. + } => Some(oauth_persistor.clone()), + _ => None, + }; + let deadline = timeout.map(|duration| Instant::now() + duration); + match self + .connect_pending_transport_with_initialize_retries( + initial_transport, + client_service.clone(), + timeout, + ) + .await + { + Ok(result) => Ok(result), + Err(error) + if oauth_persistor.is_some() && Self::is_unauthorized_initialize_error(&error) => + { + let Some(oauth_persistor) = oauth_persistor else { + return Err(error); + }; + oauth_persistor.refresh_after_unauthorized().await?; + let remaining = remaining_initialize_timeout(timeout, deadline)?; + let transport = match remaining { + Some(remaining) => time::timeout( + remaining, + Self::create_pending_transport(&self.transport_recipe), + ) + .await + .map_err(|_| initialize_timeout_error(timeout, remaining))??, + None => Self::create_pending_transport(&self.transport_recipe).await?, + }; + let remaining = remaining_initialize_timeout(timeout, deadline)?; + self.connect_pending_transport_with_initialize_retries( + transport, + client_service, + remaining, + ) + .await + } + Err(error) => Err(error), + } + } + pub(super) async fn connect_pending_transport_with_initialize_retries( &self, initial_transport: PendingTransport, @@ -109,6 +163,29 @@ impl RmcpClient { }) } + fn is_unauthorized_initialize_error(error: &anyhow::Error) -> bool { + error.chain().any(|source| { + source + .downcast_ref::() + .is_some_and(|error| Self::is_unauthorized_client_initialize_error(&error.source)) + || source + .downcast_ref::() + .is_some_and(Self::is_unauthorized_client_initialize_error) + }) + } + + fn is_unauthorized_client_initialize_error( + error: &rmcp::service::ClientInitializeError, + ) -> bool { + match error { + rmcp::service::ClientInitializeError::TransportError { error, .. } => error + .error + .downcast_ref::>() + .is_some_and(Self::is_unauthorized_streamable_http_error), + _ => false, + } + } + fn is_retryable_client_initialize_error(error: &rmcp::service::ClientInitializeError) -> bool { match error { rmcp::service::ClientInitializeError::TransportError { error, context } diff --git a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs index 8ba5168cab..4ec89ddc89 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs @@ -39,6 +39,8 @@ const EXPIRED_ACCESS_TOKEN: &str = "expired-access-token"; const REFRESH_TOKEN: &str = "valid-refresh-token"; const REFRESHED_ACCESS_TOKEN: &str = "refreshed-access-token"; const ROTATED_REFRESH_TOKEN: &str = "rotated-refresh-token"; +const FINAL_ACCESS_TOKEN: &str = "final-access-token"; +const FINAL_REFRESH_TOKEN: &str = "final-refresh-token"; const PREFLIGHT_REFRESH_ERROR: &str = "preflight refresh failed distinctly"; const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_STARTUP_SERVER_URL"; const UNREFRESHABLE_SERVER_URL: &str = "https://unrefreshable.example/mcp"; @@ -123,6 +125,125 @@ async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn recovers_initialization_and_operation_401_without_rmcp_refresh() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": [""], + }))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={REFRESH_TOKEN}" + ))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": REFRESHED_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": ROTATED_REFRESH_TOKEN, + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={ROTATED_REFRESH_TOKEN}" + ))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": FINAL_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": FINAL_REFRESH_TOKEN, + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {EXPIRED_ACCESS_TOKEN}"), + )) + .respond_with( + ResponseTemplate::new(401).insert_header("www-authenticate", "Bearer realm=test"), + ) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {REFRESHED_ACCESS_TOKEN}"), + )) + .respond_with(|request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body.get("method").and_then(Value::as_str) { + Some("initialize") => initialize_response(&body, "oauth-401-recovery-test"), + Some("notifications/initialized") => ResponseTemplate::new(202), + Some("tools/list") => ResponseTemplate::new(401) + .insert_header("www-authenticate", "Bearer realm=test"), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + }) + .expect(3) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {FINAL_ACCESS_TOKEN}"), + )) + .respond_with(|request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body.get("method").and_then(Value::as_str) { + Some("initialize") => initialize_response(&body, "oauth-401-persisted-test"), + Some("notifications/initialized") => ResponseTemplate::new(202), + Some("tools/list") => ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { "tools": [] }, + })), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + }) + .expect(3) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + let server_url = format!("{}/mcp", server.uri()); + let status = Command::new(std::env::current_exe()?) + .args([ + "oauth_401_recovery_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(CHILD_SERVER_URL_ENV, server_url) + .status() + .await?; + assert!( + status.success(), + "OAuth 401 recovery child failed: {status}" + ); + server.verify().await; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn concurrent_file_mode_startup_refreshes_once() -> anyhow::Result<()> { let server = MockServer::start().await; @@ -473,6 +594,25 @@ async fn oauth_startup_child() -> anyhow::Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by recovers_initialization_and_operation_401_without_rmcp_refresh"] +async fn oauth_401_recovery_child() -> anyhow::Result<()> { + let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; + save_no_expiry_file_mode_tokens(&server_url)?; + + let client = new_file_mode_oauth_client(&server_url).await?; + initialize_client(&client).await?; + let tools = client + .list_tools(/*params*/ None, Some(Duration::from_secs(5))) + .await?; + assert!(tools.tools.is_empty()); + client.shutdown().await; + + let reloaded_client = new_file_mode_oauth_client(&server_url).await?; + initialize_client(&reloaded_client).await?; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] #[ignore = "spawned by operation_preflight_refresh_failure_blocks_rmcp_request"] async fn operation_preflight_refresh_failure_child() -> anyhow::Result<()> { @@ -591,3 +731,59 @@ fn save_near_expiry_file_mode_tokens(server_url: &str) -> anyhow::Result<()> { )?; Ok(()) } + +fn save_no_expiry_file_mode_tokens(server_url: &str) -> anyhow::Result<()> { + let mut response = OAuthTokenResponse::new( + AccessToken::new(EXPIRED_ACCESS_TOKEN.to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + response.set_refresh_token(Some(RefreshToken::new(REFRESH_TOKEN.to_string()))); + let tokens = StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: server_url.to_string(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: None, + }; + save_oauth_tokens( + SERVER_NAME, + &tokens, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + Ok(()) +} + +async fn new_file_mode_oauth_client(server_url: &str) -> anyhow::Result { + RmcpClient::new_streamable_http_client( + SERVER_NAME, + server_url, + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await +} + +fn initialize_response(body: &Value, name: &str) -> ResponseTemplate { + ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { + "protocolVersion": body + .pointer("/params/protocolVersion") + .cloned() + .unwrap_or_else(|| json!("2025-06-18")), + "capabilities": {}, + "serverInfo": { + "name": name, + "version": "0.0.0-test", + }, + }, + })) +}