From 14c68356c790ccc32d357486d298dbd3ba39c717 Mon Sep 17 00:00:00 2001 From: Adam Perry Date: Fri, 5 Jun 2026 04:11:16 +0000 Subject: [PATCH 1/3] test(rmcp): reproduce premature OAuth 401 handling --- .../tests/streamable_http_oauth_startup.rs | 246 ++++++++++++++++-- 1 file changed, 218 insertions(+), 28 deletions(-) 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 1c18f2c98bda..2d87a2eac40a 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs @@ -1,6 +1,11 @@ mod streamable_http_test_support; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::time::Duration; +use std::time::SystemTime; +use std::time::UNIX_EPOCH; use codex_config::types::OAuthCredentialsStoreMode; use codex_exec_server::Environment; @@ -30,23 +35,16 @@ use streamable_http_test_support::initialize_client; const SERVER_NAME: &str = "test-streamable-http-oauth-startup"; const EXPIRED_ACCESS_TOKEN: &str = "expired-access-token"; +const REVOKED_ACCESS_TOKEN: &str = "revoked-access-token"; const REFRESH_TOKEN: &str = "valid-refresh-token"; const REFRESHED_ACCESS_TOKEN: &str = "refreshed-access-token"; const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_STARTUP_SERVER_URL"; +const INVALID_TOKEN_CHALLENGE: &str = r#"Bearer error="invalid_token""#; #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn refreshes_expired_persisted_token_before_initialize() -> 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": [""], - }))) - .expect(1) - .mount(&server) - .await; + mount_oauth_metadata(&server).await; Mock::given(method("POST")) .and(path("/oauth/token")) .and(body_string_contains("grant_type=refresh_token")) @@ -97,21 +95,113 @@ async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result let codex_home = TempDir::new()?; let server_url = format!("{}/mcp", server.uri()); - - // Credential storage resolves CODEX_HOME from the process environment. - // Run the client half of the test in an ignored helper test so it can use - // an isolated home without mutating the parent test runner's environment. - let status = Command::new(std::env::current_exe()?) - .args(["oauth_startup_child", "--exact", "--ignored", "--nocapture"]) - .env("CODEX_HOME", codex_home.path()) - .env(CHILD_SERVER_URL_ENV, server_url) - .status() - .await?; + let status = run_child_test(&codex_home, &server_url, "oauth_startup_child").await?; assert!(status.success(), "OAuth startup child failed: {status}"); server.verify().await; Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn initialize_401_invalid_token_does_not_refresh_or_retry() -> anyhow::Result<()> { + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + mount_unexpected_refresh(&server).await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {REVOKED_ACCESS_TOKEN}"), + )) + .respond_with( + ResponseTemplate::new(401).insert_header("www-authenticate", INVALID_TOKEN_CHALLENGE), + ) + .expect(1) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + let server_url = format!("{}/mcp", server.uri()); + let status = run_child_test( + &codex_home, + &server_url, + "initialize_401_invalid_token_child", + ) + .await?; + assert!( + status.success(), + "OAuth initialize 401 child failed: {status}" + ); + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn tool_call_401_invalid_token_does_not_refresh_or_retry() -> anyhow::Result<()> { + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + mount_unexpected_refresh(&server).await; + + let tool_call_count = Arc::new(AtomicUsize::new(0)); + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {REVOKED_ACCESS_TOKEN}"), + )) + .respond_with({ + let tool_call_count = Arc::clone(&tool_call_count); + move |request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body.get("method").and_then(Value::as_str) { + Some("initialize") => 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": { + "tools": {}, + }, + "serverInfo": { + "name": "oauth-reactive-401-test", + "version": "0.0.0-test", + }, + }, + })), + Some("notifications/initialized") => ResponseTemplate::new(202), + Some("tools/call") => { + tool_call_count.fetch_add(1, Ordering::SeqCst); + ResponseTemplate::new(401) + .insert_header("www-authenticate", INVALID_TOKEN_CHALLENGE) + } + 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 = run_child_test( + &codex_home, + &server_url, + "tool_call_401_invalid_token_child", + ) + .await?; + assert!( + status.success(), + "OAuth tool-call 401 child failed: {status}" + ); + assert_eq!(tool_call_count.load(Ordering::SeqCst), 1); + server.verify().await; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] #[ignore = "spawned by refreshes_expired_persisted_token_before_initialize"] async fn oauth_startup_child() -> anyhow::Result<()> { @@ -119,8 +209,101 @@ async fn oauth_startup_child() -> anyhow::Result<()> { // Save an expired access token with a valid refresh token so startup must // refresh before sending the initialize request. + save_test_oauth_tokens(&server_url, EXPIRED_ACCESS_TOKEN, /*expires_at*/ 0)?; + + let client = new_oauth_client(&server_url).await?; + initialize_client(&client).await?; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by initialize_401_invalid_token_does_not_refresh_or_retry"] +async fn initialize_401_invalid_token_child() -> anyhow::Result<()> { + let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; + save_test_oauth_tokens(&server_url, REVOKED_ACCESS_TOKEN, future_expiry())?; + + let client = new_oauth_client(&server_url).await?; + let error = initialize_client(&client).await.unwrap_err(); + let error_message = error.to_string(); + assert!( + error_message.contains("handshaking with MCP server failed: Send message error") + && error_message.contains("Auth required") + && error_message.contains("when send initialize request"), + "unexpected initialize error: {error:#}" + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by tool_call_401_invalid_token_does_not_refresh_or_retry"] +async fn tool_call_401_invalid_token_child() -> anyhow::Result<()> { + let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; + save_test_oauth_tokens(&server_url, REVOKED_ACCESS_TOKEN, future_expiry())?; + + let client = new_oauth_client(&server_url).await?; + initialize_client(&client).await?; + let error = client + .call_tool( + "echo".to_string(), + Some(json!({ "message": "should not be replayed" })), + /*meta*/ None, + Some(Duration::from_secs(5)), + ) + .await + .unwrap_err(); + assert!( + error.to_string().contains("Transport send error") + && error.to_string().contains("Auth required"), + "unexpected tool-call error: {error:#}" + ); + Ok(()) +} + +async fn mount_oauth_metadata(server: &MockServer) { + 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": [""], + }))) + .expect(1) + .mount(server) + .await; +} + +async fn mount_unexpected_refresh(server: &MockServer) { + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(500)) + .expect(0) + .mount(server) + .await; +} + +async fn run_child_test( + codex_home: &TempDir, + server_url: &str, + test_name: &str, +) -> anyhow::Result { + // Credential storage resolves CODEX_HOME from the process environment. + // Run the client half of the test in an ignored helper test so it can use + // an isolated home without mutating the parent test runner's environment. + Ok(Command::new(std::env::current_exe()?) + .args([test_name, "--exact", "--ignored", "--nocapture"]) + .env("CODEX_HOME", codex_home.path()) + .env(CHILD_SERVER_URL_ENV, server_url) + .status() + .await?) +} + +fn save_test_oauth_tokens( + server_url: &str, + access_token: &str, + expires_at: u64, +) -> anyhow::Result<()> { let mut response = OAuthTokenResponse::new( - AccessToken::new(EXPIRED_ACCESS_TOKEN.to_string()), + AccessToken::new(access_token.to_string()), BasicTokenType::Bearer, VendorExtraTokenFields::default(), ); @@ -128,19 +311,22 @@ async fn oauth_startup_child() -> anyhow::Result<()> { response.set_expires_in(Some(&Duration::from_secs(7200))); let tokens = StoredOAuthTokens { server_name: SERVER_NAME.to_string(), - url: server_url.clone(), + url: server_url.to_string(), client_id: "test-client-id".to_string(), token_response: WrappedOAuthTokenResponse(response), - expires_at: Some(0), + expires_at: Some(expires_at), }; save_oauth_tokens(SERVER_NAME, &tokens, OAuthCredentialsStoreMode::File)?; + Ok(()) +} +async fn new_oauth_client(server_url: &str) -> anyhow::Result { // This mirrors create_client's transport and initialization setup, except // it omits the direct bearer token. Supplying that token would bypass the // persisted OAuth credentials and the startup refresh under test. - let client = RmcpClient::new_streamable_http_client( + RmcpClient::new_streamable_http_client( SERVER_NAME, - &server_url, + server_url, /*bearer_token*/ None, /*http_headers*/ None, /*env_http_headers*/ None, @@ -148,8 +334,12 @@ async fn oauth_startup_child() -> anyhow::Result<()> { Environment::default_for_tests().get_http_client(), /*auth_provider*/ None, ) - .await?; + .await +} - initialize_client(&client).await?; - Ok(()) +fn future_expiry() -> u64 { + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default(); + now.saturating_add(Duration::from_secs(7200)).as_millis() as u64 } From 14fd2db95b6a0e9f3f89698e921bae0d9aa50111 Mon Sep 17 00:00:00 2001 From: Adam Perry Date: Fri, 5 Jun 2026 04:45:12 +0000 Subject: [PATCH 2/3] fix(rmcp): recover from premature OAuth invalid_token --- .../rmcp-client/src/http_client_adapter.rs | 10 + .../http_client_adapter/www_authenticate.rs | 25 +- .../www_authenticate_tests.rs | 22 + codex-rs/rmcp-client/src/lib.rs | 1 + codex-rs/rmcp-client/src/oauth.rs | 16 +- .../rmcp-client/src/oauth_invalid_token.rs | 33 ++ codex-rs/rmcp-client/src/rmcp_client.rs | 82 +++- .../tests/streamable_http_oauth_startup.rs | 413 ++++++++++++------ .../tests/streamable_http_test_support.rs | 9 +- 9 files changed, 449 insertions(+), 162 deletions(-) create mode 100644 codex-rs/rmcp-client/src/oauth_invalid_token.rs diff --git a/codex-rs/rmcp-client/src/http_client_adapter.rs b/codex-rs/rmcp-client/src/http_client_adapter.rs index 6f98f7892054..d089241d01ee 100644 --- a/codex-rs/rmcp-client/src/http_client_adapter.rs +++ b/codex-rs/rmcp-client/src/http_client_adapter.rs @@ -40,6 +40,7 @@ use sse_stream::SseStream; mod www_authenticate; +use self::www_authenticate::has_bearer_invalid_token_challenge; use self::www_authenticate::insufficient_scope_challenge; const EVENT_STREAM_MIME_TYPE: &str = "text/event-stream"; @@ -58,6 +59,8 @@ pub(crate) struct StreamableHttpClientAdapter { pub(crate) enum StreamableHttpClientAdapterError { #[error("streamable HTTP session expired with 404 Not Found")] SessionExpired404, + #[error("MCP OAuth access token was rejected")] + OAuthInvalidToken, #[error(transparent)] HttpRequest(#[from] ExecServerError), #[error("invalid HTTP header: {0}")] @@ -158,6 +161,13 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { StreamableHttpClientAdapterError::SessionExpired404, )); } + if response.status == StatusCode::UNAUTHORIZED.as_u16() + && has_bearer_invalid_token_challenge(&response.headers) + { + return Err(StreamableHttpError::Client( + StreamableHttpClientAdapterError::OAuthInvalidToken, + )); + } if response.status == StatusCode::UNAUTHORIZED.as_u16() && let Some(header) = response_header(&response.headers, reqwest::header::WWW_AUTHENTICATE) diff --git a/codex-rs/rmcp-client/src/http_client_adapter/www_authenticate.rs b/codex-rs/rmcp-client/src/http_client_adapter/www_authenticate.rs index 3c122c3bb786..4f29379c5872 100644 --- a/codex-rs/rmcp-client/src/http_client_adapter/www_authenticate.rs +++ b/codex-rs/rmcp-client/src/http_client_adapter/www_authenticate.rs @@ -60,6 +60,10 @@ impl BearerChallenge { Parameter::Missing | Parameter::Value(_) | Parameter::Invalid => None, } } + + fn has_error(&self, expected: &str) -> bool { + matches!(&self.error, Parameter::Value(error) if error == expected) + } } /// Finds a Bearer insufficient-scope challenge among all `WWW-Authenticate` @@ -80,6 +84,15 @@ pub(super) fn insufficient_scope_challenge( }) } +/// Returns whether any `WWW-Authenticate` response header field value contains +/// a Bearer challenge with an `invalid_token` error. +pub(super) fn has_bearer_invalid_token_challenge(headers: &[HttpHeader]) -> bool { + headers + .iter() + .filter(|header| header.name.eq_ignore_ascii_case(WWW_AUTHENTICATE.as_str())) + .any(|header| parse_bearer_challenge(&header.value, "invalid_token").is_some()) +} + /// Parses a Bearer `WWW-Authenticate` challenge with an `insufficient_scope` /// error and extracts its optional required scope. /// @@ -96,6 +109,11 @@ pub(super) fn insufficient_scope_challenge( /// /// RMCP has related parsing logic, but it is private to that crate. fn parse_bearer_insufficient_scope(header: &str) -> Option { + parse_bearer_challenge(header, "insufficient_scope") + .and_then(BearerChallenge::into_insufficient_scope) +} + +fn parse_bearer_challenge(header: &str, expected_error: &str) -> Option { let segments = split_unquoted_segments(header)?; let mut bearer_challenge: Option = None; @@ -107,9 +125,8 @@ fn parse_bearer_insufficient_scope(header: &str) -> Option Option Option> { diff --git a/codex-rs/rmcp-client/src/http_client_adapter/www_authenticate_tests.rs b/codex-rs/rmcp-client/src/http_client_adapter/www_authenticate_tests.rs index 6696655d8842..aa56c9e09acd 100644 --- a/codex-rs/rmcp-client/src/http_client_adapter/www_authenticate_tests.rs +++ b/codex-rs/rmcp-client/src/http_client_adapter/www_authenticate_tests.rs @@ -3,6 +3,7 @@ use pretty_assertions::assert_eq; use super::BearerInsufficientScope; use super::InsufficientScopeChallenge; +use super::has_bearer_invalid_token_challenge; use super::insufficient_scope_challenge; use super::parse_bearer_insufficient_scope; @@ -122,3 +123,24 @@ fn selects_bearer_challenge_from_a_later_www_authenticate_field_value() { }) ); } + +#[test] +fn detects_only_bearer_invalid_token_challenges() { + let matching = vec![HttpHeader { + name: "WWW-Authenticate".to_string(), + value: r#"Basic realm="example", Bearer error="invalid_token""#.to_string(), + }]; + assert!(has_bearer_invalid_token_challenge(&matching)); + + for value in [ + r#"Bearer realm="example""#, + r#"Bearer error="insufficient_scope""#, + r#"Basic error="invalid_token""#, + ] { + let headers = vec![HttpHeader { + name: "www-authenticate".to_string(), + value: value.to_string(), + }]; + assert!(!has_bearer_invalid_token_challenge(&headers), "{value}"); + } +} diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index e1ee18c75324..53272fa1ded6 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -5,6 +5,7 @@ mod http_client_adapter; mod in_process_transport; mod logging_client_handler; mod oauth; +mod oauth_invalid_token; mod perform_oauth_login; mod program_resolver; mod rmcp_client; diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index c348460795de..224362603bc6 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -345,10 +345,6 @@ impl OAuthPersistor { Ok(()) } - #[expect( - clippy::await_holding_invalid_type, - reason = "AuthorizationManager async access must be serialized through its mutex" - )] pub(crate) async fn refresh_if_needed(&self) -> Result<()> { let expires_at = { let guard = self.inner.last_credentials.lock().await; @@ -359,6 +355,18 @@ impl OAuthPersistor { return Ok(()); } + self.refresh().await + } + + pub(crate) async fn force_refresh(&self) -> Result<()> { + self.refresh().await + } + + #[expect( + clippy::await_holding_invalid_type, + reason = "AuthorizationManager async access must be serialized through its mutex" + )] + async fn refresh(&self) -> Result<()> { { let manager = self.inner.authorization_manager.clone(); let guard = manager.lock().await; diff --git a/codex-rs/rmcp-client/src/oauth_invalid_token.rs b/codex-rs/rmcp-client/src/oauth_invalid_token.rs new file mode 100644 index 000000000000..fe024031ae80 --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth_invalid_token.rs @@ -0,0 +1,33 @@ +use anyhow::Error; +use rmcp::service::ClientInitializeError; +use rmcp::transport::DynamicTransportError; +use rmcp::transport::streamable_http_client::StreamableHttpError; + +use crate::http_client_adapter::StreamableHttpClientAdapterError; + +#[derive(Debug, thiserror::Error)] +#[error("MCP OAuth access token was rejected; credentials refreshed, retry the request")] +pub(crate) struct RetryRequired; + +pub(crate) fn rejected_initialize_request(error: &Error) -> bool { + let Some(ClientInitializeError::TransportError { error, context }) = error + .chain() + .find_map(|source| source.downcast_ref::()) + else { + return false; + }; + + context.as_ref() == "send initialize request" && rejected_transport(error) +} + +pub(crate) fn rejected_transport(error: &DynamicTransportError) -> bool { + error + .error + .downcast_ref::>() + .is_some_and(|error| { + matches!( + error, + StreamableHttpError::Client(StreamableHttpClientAdapterError::OAuthInvalidToken) + ) + }) +} diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 90b09d724c3d..792cf72209c4 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -9,6 +9,7 @@ use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; +use anyhow::Context; use anyhow::Result; use anyhow::anyhow; use codex_api::SharedAuthProvider; @@ -41,6 +42,7 @@ use rmcp::model::ReadResourceResult; use rmcp::model::RequestId; use rmcp::model::ServerResult; use rmcp::model::Tool; +use rmcp::service::ClientInitializeError; use rmcp::service::RoleClient; use rmcp::service::RunningService; use rmcp::service::{self}; @@ -66,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_invalid_token; use crate::stdio_server_launcher::StdioServerCommand; use crate::stdio_server_launcher::StdioServerLauncher; use crate::stdio_server_launcher::StdioServerProcessHandle; @@ -223,6 +226,10 @@ enum ClientOperationError { Timeout { label: String, duration: Duration }, } +#[derive(Debug, thiserror::Error)] +#[error("handshaking with MCP server failed: {0}")] +struct McpHandshakeError(#[source] ClientInitializeError); + pub type Elicitation = CreateElicitationRequestParams; #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] @@ -396,9 +403,47 @@ impl RmcpClient { } }; - let (service, oauth_persistor) = - Self::connect_pending_transport(pending_transport, client_service.clone(), timeout) - .await?; + let startup_oauth_persistor = match &pending_transport { + PendingTransport::StreamableHttpWithOAuth { + oauth_persistor, .. + } => Some(oauth_persistor.clone()), + PendingTransport::InProcess { .. } + | PendingTransport::Stdio { .. } + | PendingTransport::StreamableHttp { .. } => None, + }; + let connect = async { + match Self::connect_pending_transport( + pending_transport, + client_service.clone(), + /*timeout*/ None, + ) + .await + { + Err(error) if oauth_invalid_token::rejected_initialize_request(&error) => { + let Some(oauth_persistor) = startup_oauth_persistor else { + return Err(error); + }; + oauth_persistor.force_refresh().await.with_context( + || "failed to refresh MCP OAuth access token after invalid_token", + )?; + let retry_transport = + Self::create_pending_transport(&self.transport_recipe).await?; + Self::connect_pending_transport( + retry_transport, + client_service.clone(), + /*timeout*/ None, + ) + .await + } + result => result, + } + }; + let (service, oauth_persistor) = match timeout { + Some(duration) => time::timeout(duration, connect).await.map_err(|_| { + anyhow!("timed out handshaking with MCP server after {duration:?}") + })??, + None => connect.await?, + }; let initialize_result_rmcp = service .peer() @@ -596,7 +641,20 @@ impl RmcpClient { } .boxed() }) - .await?; + .await; + let result = match result { + Ok(result) => result, + Err(error) if Self::is_invalid_token_operation_error(&error) => { + let Some(oauth_persistor) = self.oauth_persistor().await else { + return Err(error); + }; + oauth_persistor.force_refresh().await.with_context( + || "failed to refresh MCP OAuth access token after invalid_token", + )?; + return Err(oauth_invalid_token::RetryRequired.into()); + } + Err(error) => return Err(error), + }; self.persist_oauth_tokens().await; Ok(result) } @@ -849,10 +907,8 @@ impl RmcpClient { Some(duration) => time::timeout(duration, transport) .await .map_err(|_| anyhow!("timed out handshaking with MCP server after {duration:?}"))? - .map_err(|err| anyhow!("handshaking with MCP server failed: {err}"))?, - None => transport - .await - .map_err(|err| anyhow!("handshaking with MCP server failed: {err}"))?, + .map_err(McpHandshakeError)?, + None => transport.await.map_err(McpHandshakeError)?, }; Ok((Arc::new(service), oauth_persistor)) @@ -941,6 +997,16 @@ impl RmcpClient { }) } + fn is_invalid_token_operation_error(error: &anyhow::Error) -> bool { + let Some(ClientOperationError::Service(rmcp::service::ServiceError::TransportSend(error))) = + error.downcast_ref::() + else { + return false; + }; + + oauth_invalid_token::rejected_transport(error) + } + async fn reinitialize_after_session_expiry( &self, failed_service: &Arc>, 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 2d87a2eac40a..b6ad64fcecd8 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs @@ -1,8 +1,5 @@ mod streamable_http_test_support; -use std::sync::Arc; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; use std::time::Duration; use std::time::SystemTime; use std::time::UNIX_EPOCH; @@ -16,6 +13,7 @@ use codex_rmcp_client::save_oauth_tokens; use oauth2::AccessToken; use oauth2::RefreshToken; use oauth2::basic::BasicTokenType; +use pretty_assertions::assert_eq; use rmcp::transport::auth::OAuthTokenResponse; use rmcp::transport::auth::VendorExtraTokenFields; use serde_json::Value; @@ -31,7 +29,9 @@ use wiremock::matchers::header; use wiremock::matchers::method; use wiremock::matchers::path; +use streamable_http_test_support::expected_echo_result; use streamable_http_test_support::initialize_client; +use streamable_http_test_support::initialize_client_with_timeout; const SERVER_NAME: &str = "test-streamable-http-oauth-startup"; const EXPIRED_ACCESS_TOKEN: &str = "expired-access-token"; @@ -39,70 +39,85 @@ const REVOKED_ACCESS_TOKEN: &str = "revoked-access-token"; const REFRESH_TOKEN: &str = "valid-refresh-token"; const REFRESHED_ACCESS_TOKEN: &str = "refreshed-access-token"; const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_STARTUP_SERVER_URL"; +const CHILD_SCENARIO_ENV: &str = "MCP_TEST_OAUTH_STARTUP_SCENARIO"; const INVALID_TOKEN_CHALLENGE: &str = r#"Bearer error="invalid_token""#; +const RETRY_REQUIRED_ERROR: &str = + "MCP OAuth access token was rejected; credentials refreshed, retry the request"; + +const SCENARIO_INITIALIZE_SUCCEEDS: &str = "initialize_succeeds"; +const SCENARIO_INITIALIZE_FAILS: &str = "initialize_fails"; +const SCENARIO_INITIALIZE_TIMES_OUT: &str = "initialize_times_out"; +const SCENARIO_TOOL_CALL_RECOVERS: &str = "tool_call_recovers"; #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result<()> { let server = MockServer::start().await; mount_oauth_metadata(&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": REFRESH_TOKEN, - }))) - .expect(1) - .mount(&server) - .await; + mount_successful_refresh(&server).await; + mount_successful_initialize(&server, REFRESHED_ACCESS_TOKEN).await; + + let codex_home = TempDir::new()?; + let server_url = format!("{}/mcp", server.uri()); + let status = run_child_test(&codex_home, &server_url, "oauth_startup_child").await?; + assert!(status.success(), "OAuth startup child failed: {status}"); + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn initialize_invalid_token_refreshes_and_retries_once() -> anyhow::Result<()> { + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + mount_successful_refresh(&server).await; + mount_invalid_token_initialize(&server, REVOKED_ACCESS_TOKEN).await; + mount_successful_initialize(&server, REFRESHED_ACCESS_TOKEN).await; + + run_scenario_child(&server, SCENARIO_INITIALIZE_SUCCEEDS).await?; + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn initialize_invalid_token_retry_stops_after_second_401() -> anyhow::Result<()> { + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + mount_successful_refresh(&server).await; + mount_invalid_token_initialize(&server, REVOKED_ACCESS_TOKEN).await; + mount_invalid_token_initialize(&server, REFRESHED_ACCESS_TOKEN).await; + + run_scenario_child(&server, SCENARIO_INITIALIZE_FAILS).await?; + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn initialize_invalid_token_retry_uses_original_timeout() -> anyhow::Result<()> { + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + mount_invalid_token_initialize(&server, REVOKED_ACCESS_TOKEN).await; + mount_refresh( + &server, + successful_refresh_response().set_delay(Duration::from_secs(1)), + ) + .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") => 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": "oauth-startup-test", - "version": "0.0.0-test", - }, - }, - })), - Some("notifications/initialized") => ResponseTemplate::new(202), - method => ResponseTemplate::new(400) - .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), - } - }) - .expect(2) + .respond_with(ResponseTemplate::new(500)) + .expect(0) .mount(&server) .await; - let codex_home = TempDir::new()?; - let server_url = format!("{}/mcp", server.uri()); - let status = run_child_test(&codex_home, &server_url, "oauth_startup_child").await?; - assert!(status.success(), "OAuth startup child failed: {status}"); + run_scenario_child(&server, SCENARIO_INITIALIZE_TIMES_OUT).await?; server.verify().await; Ok(()) } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn initialize_401_invalid_token_does_not_refresh_or_retry() -> anyhow::Result<()> { +async fn initialize_generic_401_does_not_refresh_or_retry() -> anyhow::Result<()> { let server = MockServer::start().await; mount_oauth_metadata(&server).await; mount_unexpected_refresh(&server).await; @@ -113,91 +128,94 @@ async fn initialize_401_invalid_token_does_not_refresh_or_retry() -> anyhow::Res format!("Bearer {REVOKED_ACCESS_TOKEN}"), )) .respond_with( - ResponseTemplate::new(401).insert_header("www-authenticate", INVALID_TOKEN_CHALLENGE), + ResponseTemplate::new(401) + .insert_header("www-authenticate", r#"Bearer realm="example""#), ) .expect(1) .mount(&server) .await; - let codex_home = TempDir::new()?; - let server_url = format!("{}/mcp", server.uri()); - let status = run_child_test( - &codex_home, - &server_url, - "initialize_401_invalid_token_child", - ) - .await?; - assert!( - status.success(), - "OAuth initialize 401 child failed: {status}" - ); + run_scenario_child(&server, SCENARIO_INITIALIZE_FAILS).await?; server.verify().await; Ok(()) } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn tool_call_401_invalid_token_does_not_refresh_or_retry() -> anyhow::Result<()> { +async fn initialized_notification_invalid_token_does_not_refresh_or_retry() -> anyhow::Result<()> { let server = MockServer::start().await; mount_oauth_metadata(&server).await; mount_unexpected_refresh(&server).await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {REVOKED_ACCESS_TOKEN}"), + )) + .respond_with(|request: &Request| { + let body = request_body(request); + match body.get("method").and_then(Value::as_str) { + Some("initialize") => initialize_response(&body), + Some("notifications/initialized") => ResponseTemplate::new(401) + .insert_header("www-authenticate", INVALID_TOKEN_CHALLENGE), + method => unexpected_method(method), + } + }) + .expect(2) + .mount(&server) + .await; + + run_scenario_child(&server, SCENARIO_INITIALIZE_FAILS).await?; + server.verify().await; + Ok(()) +} - let tool_call_count = Arc::new(AtomicUsize::new(0)); +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn tool_call_invalid_token_refreshes_without_replay_and_next_call_succeeds() +-> anyhow::Result<()> { + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + mount_successful_refresh(&server).await; Mock::given(method("POST")) .and(path("/mcp")) .and(header( "authorization", format!("Bearer {REVOKED_ACCESS_TOKEN}"), )) - .respond_with({ - let tool_call_count = Arc::clone(&tool_call_count); - move |request: &Request| { - let body: Value = request.body_json().expect("valid JSON-RPC request"); - match body.get("method").and_then(Value::as_str) { - Some("initialize") => 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": { - "tools": {}, - }, - "serverInfo": { - "name": "oauth-reactive-401-test", - "version": "0.0.0-test", - }, - }, - })), - Some("notifications/initialized") => ResponseTemplate::new(202), - Some("tools/call") => { - tool_call_count.fetch_add(1, Ordering::SeqCst); - ResponseTemplate::new(401) - .insert_header("www-authenticate", INVALID_TOKEN_CHALLENGE) - } - method => ResponseTemplate::new(400) - .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), - } + .respond_with(|request: &Request| { + let body = request_body(request); + match body.get("method").and_then(Value::as_str) { + Some("initialize") => initialize_response(&body), + Some("notifications/initialized") => ResponseTemplate::new(202), + Some("tools/call") => ResponseTemplate::new(401) + .insert_header("www-authenticate", INVALID_TOKEN_CHALLENGE), + method => unexpected_method(method), } }) .expect(3) .mount(&server) .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {REFRESHED_ACCESS_TOKEN}"), + )) + .respond_with(|request: &Request| { + let body = request_body(request); + match body.get("method").and_then(Value::as_str) { + Some("tools/call") => ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": expected_echo_result("caller retry"), + })), + method => unexpected_method(method), + } + }) + .expect(1) + .mount(&server) + .await; - let codex_home = TempDir::new()?; - let server_url = format!("{}/mcp", server.uri()); - let status = run_child_test( - &codex_home, - &server_url, - "tool_call_401_invalid_token_child", - ) - .await?; - assert!( - status.success(), - "OAuth tool-call 401 child failed: {status}" - ); - assert_eq!(tool_call_count.load(Ordering::SeqCst), 1); + run_scenario_child(&server, SCENARIO_TOOL_CALL_RECOVERS).await?; server.verify().await; Ok(()) } @@ -206,9 +224,6 @@ async fn tool_call_401_invalid_token_does_not_refresh_or_retry() -> anyhow::Resu #[ignore = "spawned by refreshes_expired_persisted_token_before_initialize"] async fn oauth_startup_child() -> anyhow::Result<()> { let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; - - // Save an expired access token with a valid refresh token so startup must - // refresh before sending the initialize request. save_test_oauth_tokens(&server_url, EXPIRED_ACCESS_TOKEN, /*expires_at*/ 0)?; let client = new_oauth_client(&server_url).await?; @@ -217,44 +232,60 @@ async fn oauth_startup_child() -> anyhow::Result<()> { } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] -#[ignore = "spawned by initialize_401_invalid_token_does_not_refresh_or_retry"] -async fn initialize_401_invalid_token_child() -> anyhow::Result<()> { +#[ignore = "spawned by OAuth invalid_token parent tests"] +async fn oauth_invalid_token_child() -> anyhow::Result<()> { let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; + let scenario = std::env::var(CHILD_SCENARIO_ENV)?; save_test_oauth_tokens(&server_url, REVOKED_ACCESS_TOKEN, future_expiry())?; let client = new_oauth_client(&server_url).await?; - let error = initialize_client(&client).await.unwrap_err(); - let error_message = error.to_string(); - assert!( - error_message.contains("handshaking with MCP server failed: Send message error") - && error_message.contains("Auth required") - && error_message.contains("when send initialize request"), - "unexpected initialize error: {error:#}" - ); + match scenario.as_str() { + SCENARIO_INITIALIZE_SUCCEEDS => initialize_client(&client).await?, + SCENARIO_INITIALIZE_FAILS => { + initialize_client(&client).await.unwrap_err(); + } + SCENARIO_INITIALIZE_TIMES_OUT => { + let error = initialize_client_with_timeout(&client, Duration::from_millis(500)) + .await + .unwrap_err(); + assert_eq!( + error.to_string(), + "timed out handshaking with MCP server after 500ms" + ); + } + SCENARIO_TOOL_CALL_RECOVERS => { + initialize_client(&client).await?; + let error = call_echo_tool(&client, "must not be replayed") + .await + .unwrap_err(); + assert_eq!(error.to_string(), RETRY_REQUIRED_ERROR); + + let result = call_echo_tool(&client, "caller retry").await?; + assert_eq!(result, expected_echo_result("caller retry")); + } + other => anyhow::bail!("unknown OAuth invalid_token test scenario: {other}"), + } Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 1)] -#[ignore = "spawned by tool_call_401_invalid_token_does_not_refresh_or_retry"] -async fn tool_call_401_invalid_token_child() -> anyhow::Result<()> { - let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; - save_test_oauth_tokens(&server_url, REVOKED_ACCESS_TOKEN, future_expiry())?; - - let client = new_oauth_client(&server_url).await?; - initialize_client(&client).await?; - let error = client - .call_tool( - "echo".to_string(), - Some(json!({ "message": "should not be replayed" })), - /*meta*/ None, - Some(Duration::from_secs(5)), - ) - .await - .unwrap_err(); +async fn run_scenario_child(server: &MockServer, scenario: &str) -> anyhow::Result<()> { + let codex_home = TempDir::new()?; + let server_url = format!("{}/mcp", server.uri()); + let status = Command::new(std::env::current_exe()?) + .args([ + "oauth_invalid_token_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(CHILD_SERVER_URL_ENV, server_url) + .env(CHILD_SCENARIO_ENV, scenario) + .status() + .await?; assert!( - error.to_string().contains("Transport send error") - && error.to_string().contains("Auth required"), - "unexpected tool-call error: {error:#}" + status.success(), + "OAuth invalid_token child failed: {status}" ); Ok(()) } @@ -267,11 +298,37 @@ async fn mount_oauth_metadata(server: &MockServer) { "token_endpoint": format!("{}/oauth/token", server.uri()), "scopes_supported": [""], }))) + .expect(1..=2) + .mount(server) + .await; +} + +async fn mount_successful_refresh(server: &MockServer) { + mount_refresh(server, successful_refresh_response()).await; +} + +async fn mount_refresh(server: &MockServer, response: ResponseTemplate) { + 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(response) .expect(1) .mount(server) .await; } +fn successful_refresh_response() -> ResponseTemplate { + ResponseTemplate::new(200).set_body_json(json!({ + "access_token": REFRESHED_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": REFRESH_TOKEN, + })) +} + async fn mount_unexpected_refresh(server: &MockServer) { Mock::given(method("POST")) .and(path("/oauth/token")) @@ -281,14 +338,69 @@ async fn mount_unexpected_refresh(server: &MockServer) { .await; } +async fn mount_invalid_token_initialize(server: &MockServer, access_token: &str) { + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header("authorization", format!("Bearer {access_token}"))) + .respond_with( + ResponseTemplate::new(401).insert_header("www-authenticate", INVALID_TOKEN_CHALLENGE), + ) + .expect(1) + .mount(server) + .await; +} + +async fn mount_successful_initialize(server: &MockServer, access_token: &str) { + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header("authorization", format!("Bearer {access_token}"))) + .respond_with(|request: &Request| { + let body = request_body(request); + match body.get("method").and_then(Value::as_str) { + Some("initialize") => initialize_response(&body), + Some("notifications/initialized") => ResponseTemplate::new(202), + method => unexpected_method(method), + } + }) + .expect(2) + .mount(server) + .await; +} + +fn initialize_response(body: &Value) -> 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": { + "tools": {}, + }, + "serverInfo": { + "name": "oauth-invalid-token-test", + "version": "0.0.0-test", + }, + }, + })) +} + +fn unexpected_method(method: Option<&str>) -> ResponseTemplate { + ResponseTemplate::new(400).set_body_string(format!("unexpected JSON-RPC method: {method:?}")) +} + +#[expect(clippy::expect_used)] +fn request_body(request: &Request) -> Value { + request.body_json().expect("valid JSON-RPC request") +} + async fn run_child_test( codex_home: &TempDir, server_url: &str, test_name: &str, ) -> anyhow::Result { - // Credential storage resolves CODEX_HOME from the process environment. - // Run the client half of the test in an ignored helper test so it can use - // an isolated home without mutating the parent test runner's environment. Ok(Command::new(std::env::current_exe()?) .args([test_name, "--exact", "--ignored", "--nocapture"]) .env("CODEX_HOME", codex_home.path()) @@ -321,9 +433,6 @@ fn save_test_oauth_tokens( } async fn new_oauth_client(server_url: &str) -> anyhow::Result { - // This mirrors create_client's transport and initialization setup, except - // it omits the direct bearer token. Supplying that token would bypass the - // persisted OAuth credentials and the startup refresh under test. RmcpClient::new_streamable_http_client( SERVER_NAME, server_url, @@ -337,6 +446,20 @@ async fn new_oauth_client(server_url: &str) -> anyhow::Result { .await } +async fn call_echo_tool( + client: &RmcpClient, + message: &str, +) -> anyhow::Result { + client + .call_tool( + "echo".to_string(), + Some(json!({ "message": message })), + /*meta*/ None, + Some(Duration::from_secs(5)), + ) + .await +} + fn future_expiry() -> u64 { let now = SystemTime::now() .duration_since(UNIX_EPOCH) diff --git a/codex-rs/rmcp-client/tests/streamable_http_test_support.rs b/codex-rs/rmcp-client/tests/streamable_http_test_support.rs index 03ee79392cd9..f46c4198564d 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_test_support.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_test_support.rs @@ -92,10 +92,17 @@ pub(crate) async fn create_client(base_url: &str) -> anyhow::Result } pub(crate) async fn initialize_client(client: &RmcpClient) -> anyhow::Result<()> { + initialize_client_with_timeout(client, Duration::from_secs(5)).await +} + +pub(crate) async fn initialize_client_with_timeout( + client: &RmcpClient, + timeout: Duration, +) -> anyhow::Result<()> { client .initialize( init_params(), - Some(Duration::from_secs(5)), + Some(timeout), Box::new(|_, _| { async { Ok(ElicitationResponse { From 2cc6e6645342078445e69cd9be764248a4d39240 Mon Sep 17 00:00:00 2001 From: Adam Perry Date: Fri, 5 Jun 2026 05:22:59 +0000 Subject: [PATCH 3/3] fix(rmcp): scope invalid_token recovery to OAuth --- .../rmcp-client/src/http_client_adapter.rs | 11 +- .../http_client_adapter/www_authenticate.rs | 11 +- .../www_authenticate_tests.rs | 16 +- codex-rs/rmcp-client/src/lib.rs | 1 + codex-rs/rmcp-client/src/oauth.rs | 27 ++- .../rmcp-client/src/oauth_invalid_token.rs | 4 +- codex-rs/rmcp-client/src/rmcp_client.rs | 20 +- .../tests/streamable_http_oauth_startup.rs | 196 +++++++++++++++--- .../tests/streamable_http_recovery.rs | 26 +++ 9 files changed, 241 insertions(+), 71 deletions(-) diff --git a/codex-rs/rmcp-client/src/http_client_adapter.rs b/codex-rs/rmcp-client/src/http_client_adapter.rs index d089241d01ee..2894cdae33ea 100644 --- a/codex-rs/rmcp-client/src/http_client_adapter.rs +++ b/codex-rs/rmcp-client/src/http_client_adapter.rs @@ -40,8 +40,8 @@ use sse_stream::SseStream; mod www_authenticate; -use self::www_authenticate::has_bearer_invalid_token_challenge; use self::www_authenticate::insufficient_scope_challenge; +pub(crate) use self::www_authenticate::is_bearer_invalid_token_challenge; const EVENT_STREAM_MIME_TYPE: &str = "text/event-stream"; const JSON_MIME_TYPE: &str = "application/json"; @@ -59,8 +59,6 @@ pub(crate) struct StreamableHttpClientAdapter { pub(crate) enum StreamableHttpClientAdapterError { #[error("streamable HTTP session expired with 404 Not Found")] SessionExpired404, - #[error("MCP OAuth access token was rejected")] - OAuthInvalidToken, #[error(transparent)] HttpRequest(#[from] ExecServerError), #[error("invalid HTTP header: {0}")] @@ -161,13 +159,6 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { StreamableHttpClientAdapterError::SessionExpired404, )); } - if response.status == StatusCode::UNAUTHORIZED.as_u16() - && has_bearer_invalid_token_challenge(&response.headers) - { - return Err(StreamableHttpError::Client( - StreamableHttpClientAdapterError::OAuthInvalidToken, - )); - } if response.status == StatusCode::UNAUTHORIZED.as_u16() && let Some(header) = response_header(&response.headers, reqwest::header::WWW_AUTHENTICATE) diff --git a/codex-rs/rmcp-client/src/http_client_adapter/www_authenticate.rs b/codex-rs/rmcp-client/src/http_client_adapter/www_authenticate.rs index 4f29379c5872..eb23b4b35400 100644 --- a/codex-rs/rmcp-client/src/http_client_adapter/www_authenticate.rs +++ b/codex-rs/rmcp-client/src/http_client_adapter/www_authenticate.rs @@ -84,13 +84,10 @@ pub(super) fn insufficient_scope_challenge( }) } -/// Returns whether any `WWW-Authenticate` response header field value contains -/// a Bearer challenge with an `invalid_token` error. -pub(super) fn has_bearer_invalid_token_challenge(headers: &[HttpHeader]) -> bool { - headers - .iter() - .filter(|header| header.name.eq_ignore_ascii_case(WWW_AUTHENTICATE.as_str())) - .any(|header| parse_bearer_challenge(&header.value, "invalid_token").is_some()) +/// Returns whether a `WWW-Authenticate` field value contains a Bearer +/// challenge with an `invalid_token` error. +pub(crate) fn is_bearer_invalid_token_challenge(header: &str) -> bool { + parse_bearer_challenge(header, "invalid_token").is_some() } /// Parses a Bearer `WWW-Authenticate` challenge with an `insufficient_scope` diff --git a/codex-rs/rmcp-client/src/http_client_adapter/www_authenticate_tests.rs b/codex-rs/rmcp-client/src/http_client_adapter/www_authenticate_tests.rs index aa56c9e09acd..60a96a6b6167 100644 --- a/codex-rs/rmcp-client/src/http_client_adapter/www_authenticate_tests.rs +++ b/codex-rs/rmcp-client/src/http_client_adapter/www_authenticate_tests.rs @@ -3,8 +3,8 @@ use pretty_assertions::assert_eq; use super::BearerInsufficientScope; use super::InsufficientScopeChallenge; -use super::has_bearer_invalid_token_challenge; use super::insufficient_scope_challenge; +use super::is_bearer_invalid_token_challenge; use super::parse_bearer_insufficient_scope; #[test] @@ -126,21 +126,15 @@ fn selects_bearer_challenge_from_a_later_www_authenticate_field_value() { #[test] fn detects_only_bearer_invalid_token_challenges() { - let matching = vec![HttpHeader { - name: "WWW-Authenticate".to_string(), - value: r#"Basic realm="example", Bearer error="invalid_token""#.to_string(), - }]; - assert!(has_bearer_invalid_token_challenge(&matching)); + assert!(is_bearer_invalid_token_challenge( + r#"Basic realm="example", Bearer error="invalid_token""# + )); for value in [ r#"Bearer realm="example""#, r#"Bearer error="insufficient_scope""#, r#"Basic error="invalid_token""#, ] { - let headers = vec![HttpHeader { - name: "www-authenticate".to_string(), - value: value.to_string(), - }]; - assert!(!has_bearer_invalid_token_challenge(&headers), "{value}"); + assert!(!is_bearer_invalid_token_challenge(value), "{value}"); } } diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index 53272fa1ded6..55bd0681c7b9 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -18,6 +18,7 @@ pub use auth_status::discover_streamable_http_oauth; pub use auth_status::supports_oauth_login; pub use codex_protocol::protocol::McpAuthStatus; pub use in_process_transport::InProcessTransportFactory; +pub use oauth::OAUTH_REFRESH_REAUTHENTICATION_REQUIRED_ERROR; pub use oauth::StoredOAuthTokens; pub use oauth::WrappedOAuthTokenResponse; pub use oauth::delete_oauth_tokens; diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index 224362603bc6..a503ce3f342a 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -19,6 +19,7 @@ use anyhow::Context; use anyhow::Error; use anyhow::Result; +use anyhow::anyhow; use codex_config::types::OAuthCredentialsStoreMode; use oauth2::AccessToken; use oauth2::RefreshToken; @@ -45,6 +46,7 @@ use tracing::warn; use codex_keyring_store::DefaultKeyringStore; use codex_keyring_store::KeyringStore; +use rmcp::transport::auth::AuthError; use rmcp::transport::auth::AuthorizationManager; use tokio::sync::Mutex; @@ -52,6 +54,9 @@ use codex_utils_home_dir::find_codex_home; const KEYRING_SERVICE: &str = "Codex MCP Credentials"; const REFRESH_SKEW_MILLIS: u64 = 30_000; +const MISSING_REFRESH_TOKEN_ERROR: &str = "No refresh token available"; +pub const OAUTH_REFRESH_REAUTHENTICATION_REQUIRED_ERROR: &str = + "OAuth refresh token was rejected; reauthentication required"; #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct StoredOAuthTokens { @@ -370,18 +375,28 @@ impl OAuthPersistor { { let manager = self.inner.authorization_manager.clone(); let guard = manager.lock().await; - guard.refresh_token().await.with_context(|| { - format!( - "failed to refresh OAuth tokens for server {}", - self.inner.server_name - ) - })?; + if let Err(error) = guard.refresh_token().await { + return Err(refresh_error(error, &self.inner.server_name)); + } } self.persist_if_needed().await } } +fn refresh_error(error: AuthError, server_name: &str) -> Error { + match error { + AuthError::TokenRefreshFailed(message) + if message == MISSING_REFRESH_TOKEN_ERROR || message.contains("invalid_grant") => + { + anyhow!(OAUTH_REFRESH_REAUTHENTICATION_REQUIRED_ERROR) + } + error => Error::new(error).context(format!( + "failed to refresh OAuth tokens for server {server_name}" + )), + } +} + const FALLBACK_FILENAME: &str = ".credentials.json"; const MCP_SERVER_TYPE: &str = "http"; diff --git a/codex-rs/rmcp-client/src/oauth_invalid_token.rs b/codex-rs/rmcp-client/src/oauth_invalid_token.rs index fe024031ae80..6b6e694bb677 100644 --- a/codex-rs/rmcp-client/src/oauth_invalid_token.rs +++ b/codex-rs/rmcp-client/src/oauth_invalid_token.rs @@ -4,6 +4,7 @@ use rmcp::transport::DynamicTransportError; use rmcp::transport::streamable_http_client::StreamableHttpError; use crate::http_client_adapter::StreamableHttpClientAdapterError; +use crate::http_client_adapter::is_bearer_invalid_token_challenge; #[derive(Debug, thiserror::Error)] #[error("MCP OAuth access token was rejected; credentials refreshed, retry the request")] @@ -27,7 +28,8 @@ pub(crate) fn rejected_transport(error: &DynamicTransportError) -> bool { .is_some_and(|error| { matches!( error, - StreamableHttpError::Client(StreamableHttpClientAdapterError::OAuthInvalidToken) + StreamableHttpError::AuthRequired(auth) + if is_bearer_invalid_token_challenge(&auth.www_authenticate_header) ) }) } diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 792cf72209c4..ff6d5a9f6cea 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -9,7 +9,6 @@ use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; -use anyhow::Context; use anyhow::Result; use anyhow::anyhow; use codex_api::SharedAuthProvider; @@ -419,13 +418,14 @@ impl RmcpClient { ) .await { - Err(error) if oauth_invalid_token::rejected_initialize_request(&error) => { + Err(error) => { let Some(oauth_persistor) = startup_oauth_persistor else { return Err(error); }; - oauth_persistor.force_refresh().await.with_context( - || "failed to refresh MCP OAuth access token after invalid_token", - )?; + if !oauth_invalid_token::rejected_initialize_request(&error) { + return Err(error); + } + oauth_persistor.force_refresh().await?; let retry_transport = Self::create_pending_transport(&self.transport_recipe).await?; Self::connect_pending_transport( @@ -644,16 +644,16 @@ impl RmcpClient { .await; let result = match result { Ok(result) => result, - Err(error) if Self::is_invalid_token_operation_error(&error) => { + Err(error) => { let Some(oauth_persistor) = self.oauth_persistor().await else { return Err(error); }; - oauth_persistor.force_refresh().await.with_context( - || "failed to refresh MCP OAuth access token after invalid_token", - )?; + if !Self::is_invalid_token_operation_error(&error) { + return Err(error); + } + oauth_persistor.force_refresh().await?; return Err(oauth_invalid_token::RetryRequired.into()); } - Err(error) => return Err(error), }; self.persist_oauth_tokens().await; Ok(result) 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 b6ad64fcecd8..cb7f190ccc3d 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs @@ -1,11 +1,13 @@ mod streamable_http_test_support; +use std::fs; use std::time::Duration; use std::time::SystemTime; use std::time::UNIX_EPOCH; use codex_config::types::OAuthCredentialsStoreMode; use codex_exec_server::Environment; +use codex_rmcp_client::OAUTH_REFRESH_REAUTHENTICATION_REQUIRED_ERROR; use codex_rmcp_client::RmcpClient; use codex_rmcp_client::StoredOAuthTokens; use codex_rmcp_client::WrappedOAuthTokenResponse; @@ -38,6 +40,7 @@ const EXPIRED_ACCESS_TOKEN: &str = "expired-access-token"; const REVOKED_ACCESS_TOKEN: &str = "revoked-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 CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_STARTUP_SERVER_URL"; const CHILD_SCENARIO_ENV: &str = "MCP_TEST_OAUTH_STARTUP_SCENARIO"; const INVALID_TOKEN_CHALLENGE: &str = r#"Bearer error="invalid_token""#; @@ -46,7 +49,10 @@ const RETRY_REQUIRED_ERROR: &str = const SCENARIO_INITIALIZE_SUCCEEDS: &str = "initialize_succeeds"; const SCENARIO_INITIALIZE_FAILS: &str = "initialize_fails"; +const SCENARIO_INITIALIZE_REAUTH_REQUIRED: &str = "initialize_reauth_required"; const SCENARIO_INITIALIZE_TIMES_OUT: &str = "initialize_times_out"; +const SCENARIO_TOOL_CALL_FAILS: &str = "tool_call_fails"; +const SCENARIO_TOOL_CALL_REAUTH_REQUIRED: &str = "tool_call_reauth_required"; const SCENARIO_TOOL_CALL_RECOVERS: &str = "tool_call_recovers"; #[tokio::test(flavor = "multi_thread", worker_threads = 1)] @@ -97,7 +103,7 @@ async fn initialize_invalid_token_retry_uses_original_timeout() -> anyhow::Resul mount_invalid_token_initialize(&server, REVOKED_ACCESS_TOKEN).await; mount_refresh( &server, - successful_refresh_response().set_delay(Duration::from_secs(1)), + successful_refresh_response(REFRESH_TOKEN).set_delay(Duration::from_secs(1)), ) .await; Mock::given(method("POST")) @@ -141,30 +147,26 @@ async fn initialize_generic_401_does_not_refresh_or_retry() -> anyhow::Result<() } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn initialized_notification_invalid_token_does_not_refresh_or_retry() -> anyhow::Result<()> { +async fn initialize_invalid_token_rejected_refresh_requires_reauthentication() -> anyhow::Result<()> +{ let server = MockServer::start().await; mount_oauth_metadata(&server).await; - mount_unexpected_refresh(&server).await; + mount_invalid_token_initialize(&server, REVOKED_ACCESS_TOKEN).await; Mock::given(method("POST")) - .and(path("/mcp")) - .and(header( - "authorization", - format!("Bearer {REVOKED_ACCESS_TOKEN}"), - )) - .respond_with(|request: &Request| { - let body = request_body(request); - match body.get("method").and_then(Value::as_str) { - Some("initialize") => initialize_response(&body), - Some("notifications/initialized") => ResponseTemplate::new(401) - .insert_header("www-authenticate", INVALID_TOKEN_CHALLENGE), - method => unexpected_method(method), - } - }) - .expect(2) + .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(400).set_body_json(json!({ + "error": "invalid_grant", + "error_description": "refresh token revoked", + }))) + .expect(1) .mount(&server) .await; - run_scenario_child(&server, SCENARIO_INITIALIZE_FAILS).await?; + run_scenario_child(&server, SCENARIO_INITIALIZE_REAUTH_REQUIRED).await?; server.verify().await; Ok(()) } @@ -174,7 +176,7 @@ async fn tool_call_invalid_token_refreshes_without_replay_and_next_call_succeeds -> anyhow::Result<()> { let server = MockServer::start().await; mount_oauth_metadata(&server).await; - mount_successful_refresh(&server).await; + mount_successful_refresh_with_token(&server, ROTATED_REFRESH_TOKEN).await; Mock::given(method("POST")) .and(path("/mcp")) .and(header( @@ -215,7 +217,34 @@ async fn tool_call_invalid_token_refreshes_without_replay_and_next_call_succeeds .mount(&server) .await; - run_scenario_child(&server, SCENARIO_TOOL_CALL_RECOVERS).await?; + let codex_home = TempDir::new()?; + run_scenario_child_in_home(&server, SCENARIO_TOOL_CALL_RECOVERS, &codex_home).await?; + assert_persisted_refresh_token(&codex_home, ROTATED_REFRESH_TOKEN)?; + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn tool_call_invalid_token_without_refresh_token_requires_reauthentication() +-> anyhow::Result<()> { + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + mount_unexpected_refresh(&server).await; + mount_tool_call_failure(&server, INVALID_TOKEN_CHALLENGE).await; + + run_scenario_child(&server, SCENARIO_TOOL_CALL_REAUTH_REQUIRED).await?; + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn tool_call_generic_401_does_not_refresh_or_replay() -> anyhow::Result<()> { + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + mount_unexpected_refresh(&server).await; + mount_tool_call_failure(&server, r#"Bearer realm="example""#).await; + + run_scenario_child(&server, SCENARIO_TOOL_CALL_FAILS).await?; server.verify().await; Ok(()) } @@ -236,7 +265,15 @@ async fn oauth_startup_child() -> anyhow::Result<()> { async fn oauth_invalid_token_child() -> anyhow::Result<()> { let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; let scenario = std::env::var(CHILD_SCENARIO_ENV)?; - save_test_oauth_tokens(&server_url, REVOKED_ACCESS_TOKEN, future_expiry())?; + if scenario == SCENARIO_TOOL_CALL_REAUTH_REQUIRED { + save_test_oauth_tokens_without_refresh_token( + &server_url, + REVOKED_ACCESS_TOKEN, + future_expiry(), + )?; + } else { + save_test_oauth_tokens(&server_url, REVOKED_ACCESS_TOKEN, future_expiry())?; + } let client = new_oauth_client(&server_url).await?; match scenario.as_str() { @@ -244,6 +281,13 @@ async fn oauth_invalid_token_child() -> anyhow::Result<()> { SCENARIO_INITIALIZE_FAILS => { initialize_client(&client).await.unwrap_err(); } + SCENARIO_INITIALIZE_REAUTH_REQUIRED => { + let error = initialize_client(&client).await.unwrap_err(); + assert_eq!( + error.to_string(), + OAUTH_REFRESH_REAUTHENTICATION_REQUIRED_ERROR + ); + } SCENARIO_INITIALIZE_TIMES_OUT => { let error = initialize_client_with_timeout(&client, Duration::from_millis(500)) .await @@ -253,6 +297,22 @@ async fn oauth_invalid_token_child() -> anyhow::Result<()> { "timed out handshaking with MCP server after 500ms" ); } + SCENARIO_TOOL_CALL_FAILS => { + initialize_client(&client).await?; + call_echo_tool(&client, "must not be replayed") + .await + .unwrap_err(); + } + SCENARIO_TOOL_CALL_REAUTH_REQUIRED => { + initialize_client(&client).await?; + let error = call_echo_tool(&client, "must not be replayed") + .await + .unwrap_err(); + assert_eq!( + error.to_string(), + OAUTH_REFRESH_REAUTHENTICATION_REQUIRED_ERROR + ); + } SCENARIO_TOOL_CALL_RECOVERS => { initialize_client(&client).await?; let error = call_echo_tool(&client, "must not be replayed") @@ -270,6 +330,14 @@ async fn oauth_invalid_token_child() -> anyhow::Result<()> { async fn run_scenario_child(server: &MockServer, scenario: &str) -> anyhow::Result<()> { let codex_home = TempDir::new()?; + run_scenario_child_in_home(server, scenario, &codex_home).await +} + +async fn run_scenario_child_in_home( + server: &MockServer, + scenario: &str, + codex_home: &TempDir, +) -> anyhow::Result<()> { let server_url = format!("{}/mcp", server.uri()); let status = Command::new(std::env::current_exe()?) .args([ @@ -304,7 +372,11 @@ async fn mount_oauth_metadata(server: &MockServer) { } async fn mount_successful_refresh(server: &MockServer) { - mount_refresh(server, successful_refresh_response()).await; + mount_successful_refresh_with_token(server, REFRESH_TOKEN).await; +} + +async fn mount_successful_refresh_with_token(server: &MockServer, refresh_token: &str) { + mount_refresh(server, successful_refresh_response(refresh_token)).await; } async fn mount_refresh(server: &MockServer, response: ResponseTemplate) { @@ -320,12 +392,12 @@ async fn mount_refresh(server: &MockServer, response: ResponseTemplate) { .await; } -fn successful_refresh_response() -> ResponseTemplate { +fn successful_refresh_response(refresh_token: &str) -> ResponseTemplate { ResponseTemplate::new(200).set_body_json(json!({ "access_token": REFRESHED_ACCESS_TOKEN, "token_type": "Bearer", "expires_in": 7200, - "refresh_token": REFRESH_TOKEN, + "refresh_token": refresh_token, })) } @@ -350,6 +422,29 @@ async fn mount_invalid_token_initialize(server: &MockServer, access_token: &str) .await; } +async fn mount_tool_call_failure(server: &MockServer, challenge: &'static str) { + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {REVOKED_ACCESS_TOKEN}"), + )) + .respond_with(move |request: &Request| { + let body = request_body(request); + match body.get("method").and_then(Value::as_str) { + Some("initialize") => initialize_response(&body), + Some("notifications/initialized") => ResponseTemplate::new(202), + Some("tools/call") => { + ResponseTemplate::new(401).insert_header("www-authenticate", challenge) + } + method => unexpected_method(method), + } + }) + .expect(3) + .mount(server) + .await; +} + async fn mount_successful_initialize(server: &MockServer, access_token: &str) { Mock::given(method("POST")) .and(path("/mcp")) @@ -413,13 +508,40 @@ fn save_test_oauth_tokens( server_url: &str, access_token: &str, expires_at: u64, +) -> anyhow::Result<()> { + save_test_oauth_tokens_with_refresh_token( + server_url, + access_token, + expires_at, + /*refresh_token*/ Some(REFRESH_TOKEN), + ) +} + +fn save_test_oauth_tokens_without_refresh_token( + server_url: &str, + access_token: &str, + expires_at: u64, +) -> anyhow::Result<()> { + save_test_oauth_tokens_with_refresh_token( + server_url, + access_token, + expires_at, + /*refresh_token*/ None, + ) +} + +fn save_test_oauth_tokens_with_refresh_token( + server_url: &str, + access_token: &str, + expires_at: u64, + refresh_token: Option<&str>, ) -> anyhow::Result<()> { let mut response = OAuthTokenResponse::new( AccessToken::new(access_token.to_string()), BasicTokenType::Bearer, VendorExtraTokenFields::default(), ); - response.set_refresh_token(Some(RefreshToken::new(REFRESH_TOKEN.to_string()))); + response.set_refresh_token(refresh_token.map(|token| RefreshToken::new(token.to_string()))); response.set_expires_in(Some(&Duration::from_secs(7200))); let tokens = StoredOAuthTokens { server_name: SERVER_NAME.to_string(), @@ -432,6 +554,28 @@ fn save_test_oauth_tokens( Ok(()) } +fn assert_persisted_refresh_token( + codex_home: &TempDir, + expected_refresh_token: &str, +) -> anyhow::Result<()> { + let credentials: Value = serde_json::from_str(&fs::read_to_string( + codex_home.path().join(".credentials.json"), + )?)?; + let entries = credentials + .as_object() + .ok_or_else(|| anyhow::anyhow!("persisted OAuth credentials must be an object"))?; + assert_eq!(entries.len(), 1); + let entry = entries + .values() + .next() + .ok_or_else(|| anyhow::anyhow!("persisted OAuth credentials must contain one entry"))?; + assert_eq!( + entry.get("refresh_token"), + Some(&json!(expected_refresh_token)) + ); + Ok(()) +} + async fn new_oauth_client(server_url: &str) -> anyhow::Result { RmcpClient::new_streamable_http_client( SERVER_NAME, diff --git a/codex-rs/rmcp-client/tests/streamable_http_recovery.rs b/codex-rs/rmcp-client/tests/streamable_http_recovery.rs index 087d3d00df68..2801c97822cd 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_recovery.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_recovery.rs @@ -57,6 +57,32 @@ async fn streamable_http_401_does_not_trigger_recovery() -> anyhow::Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn static_bearer_invalid_token_retains_auth_required_semantics() -> anyhow::Result<()> { + let (_server, base_url) = spawn_streamable_http_server().await?; + let client = create_client(&base_url).await?; + + arm_session_post_failure( + &base_url, + /*status*/ 401, + /*remaining*/ 1, + /*www_authenticate_headers*/ &[r#"Bearer error="invalid_token""#], + ) + .await?; + + let error = call_echo_tool(&client, "unauthorized").await.unwrap_err(); + assert!( + error + .chain() + .any(|source| source.to_string().ends_with("Auth required")), + "static bearer invalid_token should remain AuthRequired: {error:#}" + ); + + let subsequent = call_echo_tool(&client, "subsequent").await?; + assert_eq!(subsequent, expected_echo_result("subsequent")); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn streamable_http_403_scope_challenge_returns_insufficient_scope() -> anyhow::Result<()> { let (_server, base_url) = spawn_streamable_http_server().await?;