diff --git a/codex-rs/rmcp-client/src/http_client_adapter.rs b/codex-rs/rmcp-client/src/http_client_adapter.rs index 6f98f7892054..2894cdae33ea 100644 --- a/codex-rs/rmcp-client/src/http_client_adapter.rs +++ b/codex-rs/rmcp-client/src/http_client_adapter.rs @@ -41,6 +41,7 @@ use sse_stream::SseStream; mod www_authenticate; 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"; 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..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 @@ -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,12 @@ pub(super) fn insufficient_scope_challenge( }) } +/// 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` /// error and extracts its optional required scope. /// @@ -96,6 +106,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 +122,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..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 @@ -4,6 +4,7 @@ use pretty_assertions::assert_eq; use super::BearerInsufficientScope; use super::InsufficientScopeChallenge; use super::insufficient_scope_challenge; +use super::is_bearer_invalid_token_challenge; use super::parse_bearer_insufficient_scope; #[test] @@ -122,3 +123,18 @@ fn selects_bearer_challenge_from_a_later_www_authenticate_field_value() { }) ); } + +#[test] +fn detects_only_bearer_invalid_token_challenges() { + 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""#, + ] { + 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 e1ee18c75324..55bd0681c7b9 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; @@ -17,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 c348460795de..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 { @@ -345,10 +350,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,21 +360,43 @@ 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; - 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 new file mode 100644 index 000000000000..6b6e694bb677 --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth_invalid_token.rs @@ -0,0 +1,35 @@ +use anyhow::Error; +use rmcp::service::ClientInitializeError; +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")] +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::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 90b09d724c3d..ff6d5a9f6cea 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -41,6 +41,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 +67,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 +225,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 +402,48 @@ 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) => { + let Some(oauth_persistor) = startup_oauth_persistor else { + return Err(error); + }; + 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( + 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) => { + let Some(oauth_persistor) = self.oauth_persistor().await else { + return Err(error); + }; + if !Self::is_invalid_token_operation_error(&error) { + return Err(error); + } + oauth_persistor.force_refresh().await?; + return Err(oauth_invalid_token::RetryRequired.into()); + } + }; 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 1c18f2c98bda..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,9 +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; @@ -11,6 +15,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; @@ -26,42 +31,171 @@ 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"; +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""#; +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_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)] 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": [""], - }))) + mount_oauth_metadata(&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(REFRESH_TOKEN).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; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {REVOKED_ACCESS_TOKEN}"), + )) + .respond_with( + ResponseTemplate::new(401) + .insert_header("www-authenticate", r#"Bearer realm="example""#), + ) .expect(1) .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 initialize_invalid_token_rejected_refresh_requires_reauthentication() -> anyhow::Result<()> +{ + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + mount_invalid_token_initialize(&server, REVOKED_ACCESS_TOKEN).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, + .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_REAUTH_REQUIRED).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_with_token(&server, ROTATED_REFRESH_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(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( @@ -69,45 +203,48 @@ async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result format!("Bearer {REFRESHED_ACCESS_TOKEN}"), )) .respond_with(|request: &Request| { - let body: Value = request.body_json().expect("valid JSON-RPC request"); + let body = request_body(request); match body.get("method").and_then(Value::as_str) { - Some("initialize") => ResponseTemplate::new(200).set_body_json(json!({ + Some("tools/call") => 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", - }, - }, + "result": expected_echo_result("caller retry"), })), - Some("notifications/initialized") => ResponseTemplate::new(202), - method => ResponseTemplate::new(400) - .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + method => unexpected_method(method), } }) - .expect(2) + .expect(1) .mount(&server) .await; let codex_home = TempDir::new()?; - let server_url = format!("{}/mcp", server.uri()); + 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(()) +} - // 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?; - assert!(status.success(), "OAuth startup child failed: {status}"); +#[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(()) } @@ -116,31 +253,333 @@ async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result #[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_test_oauth_tokens(&server_url, EXPIRED_ACCESS_TOKEN, /*expires_at*/ 0)?; + + let client = new_oauth_client(&server_url).await?; + initialize_client(&client).await?; + Ok(()) +} - // Save an expired access token with a valid refresh token so startup must - // refresh before sending the initialize request. +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[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)?; + 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() { + SCENARIO_INITIALIZE_SUCCEEDS => initialize_client(&client).await?, + 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 + .unwrap_err(); + assert_eq!( + error.to_string(), + "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") + .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(()) +} + +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([ + "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!( + status.success(), + "OAuth invalid_token child failed: {status}" + ); + 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..=2) + .mount(server) + .await; +} + +async fn mount_successful_refresh(server: &MockServer) { + 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) { + 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(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, + })) +} + +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 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_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")) + .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 { + 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<()> { + 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(EXPIRED_ACCESS_TOKEN.to_string()), + 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(), - 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(()) +} - // 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( +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, - &server_url, + server_url, /*bearer_token*/ None, /*http_headers*/ None, /*env_http_headers*/ None, @@ -148,8 +587,26 @@ async fn oauth_startup_child() -> anyhow::Result<()> { Environment::default_for_tests().get_http_client(), /*auth_provider*/ None, ) - .await?; + .await +} - initialize_client(&client).await?; - Ok(()) +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) + .unwrap_or_default(); + now.saturating_add(Duration::from_secs(7200)).as_millis() as u64 } 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?; 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 {