From b3ffe3d00177a6fce0d9895b245ccf9203b2570f Mon Sep 17 00:00:00 2001 From: jif Date: Thu, 6 Aug 2026 17:20:57 +0000 Subject: [PATCH] Recover MCP servers after OAuth reauthentication (#37337) ## Why An OAuth-backed Streamable HTTP MCP server can fail startup when its stored credentials are rejected. If a client then completes OAuth sign-in and replaces those credentials, the failed server needs to become available without restarting the Codex session. ## What changed - Track the credential store and OAuth tokens used for each MCP connection, and detect replacements after authentication failures. - Refresh affected MCP servers before the next turn and include their recovered tools in the model request. - Classify HTTP authentication challenges and rejected refresh tokens as requiring reauthentication, with sign-in guidance appropriate to local and client-managed environments. ## Testing - Cover recovery from expired credentials through an externally written replacement token. - Cover authentication-error classification, startup messaging, and `Auto` store migration from file credentials to the keyring. GitOrigin-RevId: 9b090089d4e7a44ffc182f4e25da20a18e97b70f --- codex-rs/codex-mcp/src/connection_manager.rs | 56 +++++++++++++- .../src/connection_manager/startup.rs | 26 ++++--- .../codex-mcp/src/connection_manager_tests.rs | 35 +++++++-- codex-rs/codex-mcp/src/runtime.rs | 23 ++++++ codex-rs/codex-mcp/src/server.rs | 54 +++++++++++++- codex-rs/core/src/session/mcp.rs | 23 +++++- codex-rs/core/tests/suite/rmcp_client.rs | 73 ++++++++++++++++--- .../src/bin/test_streamable_http_server.rs | 14 +++- codex-rs/rmcp-client/src/lib.rs | 2 + codex-rs/rmcp-client/src/oauth.rs | 52 ++++++++++++- codex-rs/rmcp-client/src/rmcp_client.rs | 2 +- codex-rs/rmcp-client/src/startup_error.rs | 36 ++++++--- .../src/streamable_http_retry_tests.rs | 32 ++++++++ .../streamable_http_oauth_store_pinning.rs | 37 ++++++++-- 14 files changed, 408 insertions(+), 57 deletions(-) diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 3beb4d4895..d342d33a65 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -97,7 +97,7 @@ impl McpServerConnection { return Some(client); }; let reusable = match client.client.managed_oauth_credentials().await { - Some(live_credentials) => &live_credentials == desired_credentials, + Some(live_credentials) => live_credentials.as_ref() == desired_credentials, None => current .oauth_credentials() .is_ok_and(|startup_credentials| startup_credentials == desired_credentials), @@ -645,6 +645,60 @@ impl McpConnectionSet { self.servers.contains_key(server_name) } + pub(crate) async fn authentication_failed_servers(&self) -> Vec { + let mut failed_servers = Vec::new(); + for (server_name, view) in &self.servers { + if view + .connection + .client + .startup_complete + .load(Ordering::Acquire) + && let Err(error) = view.connection.client().await + && error.is_authentication_required() + { + failed_servers.push(server_name.clone()); + } + } + failed_servers + } + + pub(crate) async fn updated_oauth_credentials_after_auth_failure( + &self, + config: &crate::McpConfig, + ) -> Vec { + let mut candidates = Vec::new(); + for server_name in self.authentication_failed_servers().await { + if let Some(view) = self.servers.get(&server_name) + && let Some(identity) = view.connection.identity.as_ref() + && let Some(server) = config.mcp_server_catalog.server(&server_name) + { + candidates.push((server_name, identity.clone(), server.config().clone())); + } + } + if candidates.is_empty() { + return Vec::new(); + } + + match tokio::task::spawn_blocking(move || { + candidates + .into_iter() + .filter_map(|(server_name, identity, config)| { + identity + .oauth_credentials_changed(&server_name, &config) + .then_some(server_name) + }) + .collect() + }) + .await + { + Ok(recovered_servers) => recovered_servers, + Err(error) => { + warn!(%error, "failed to inspect stored MCP OAuth credentials"); + Vec::new() + } + } + } + pub(crate) async fn wait_for_server_startup(&self, server_name: &str) -> bool { let Some(view) = self.servers.get(server_name) else { return false; diff --git a/codex-rs/codex-mcp/src/connection_manager/startup.rs b/codex-rs/codex-mcp/src/connection_manager/startup.rs index b06401d83e..fc933e1574 100644 --- a/codex-rs/codex-mcp/src/connection_manager/startup.rs +++ b/codex-rs/codex-mcp/src/connection_manager/startup.rs @@ -59,15 +59,14 @@ pub(super) fn mcp_startup_failure_reason( return None; } match auth_state { - Some(McpAuthState::LoggedOut(McpLoginRequirement::Reauthentication)) => { - Some(McpStartupFailureReason::ReauthenticationRequired) - } + Some( + McpAuthState::LoggedOut(McpLoginRequirement::Reauthentication) | McpAuthState::OAuth, + ) => Some(McpStartupFailureReason::ReauthenticationRequired), Some( McpAuthState::Unsupported | McpAuthState::Unknown | McpAuthState::LoggedOut(McpLoginRequirement::Login) - | McpAuthState::BearerToken - | McpAuthState::OAuth, + | McpAuthState::BearerToken, ) | None => None, } @@ -91,13 +90,18 @@ pub(super) fn mcp_init_error_display( format!( "GitHub MCP does not support OAuth. Log in by adding a personal access token (https://github.com/settings/personal-access-tokens) to your environment and config.toml:\n[mcp_servers.{server_name}]\nbearer_token_env_var = CODEX_GITHUB_PERSONAL_ACCESS_TOKEN" ) - } else if matches!( - error, - StartupOutcomeError::Failed { error, .. } if error.contains("Auth required") - ) { - format!( - "The {server_name} MCP server is not logged in. Run `codex mcp login {server_name}`." + } else if error.is_authentication_required() + || matches!( + error, + StartupOutcomeError::Failed { error, .. } if error.contains("Auth required") ) + { + let recovery_hint = if config.is_some_and(|config| !config.is_local_environment()) { + "Use your client's MCP OAuth sign-in flow.".to_string() + } else { + format!("Run `codex mcp login {server_name}`.") + }; + format!("The {server_name} MCP server is not logged in. {recovery_hint}") } else if matches!( error, StartupOutcomeError::Failed { error, .. } diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 87b9d2e138..571bbb808a 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -3305,15 +3305,33 @@ fn mcp_init_error_display_prompts_for_github_pat() { #[test] fn mcp_init_error_display_prompts_for_login_when_auth_required() { let server_name = "example"; - let err: StartupOutcomeError = anyhow::anyhow!("Auth required for server").into(); - - let display = mcp_init_error_display(server_name, /*config*/ None, &err); - let expected = format!( "The {server_name} MCP server is not logged in. Run `codex mcp login {server_name}`." ); + let executor_config: McpServerConfig = serde_json::from_value(serde_json::json!({ + "url": "https://example.com/mcp", + "environment_id": "executor-1", + })) + .expect("executor MCP configuration should deserialize"); - assert_eq!(expected, display); + for error in [ + anyhow::anyhow!("Auth required for server").into(), + StartupOutcomeError::Failed { + error: "OAuth refresh token was rejected: invalid_grant".to_string(), + is_authentication_required: true, + }, + ] { + let display = mcp_init_error_display(server_name, /*config*/ None, &error); + assert_eq!(expected, display); + + let executor_display = mcp_init_error_display(server_name, Some(&executor_config), &error); + assert_eq!( + format!( + "The {server_name} MCP server is not logged in. Use your client's MCP OAuth sign-in flow." + ), + executor_display + ); + } } #[test] @@ -3340,7 +3358,12 @@ fn mcp_startup_failure_reason_requires_existing_oauth_and_auth_failure() { ), (Some(McpAuthState::Unsupported), true, None), (Some(McpAuthState::BearerToken), true, None), - (Some(McpAuthState::OAuth), true, None), + ( + Some(McpAuthState::OAuth), + true, + Some(McpStartupFailureReason::ReauthenticationRequired), + ), + (Some(McpAuthState::OAuth), false, None), (None, true, None), ] { let error = StartupOutcomeError::Failed { diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index 6d9ba4316f..c7c1a44bdd 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -305,6 +305,29 @@ impl McpRuntime { } } + /// Detects newly saved credentials for servers whose startup failed authentication. + pub async fn updated_oauth_credentials_after_auth_failure(&self) -> Vec { + let current = self.current.load_full(); + let Some(config) = current.config.as_ref() else { + return Vec::new(); + }; + current + .connections + .updated_oauth_credentials_after_auth_failure(config) + .await + } + + /// Checks the current generation before retrying servers detected outside the refresh gate. + pub async fn has_authentication_failed_servers(&self, server_names: &[String]) -> bool { + self.current + .load_full() + .connections + .authentication_failed_servers() + .await + .into_iter() + .any(|server_name| server_names.contains(&server_name)) + } + /// Waits for the selected server without capturing an execution binding. pub async fn wait_for_server_startup(&self, server: &str) { self.current diff --git a/codex-rs/codex-mcp/src/server.rs b/codex-rs/codex-mcp/src/server.rs index e12f906e97..b788051d2c 100644 --- a/codex-rs/codex-mcp/src/server.rs +++ b/codex-rs/codex-mcp/src/server.rs @@ -15,7 +15,9 @@ use codex_connectors::ConnectorRuntimeContextKey; use codex_exec_server::Environment; use codex_login::CodexAuth; use codex_protocol::mcp::ClientMcpExtensions; +use codex_rmcp_client::StoredOAuthCredentialSnapshot; use codex_rmcp_client::StoredOAuthTokens; +use codex_rmcp_client::stored_oauth_credential_snapshot; use codex_rmcp_client::stored_oauth_credentials; use rmcp::model::ElicitationCapability; use tracing::warn; @@ -96,7 +98,7 @@ pub(crate) struct McpServerConnectionIdentity { transport: McpServerTransportConfig, environment_id: String, oauth_store: Option<(OAuthCredentialsStoreMode, AuthKeyringBackendKind)>, - oauth_credentials: Result, String>, + oauth_credentials: Result, String>, resolved_environment: Result>, String>, local_stdio_fallback_cwd: Option, referenced_environment_variables: Vec<(String, Option)>, @@ -143,7 +145,7 @@ impl McpServerConnectionIdentity { None }; let oauth_credentials = stored_oauth_url.map_or(Ok(None), |url| { - stored_oauth_credentials( + stored_oauth_credential_snapshot( config.oauth_credential_name(server_name).as_ref(), url, store_mode, @@ -211,8 +213,52 @@ impl McpServerConnectionIdentity { && self.agent_plugin == other.agent_plugin } - pub(crate) fn oauth_credentials(&self) -> Result<&Option, &String> { - self.oauth_credentials.as_ref() + pub(crate) fn oauth_credentials(&self) -> Result, &String> { + self.oauth_credentials.as_ref().map(|credentials| { + credentials + .as_ref() + .map(StoredOAuthCredentialSnapshot::credentials) + }) + } + + pub(crate) fn oauth_credentials_changed( + &self, + server_name: &str, + config: &McpServerConfig, + ) -> bool { + let Some((store_mode, keyring_backend_kind)) = self.oauth_store else { + return false; + }; + let McpServerTransportConfig::StreamableHttp { url, .. } = &self.transport else { + return false; + }; + + let credential_name = config.oauth_credential_name(server_name); + let current_credentials = match self.oauth_credentials.as_ref() { + Ok(Some(credentials)) => credentials.reload( + credential_name.as_ref(), + url, + store_mode, + keyring_backend_kind, + ), + Ok(None) | Err(_) => stored_oauth_credentials( + credential_name.as_ref(), + url, + store_mode, + keyring_backend_kind, + ), + }; + + match current_credentials { + Ok(Some(current_credentials)) => { + self.oauth_credentials() != Ok(Some(¤t_credentials)) + } + Ok(None) => false, + Err(error) => { + warn!(server_name, %error, "failed to read stored MCP OAuth credentials"); + false + } + } } } diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index eb70a4cd06..3c4a2e8443 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -302,11 +302,32 @@ impl Session { { self.mark_mcp_runtime_dirty(); } + + let recovered_oauth_servers = self + .services + .mcp_runtime + .updated_oauth_credentials_after_auth_failure() + .await; + if !recovered_oauth_servers.is_empty() + && let Ok(_refresh) = self.mcp_refresh.acquire().await + && self + .services + .mcp_runtime + .has_authentication_failed_servers(&recovered_oauth_servers) + .await + { + self.mark_mcp_runtime_dirty(); + } self.refresh_mcp_if_dirty().await; + let required_servers = required_servers + .iter() + .chain(&recovered_oauth_servers) + .cloned() + .collect::>(); if let Some(binding) = self .services .mcp_runtime - .current_binding_with_required_servers(required_servers) + .current_binding_with_required_servers(&required_servers) .await { return binding; diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index 6efcf10a02..7085cd29ca 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -50,6 +50,7 @@ use codex_protocol::openai_models::TruncationPolicyConfig; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::McpInvocation; +use codex_protocol::protocol::McpStartupFailureReason; use codex_protocol::protocol::McpStartupStatus; use codex_protocol::protocol::McpToolCallBeginEvent; use codex_protocol::protocol::Op; @@ -3077,7 +3078,7 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> { let server_name = "rmcp_http_oauth"; let namespace = format!("mcp__{server_name}"); - mount_sse_once( + let response_mock = mount_sse_once( &server, responses::sse(vec![ responses::ev_response_created("resp-1"), @@ -3130,13 +3131,15 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> { credential_name.as_ref(), &server_url, client_id, - expected_token, + "expired-access-token", refresh_token, + OAuthCredentialExpiry::Expired, )?; // Phase 4: configure Codex with the OAuth-backed Streamable HTTP MCP // server and build the fixture in the active local or remote-aware mode. let fixture = test_codex() + .with_model_info_override("gpt-5.4", |model| model.supports_search_tool = false) .with_home(temp_home.clone()) .with_config(move |config| { // Keep OAuth credentials isolated to this test home because Bazel @@ -3160,9 +3163,35 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> { }) .build_with_auto_env(&server) .await?; - // Phase 5: wait for MCP startup before the turn is submitted, which keeps - // failures tied to server startup/discovery. - wait_for_mcp_server(&fixture.codex, server_name).await?; + // Phase 5: replace rejected credentials as an external OAuth login would. + let mut failure_reason = None; + let startup = wait_for_event(&fixture.codex, |event| { + if let EventMsg::McpStartupUpdate(update) = event + && update.server == server_name + && let McpStartupStatus::Failed { reason, .. } = &update.status + { + failure_reason = *reason; + } + matches!(event, EventMsg::McpStartupComplete(_)) + }) + .await; + let EventMsg::McpStartupComplete(startup) = startup else { + unreachable!("event guard guarantees McpStartupComplete"); + }; + assert_eq!(startup.failed.len(), 1); + assert_eq!(startup.failed[0].server, server_name); + assert_eq!( + failure_reason, + Some(McpStartupFailureReason::ReauthenticationRequired), + ); + write_fallback_oauth_tokens( + credential_name.as_ref(), + http_server.url(), + client_id, + expected_token, + refresh_token, + OAuthCredentialExpiry::Valid, + )?; // Phase 6: submit the user turn that should invoke the OAuth-backed tool. fixture @@ -3175,12 +3204,18 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> { // Phase 7: assert Codex begins the expected tool invocation. let begin_event = wait_for_event(&fixture.codex, |ev| { - matches!(ev, EventMsg::McpToolCallBegin(_)) + matches!( + ev, + EventMsg::McpToolCallBegin(_) + | EventMsg::Error(_) + | EventMsg::TurnAborted(_) + | EventMsg::TurnComplete(_) + ) }) .await; let EventMsg::McpToolCallBegin(begin) = begin_event else { - unreachable!("event guard guarantees McpToolCallBegin"); + anyhow::bail!("OAuth MCP recovery ended before the tool was called: {begin_event:?}"); }; assert_eq!(begin.invocation.server, server_name); assert_eq!(begin.invocation.tool, "echo"); @@ -3227,6 +3262,11 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> { // placement-aware MCP server. wait_for_event(&fixture.codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await; + let request = response_mock.single_request().body_json(); + assert!( + responses::namespace_child_tool(&request, &namespace, "echo").is_some(), + "the recovered MCP tool must be advertised to the model" + ); server.verify().await; http_server.shutdown().await; @@ -3584,18 +3624,27 @@ fn streamable_http_metadata_url(server_url: &str) -> String { format!("{base_url}{STREAMABLE_HTTP_METADATA_PATH}") } +enum OAuthCredentialExpiry { + Valid, + Expired, +} + fn write_fallback_oauth_tokens( server_name: &str, server_url: &str, client_id: &str, access_token: &str, refresh_token: &str, + expiry: OAuthCredentialExpiry, ) -> anyhow::Result<()> { - let expires_at = SystemTime::now() - .checked_add(Duration::from_secs(3600)) - .ok_or_else(|| anyhow::anyhow!("failed to compute expiry time"))? - .duration_since(UNIX_EPOCH)? - .as_millis() as u64; + let expires_at = match expiry { + OAuthCredentialExpiry::Valid => SystemTime::now() + .checked_add(Duration::from_secs(3600)) + .ok_or_else(|| anyhow::anyhow!("failed to compute expiry time"))? + .duration_since(UNIX_EPOCH)? + .as_millis() as u64, + OAuthCredentialExpiry::Expired => 0, + }; let tokens = serde_json::from_value(json!({ "server_name": server_name, diff --git a/codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs b/codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs index e4e5dd98a9..c63026af22 100644 --- a/codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs +++ b/codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs @@ -177,6 +177,18 @@ async fn main() -> Result<(), Box> { } }), ) + .route( + "/oauth/token", + post(|| async { + ( + StatusCode::BAD_REQUEST, + Json(json!({ + "error": "invalid_grant", + "error_description": "refresh token expired or revoked", + })), + ) + }), + ) .nest_service( "/mcp", StreamableHttpService::new( @@ -386,7 +398,7 @@ async fn require_bearer( request: Request, next: Next, ) -> Result { - if request.uri().path().contains("/.well-known/") { + if request.uri().path().contains("/.well-known/") || request.uri().path() == "/oauth/token" { return Ok(next.run(request).await); } if request diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index 80ac068684..e476e4b029 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -26,10 +26,12 @@ pub use auth_status::discover_streamable_http_oauth; pub use codex_protocol::protocol::McpAuthStatus; pub use http_client_adapter::StreamableHttpRedirectMode; pub use in_process_transport::InProcessTransportFactory; +pub use oauth::StoredOAuthCredentialSnapshot; pub use oauth::StoredOAuthTokens; pub use oauth::WrappedOAuthTokenResponse; pub use oauth::delete_oauth_tokens; pub use oauth::save_oauth_tokens; +pub use oauth::stored_oauth_credential_snapshot; pub use oauth::stored_oauth_credentials; pub use perform_oauth_login::OAuthProviderError; pub use perform_oauth_login::OauthLoginHandle; diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index 1d36887fc5..913108b8fe 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -87,6 +87,38 @@ pub struct StoredOAuthTokens { pub expires_at: Option, } +/// OAuth credentials paired with the concrete store selected for their client lifecycle. +#[derive(Debug, Clone, PartialEq)] +pub struct StoredOAuthCredentialSnapshot { + credentials: StoredOAuthTokens, + store: ResolvedOAuthCredentialStore, +} + +impl StoredOAuthCredentialSnapshot { + /// Returns the normalized credentials originally read from the selected store. + pub fn credentials(&self) -> &StoredOAuthTokens { + &self.credentials + } + + /// Rereads the selected authority, allowing Auto login to migrate File credentials to keyring. + pub fn reload( + &self, + server_name: &str, + url: &str, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, + ) -> Result> { + if self.store == ResolvedOAuthCredentialStore::File + && store_mode == OAuthCredentialsStoreMode::Auto + { + return stored_oauth_credentials(server_name, url, store_mode, keyring_backend_kind); + } + + let credentials = self.store.load(&DefaultKeyringStore, server_name, url)?; + Ok(normalized_oauth_credentials(credentials.as_ref())) + } +} + /// Wrap OAuthTokenResponse to allow for partial equality comparison. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct WrappedOAuthTokenResponse(pub OAuthTokenResponse); @@ -134,6 +166,19 @@ pub fn stored_oauth_credentials( store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, ) -> Result> { + Ok( + stored_oauth_credential_snapshot(server_name, url, store_mode, keyring_backend_kind)? + .map(|snapshot| snapshot.credentials), + ) +} + +/// Loads OAuth credentials together with the concrete authority selected by store policy. +pub fn stored_oauth_credential_snapshot( + server_name: &str, + url: &str, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> Result> { let Some(resolved) = resolve_oauth_tokens_from_store_policy( &DefaultKeyringStore, server_name, @@ -144,7 +189,12 @@ pub fn stored_oauth_credentials( else { return Ok(None); }; - Ok(normalized_oauth_credentials(Some(&resolved.tokens))) + let mut credentials = resolved.tokens; + credentials.token_response.0.set_expires_in(None); + Ok(Some(StoredOAuthCredentialSnapshot { + credentials, + store: resolved.store, + })) } fn normalized_oauth_credentials(tokens: Option<&StoredOAuthTokens>) -> Option { diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 6ce310c0a2..7a6a7d5d98 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -251,7 +251,7 @@ where } #[derive(Debug, thiserror::Error)] -enum ClientOperationError { +pub(crate) enum ClientOperationError { #[error(transparent)] Service(#[from] rmcp::service::ServiceError), #[error("timed out awaiting {label} after {duration:.0?}")] diff --git a/codex-rs/rmcp-client/src/startup_error.rs b/codex-rs/rmcp-client/src/startup_error.rs index c74f637e3d..31dd3a6e70 100644 --- a/codex-rs/rmcp-client/src/startup_error.rs +++ b/codex-rs/rmcp-client/src/startup_error.rs @@ -1,9 +1,12 @@ use anyhow::Error; use rmcp::service::ClientInitializeError; +use rmcp::service::ServiceError; +use rmcp::transport::DynamicTransportError; use rmcp::transport::auth::AuthError; use rmcp::transport::streamable_http_client::StreamableHttpError; use crate::http_client_adapter::StreamableHttpClientAdapterError; +use crate::rmcp_client::ClientOperationError; /// Returns whether an RMCP client error indicates that authentication is required. /// @@ -17,24 +20,33 @@ pub fn is_authentication_required_error(error: &Error) -> bool { .is_some_and(auth_error_requires_authentication) || source .downcast_ref::() - .is_some_and(client_initialize_error_requires_authentication) + .is_some_and(|error| { + matches!( + error, + ClientInitializeError::TransportError { error, .. } + if transport_error_requires_authentication(error) + ) + }) + || source + .downcast_ref::() + .is_some_and(|error| { + matches!( + error, + ClientOperationError::Service(ServiceError::TransportSend(error)) + if transport_error_requires_authentication(error) + ) + }) }) } -fn client_initialize_error_requires_authentication(error: &ClientInitializeError) -> bool { - let ClientInitializeError::TransportError { error, .. } = error else { - return false; - }; - +fn transport_error_requires_authentication(error: &DynamicTransportError) -> bool { error .error .downcast_ref::>() - .is_some_and(|error| { - matches!( - error, - StreamableHttpError::Auth(auth_error) - if auth_error_requires_authentication(auth_error) - ) + .is_some_and(|error| match error { + StreamableHttpError::AuthRequired(_) => true, + StreamableHttpError::Auth(auth_error) => auth_error_requires_authentication(auth_error), + _ => false, }) } diff --git a/codex-rs/rmcp-client/src/streamable_http_retry_tests.rs b/codex-rs/rmcp-client/src/streamable_http_retry_tests.rs index 69b97e1a82..efda30dad6 100644 --- a/codex-rs/rmcp-client/src/streamable_http_retry_tests.rs +++ b/codex-rs/rmcp-client/src/streamable_http_retry_tests.rs @@ -3,9 +3,11 @@ use std::any::TypeId; use codex_exec_server::ExecServerError; use pretty_assertions::assert_eq; use rmcp::transport::DynamicTransportError; +use rmcp::transport::streamable_http_client::AuthRequiredError; use rmcp::transport::streamable_http_client::StreamableHttpError; use crate::http_client_adapter::StreamableHttpClientAdapterError; +use crate::rmcp_client::ClientOperationError; use super::*; @@ -59,6 +61,36 @@ fn retryable_streamable_http_error_includes_remote_body_stream_failure() { ); } +#[test] +fn startup_http_authentication_challenges_require_reauthorization() { + let transport_error = || { + DynamicTransportError::from_parts( + "streamable_http", + TypeId::of::<()>(), + Box::new( + StreamableHttpError::::AuthRequired( + AuthRequiredError::new("Bearer error=\"invalid_token\"".to_string()), + ), + ), + ) + }; + let errors = [ + anyhow::Error::new(rmcp::service::ClientInitializeError::TransportError { + error: transport_error(), + context: "send initialize request".into(), + }), + anyhow::Error::new(ClientOperationError::from( + rmcp::service::ServiceError::TransportSend(transport_error()), + )), + ]; + + for error in errors { + assert!(crate::startup_error::is_authentication_required_error( + &error + )); + } +} + fn retryable_initialize_error(context: &'static str) -> rmcp::service::ClientInitializeError { rmcp::service::ClientInitializeError::TransportError { error: DynamicTransportError::from_parts( diff --git a/codex-rs/rmcp-client/tests/streamable_http_oauth_store_pinning.rs b/codex-rs/rmcp-client/tests/streamable_http_oauth_store_pinning.rs index d92df03dd2..7b6c111ac3 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_store_pinning.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_store_pinning.rs @@ -19,6 +19,7 @@ use codex_rmcp_client::RmcpClient; use codex_rmcp_client::StoredOAuthTokens; use codex_rmcp_client::WrappedOAuthTokenResponse; use codex_rmcp_client::save_oauth_tokens; +use codex_rmcp_client::stored_oauth_credential_snapshot; use futures::future::BoxFuture; use keyring::credential::Credential; use keyring::credential::CredentialApi; @@ -209,13 +210,6 @@ async fn auto_store_remains_pinned_across_session_recovery_child() -> anyhow::Re let base_url = std::env::var(CHILD_SERVER_URL_ENV)?; let server_url = format!("{base_url}/mcp"); - let keyring_tokens = stored_tokens(&server_url, KEYRING_ACCESS_TOKEN); - save_oauth_tokens( - SERVER_NAME, - &keyring_tokens, - OAuthCredentialsStoreMode::Keyring, - AuthKeyringBackendKind::Direct, - )?; let file_tokens = stored_tokens(&server_url, FILE_ACCESS_TOKEN); save_oauth_tokens( SERVER_NAME, @@ -223,6 +217,35 @@ async fn auto_store_remains_pinned_across_session_recovery_child() -> anyhow::Re OAuthCredentialsStoreMode::File, AuthKeyringBackendKind::Direct, )?; + let file_snapshot = stored_oauth_credential_snapshot( + SERVER_NAME, + &server_url, + OAuthCredentialsStoreMode::Auto, + AuthKeyringBackendKind::Direct, + )? + .expect("Auto should initially resolve the fallback File"); + let keyring_tokens = stored_tokens(&server_url, KEYRING_ACCESS_TOKEN); + save_oauth_tokens( + SERVER_NAME, + &keyring_tokens, + OAuthCredentialsStoreMode::Auto, + AuthKeyringBackendKind::Direct, + )?; + save_oauth_tokens( + SERVER_NAME, + &file_tokens, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::Direct, + )?; + assert_eq!( + file_snapshot.reload( + SERVER_NAME, + &server_url, + OAuthCredentialsStoreMode::Auto, + AuthKeyringBackendKind::Direct, + )?, + Some(keyring_tokens.clone()), + ); let http_client = RecordingHttpClient::new(Environment::default_for_tests().get_http_client()); let client = RmcpClient::new_streamable_http_client(