diff --git a/codex-rs/rmcp-client/src/http_client_adapter.rs b/codex-rs/rmcp-client/src/http_client_adapter.rs index 6f98f78920..d089241d01 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 3c122c3bb7..4f29379c58 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 6696655d88..aa56c9e09a 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 e1ee18c753..53272fa1de 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 c348460795..224362603b 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 0000000000..fe024031ae --- /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 90b09d724c..792cf72209 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 2d87a2eac4..b6ad64fcec 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,59 +39,22 @@ 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; - 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) - .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()); @@ -102,7 +65,59 @@ async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn initialize_401_invalid_token_does_not_refresh_or_retry() -> anyhow::Result<()> { +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(ResponseTemplate::new(500)) + .expect(0) + .mount(&server) + .await; + + run_scenario_child(&server, SCENARIO_INITIALIZE_TIMES_OUT).await?; + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +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; - - 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:?}")), - } + .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(()) +} + +#[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(|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 03ee79392c..f46c419856 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 {