From 45c6f675965dd4b69b49bcae680824f25d7bd8f6 Mon Sep 17 00:00:00 2001 From: "Adam Perry @ OpenAI" Date: Thu, 4 Jun 2026 21:13:54 -0700 Subject: [PATCH 1/3] test(rmcp): cover short-lived OAuth token lifecycle --- .../tests/streamable_http_oauth_lifecycle.rs | 395 ++++++++++++++++++ 1 file changed, 395 insertions(+) create mode 100644 codex-rs/rmcp-client/tests/streamable_http_oauth_lifecycle.rs diff --git a/codex-rs/rmcp-client/tests/streamable_http_oauth_lifecycle.rs b/codex-rs/rmcp-client/tests/streamable_http_oauth_lifecycle.rs new file mode 100644 index 000000000000..6e5c281d89b2 --- /dev/null +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_lifecycle.rs @@ -0,0 +1,395 @@ +mod streamable_http_test_support; + +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use codex_config::types::OAuthCredentialsStoreMode; +use codex_exec_server::Environment; +use codex_rmcp_client::RmcpClient; +use codex_rmcp_client::StoredOAuthTokens; +use codex_rmcp_client::WrappedOAuthTokenResponse; +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; +use serde_json::json; +use tempfile::TempDir; +use tokio::process::Command; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::Request; +use wiremock::Respond; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; +use wiremock::matchers::path; + +use streamable_http_test_support::initialize_client; + +const SERVER_NAME: &str = "test-streamable-http-oauth-lifecycle"; +const INITIAL_ACCESS_TOKEN: &str = "initial-expired-access-token"; +const INITIAL_REFRESH_TOKEN: &str = "initial-refresh-token"; +const SHORT_LIVED_ACCESS_TOKEN: &str = "short-lived-access-token"; +const SHORT_LIVED_REFRESH_TOKEN: &str = "short-lived-refresh-token"; +const LONG_LIVED_ACCESS_TOKEN: &str = "long-lived-access-token"; +const LONG_LIVED_REFRESH_TOKEN: &str = "long-lived-refresh-token"; +const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_LIFECYCLE_SERVER_URL"; +const CHILD_CHECKPOINT_URL_ENV: &str = "MCP_TEST_OAUTH_LIFECYCLE_CHECKPOINT_URL"; +const CHILD_SCENARIO_ENV: &str = "MCP_TEST_OAUTH_LIFECYCLE_SCENARIO"; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Scenario { + RefreshSucceeds, + RefreshTokenExpired, +} + +impl Scenario { + fn as_env(self) -> &'static str { + match self { + Self::RefreshSucceeds => "refresh-succeeds", + Self::RefreshTokenExpired => "refresh-token-expired", + } + } + + fn from_env(value: &str) -> anyhow::Result { + match value { + "refresh-succeeds" => Ok(Self::RefreshSucceeds), + "refresh-token-expired" => Ok(Self::RefreshTokenExpired), + _ => anyhow::bail!("unknown OAuth lifecycle scenario: {value}"), + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +enum TimelineEvent { + Refresh(String), + Mcp { + method: String, + authorization: String, + }, + IdleStarted, + IdleFinished, +} + +#[derive(Clone)] +struct Timeline { + events: Arc>>, +} + +impl Timeline { + fn new() -> Self { + Self { + events: Arc::new(Mutex::new(Vec::new())), + } + } + + fn push(&self, event: TimelineEvent) { + self.events + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(event); + } + + fn snapshot(&self) -> Vec { + self.events + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() + } +} + +struct TokenResponder { + scenario: Scenario, + timeline: Timeline, + calls: AtomicUsize, +} + +impl Respond for TokenResponder { + fn respond(&self, request: &Request) -> ResponseTemplate { + let body = String::from_utf8_lossy(&request.body); + let refresh_token = if body.contains(INITIAL_REFRESH_TOKEN) { + INITIAL_REFRESH_TOKEN + } else if body.contains(SHORT_LIVED_REFRESH_TOKEN) { + SHORT_LIVED_REFRESH_TOKEN + } else { + panic!("unexpected refresh request body: {body}"); + }; + self.timeline + .push(TimelineEvent::Refresh(refresh_token.to_string())); + + match self.calls.fetch_add(1, Ordering::SeqCst) { + 0 => ResponseTemplate::new(200).set_body_json(json!({ + "access_token": SHORT_LIVED_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 32, + "refresh_token": SHORT_LIVED_REFRESH_TOKEN, + })), + 1 if self.scenario == Scenario::RefreshSucceeds => ResponseTemplate::new(200) + .set_body_json(json!({ + "access_token": LONG_LIVED_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": LONG_LIVED_REFRESH_TOKEN, + })), + 1 | 2 if self.scenario == Scenario::RefreshTokenExpired => ResponseTemplate::new(400) + .set_body_json(json!({ + "error": "invalid_grant", + "error_description": "refresh token expired", + })), + call => panic!("unexpected OAuth token request #{call}"), + } + } +} + +#[derive(Clone)] +struct McpResponder { + timeline: Timeline, +} + +impl Respond for McpResponder { + fn respond(&self, request: &Request) -> ResponseTemplate { + let body: Value = match request.body_json() { + Ok(body) => body, + Err(error) => panic!("invalid JSON-RPC request: {error}"), + }; + let method = match body.get("method").and_then(Value::as_str) { + Some(method) => method, + None => panic!("JSON-RPC request missing method: {body}"), + }; + let authorization = match request + .headers + .get("authorization") + .and_then(|value| value.to_str().ok()) + { + Some(authorization) => authorization.to_string(), + None => panic!("MCP request missing authorization header"), + }; + self.timeline.push(TimelineEvent::Mcp { + method: method.to_string(), + authorization, + }); + + match method { + "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-lifecycle-test", + "version": "0.0.0-test", + }, + }, + })), + "notifications/initialized" => ResponseTemplate::new(202), + "tools/list" => ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { + "tools": [], + }, + })), + _ => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method}")), + } + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn refreshes_only_when_the_next_operation_starts_after_idle() -> anyhow::Result<()> { + let timeline = run_scenario(Scenario::RefreshSucceeds).await?; + + assert_eq!( + timeline, + vec![ + TimelineEvent::Refresh(INITIAL_REFRESH_TOKEN.to_string()), + TimelineEvent::Mcp { + method: "initialize".to_string(), + authorization: format!("Bearer {SHORT_LIVED_ACCESS_TOKEN}"), + }, + TimelineEvent::Mcp { + method: "notifications/initialized".to_string(), + authorization: format!("Bearer {SHORT_LIVED_ACCESS_TOKEN}"), + }, + TimelineEvent::IdleStarted, + TimelineEvent::IdleFinished, + TimelineEvent::Refresh(SHORT_LIVED_REFRESH_TOKEN.to_string()), + TimelineEvent::Mcp { + method: "tools/list".to_string(), + authorization: format!("Bearer {LONG_LIVED_ACCESS_TOKEN}"), + }, + ] + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn expired_refresh_token_before_next_operation_requires_reauthorization() -> anyhow::Result<()> +{ + let timeline = run_scenario(Scenario::RefreshTokenExpired).await?; + + assert_eq!( + timeline, + vec![ + TimelineEvent::Refresh(INITIAL_REFRESH_TOKEN.to_string()), + TimelineEvent::Mcp { + method: "initialize".to_string(), + authorization: format!("Bearer {SHORT_LIVED_ACCESS_TOKEN}"), + }, + TimelineEvent::Mcp { + method: "notifications/initialized".to_string(), + authorization: format!("Bearer {SHORT_LIVED_ACCESS_TOKEN}"), + }, + TimelineEvent::IdleStarted, + TimelineEvent::IdleFinished, + TimelineEvent::Refresh(SHORT_LIVED_REFRESH_TOKEN.to_string()), + TimelineEvent::Refresh(SHORT_LIVED_REFRESH_TOKEN.to_string()), + ] + ); + Ok(()) +} + +async fn run_scenario(scenario: Scenario) -> anyhow::Result> { + let server = MockServer::start().await; + let timeline = Timeline::new(); + + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": [""], + }))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(TokenResponder { + scenario, + timeline: timeline.clone(), + calls: AtomicUsize::new(0), + }) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(McpResponder { + timeline: timeline.clone(), + }) + .mount(&server) + .await; + + for (path, event) in [ + ("/checkpoint/idle-started", TimelineEvent::IdleStarted), + ("/checkpoint/idle-finished", TimelineEvent::IdleFinished), + ] { + let timeline = timeline.clone(); + Mock::given(method("POST")) + .and(wiremock::matchers::path(path)) + .respond_with(move |_: &Request| { + timeline.push(event.clone()); + ResponseTemplate::new(204) + }) + .mount(&server) + .await; + } + + let codex_home = TempDir::new()?; + let status = Command::new(std::env::current_exe()?) + .args([ + "oauth_lifecycle_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(CHILD_SERVER_URL_ENV, format!("{}/mcp", server.uri())) + .env(CHILD_CHECKPOINT_URL_ENV, server.uri()) + .env(CHILD_SCENARIO_ENV, scenario.as_env()) + .status() + .await?; + assert!(status.success(), "OAuth lifecycle child failed: {status}"); + + Ok(timeline.snapshot()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by OAuth lifecycle parent tests"] +async fn oauth_lifecycle_child() -> anyhow::Result<()> { + let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; + let checkpoint_url = std::env::var(CHILD_CHECKPOINT_URL_ENV)?; + let scenario = Scenario::from_env(&std::env::var(CHILD_SCENARIO_ENV)?)?; + + let mut response = OAuthTokenResponse::new( + AccessToken::new(INITIAL_ACCESS_TOKEN.to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + response.set_refresh_token(Some(RefreshToken::new(INITIAL_REFRESH_TOKEN.to_string()))); + response.set_expires_in(Some(&Duration::from_secs(7200))); + save_oauth_tokens( + SERVER_NAME, + &StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: server_url.clone(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: Some(0), + }, + OAuthCredentialsStoreMode::File, + )?; + + let client = RmcpClient::new_streamable_http_client( + SERVER_NAME, + &server_url, + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await?; + initialize_client(&client).await?; + + post_checkpoint(&checkpoint_url, "idle-started").await?; + tokio::time::sleep(Duration::from_secs(3)).await; + post_checkpoint(&checkpoint_url, "idle-finished").await?; + + let result = client + .list_tools(/*params*/ None, Some(Duration::from_secs(5))) + .await; + match scenario { + Scenario::RefreshSucceeds => { + assert_eq!(result?.tools, Vec::new()); + } + Scenario::RefreshTokenExpired => { + let error = result.expect_err("expired refresh token should fail the operation"); + assert!( + error.to_string().contains("authorization required"), + "unexpected expired refresh token error: {error:#}" + ); + } + } + + Ok(()) +} + +async fn post_checkpoint(base_url: &str, checkpoint: &str) -> anyhow::Result<()> { + reqwest::Client::new() + .post(format!("{base_url}/checkpoint/{checkpoint}")) + .send() + .await? + .error_for_status()?; + Ok(()) +} From ed1ef2a09d51487a1570bdee523bbe90f59ec970 Mon Sep 17 00:00:00 2001 From: Adam Perry Date: Fri, 5 Jun 2026 04:48:28 +0000 Subject: [PATCH 2/3] fix(rmcp): propagate proactive OAuth refresh failures --- codex-rs/rmcp-client/src/rmcp_client.rs | 29 +- .../src/rmcp_client_oauth_tests.rs | 300 ++++++++++++++++++ .../tests/streamable_http_oauth_lifecycle.rs | 190 +++-------- 3 files changed, 367 insertions(+), 152 deletions(-) create mode 100644 codex-rs/rmcp-client/src/rmcp_client_oauth_tests.rs diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 90b09d724c3d..5399789f63e7 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -439,7 +439,7 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let result = self .run_service_operation("tools/list", timeout, move |service| { let params = params.clone(); @@ -455,7 +455,7 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let result = self .run_service_operation("tools/list", timeout, move |service| { let params = params.clone(); @@ -500,7 +500,7 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let result = self .run_service_operation("resources/list", timeout, move |service| { let params = params.clone(); @@ -516,7 +516,7 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let result = self .run_service_operation("resources/templates/list", timeout, move |service| { let params = params.clone(); @@ -532,7 +532,7 @@ impl RmcpClient { params: ReadResourceRequestParams, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let result = self .run_service_operation("resources/read", timeout, move |service| { let params = params.clone(); @@ -550,7 +550,7 @@ impl RmcpClient { meta: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let arguments = match arguments { Some(Value::Object(map)) => Some(map), Some(other) => { @@ -606,7 +606,7 @@ impl RmcpClient { method: &str, params: Option, ) -> Result<()> { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; self.run_service_operation( "notifications/custom", /*timeout*/ None, @@ -636,7 +636,7 @@ impl RmcpClient { method: &str, params: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let response = self .run_service_operation("requests/custom", /*timeout*/ None, move |service| { let params = params.clone(); @@ -700,12 +700,11 @@ impl RmcpClient { } } - async fn refresh_oauth_if_needed(&self) { - if let Some(runtime) = self.oauth_persistor().await - && let Err(error) = runtime.refresh_if_needed().await - { - warn!("failed to refresh OAuth tokens: {error}"); + async fn refresh_oauth_if_needed(&self) -> Result<()> { + if let Some(runtime) = self.oauth_persistor().await { + runtime.refresh_if_needed().await?; } + Ok(()) } async fn create_pending_transport( @@ -1061,6 +1060,10 @@ async fn create_oauth_transport_and_runtime( Ok((transport, runtime)) } +#[cfg(test)] +#[path = "rmcp_client_oauth_tests.rs"] +mod oauth_tests; + #[cfg(test)] mod tests { use std::time::Duration; diff --git a/codex-rs/rmcp-client/src/rmcp_client_oauth_tests.rs b/codex-rs/rmcp-client/src/rmcp_client_oauth_tests.rs new file mode 100644 index 000000000000..b2d87902048d --- /dev/null +++ b/codex-rs/rmcp-client/src/rmcp_client_oauth_tests.rs @@ -0,0 +1,300 @@ +use std::io; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use codex_config::types::OAuthCredentialsStoreMode; +use codex_exec_server::Environment; +use futures::FutureExt; +use oauth2::AccessToken; +use oauth2::RefreshToken; +use oauth2::basic::BasicTokenType; +use pretty_assertions::assert_eq; +use rmcp::ErrorData as McpError; +use rmcp::handler::server::ServerHandler; +use rmcp::model::ClientCapabilities; +use rmcp::model::Implementation; +use rmcp::model::ListToolsResult; +use rmcp::model::ProtocolVersion; +use rmcp::model::ServerCapabilities; +use rmcp::model::ServerInfo; +use rmcp::service::RequestContext; +use rmcp::service::RoleServer; +use rmcp::transport::auth::OAuthTokenResponse; +use rmcp::transport::auth::VendorExtraTokenFields; +use serde_json::json; +use tempfile::TempDir; +use tokio::io::DuplexStream; +use tokio::process::Command; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::Request; +use wiremock::Respond; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; +use wiremock::matchers::path; + +use super::*; +use crate::oauth::WrappedOAuthTokenResponse; + +const SERVER_NAME: &str = "request-time-refresh-test"; +const INITIAL_ACCESS_TOKEN: &str = "initial-access-token"; +const INITIAL_REFRESH_TOKEN: &str = "initial-refresh-token"; +const REFRESHED_ACCESS_TOKEN: &str = "refreshed-access-token"; +const REFRESHED_REFRESH_TOKEN: &str = "refreshed-refresh-token"; +const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_REQUEST_TIME_REFRESH_SERVER_URL"; +const CHILD_SCENARIO_ENV: &str = "MCP_TEST_REQUEST_TIME_REFRESH_SCENARIO"; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Scenario { + RefreshSucceeds, + ProviderRejectsRefresh, +} + +impl Scenario { + fn as_env(self) -> &'static str { + match self { + Self::RefreshSucceeds => "refresh-succeeds", + Self::ProviderRejectsRefresh => "provider-rejects-refresh", + } + } + + fn from_env(value: &str) -> anyhow::Result { + match value { + "refresh-succeeds" => Ok(Self::RefreshSucceeds), + "provider-rejects-refresh" => Ok(Self::ProviderRejectsRefresh), + _ => anyhow::bail!("unknown request-time refresh scenario: {value}"), + } + } +} + +#[derive(Clone, Copy)] +struct TokenResponder { + scenario: Scenario, +} + +impl Respond for TokenResponder { + fn respond(&self, request: &Request) -> ResponseTemplate { + let body = String::from_utf8_lossy(&request.body); + assert!( + body.contains(INITIAL_REFRESH_TOKEN), + "unexpected refresh request body: {body}" + ); + + match self.scenario { + Scenario::RefreshSucceeds => ResponseTemplate::new(200).set_body_json(json!({ + "access_token": REFRESHED_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": REFRESHED_REFRESH_TOKEN, + })), + Scenario::ProviderRejectsRefresh => ResponseTemplate::new(400).set_body_json(json!({ + "error": "invalid_grant", + "error_description": "provider rejected refresh", + })), + } + } +} + +#[derive(Clone)] +struct CountingServer { + list_tools_calls: Arc, +} + +impl ServerHandler for CountingServer { + fn get_info(&self) -> ServerInfo { + ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) + } + + fn list_tools( + &self, + _request: Option, + _context: RequestContext, + ) -> impl Future> + Send + '_ { + self.list_tools_calls.fetch_add(1, Ordering::SeqCst); + async { + Ok(ListToolsResult { + tools: Vec::new(), + next_cursor: None, + meta: None, + }) + } + } +} + +#[derive(Clone)] +struct TestTransportFactory { + server: CountingServer, +} + +impl InProcessTransportFactory for TestTransportFactory { + fn open(&self) -> BoxFuture<'static, io::Result> { + let server = self.server.clone(); + async move { + let (client_transport, server_transport) = tokio::io::duplex(64 * 1024); + let _server_task = tokio::spawn(async move { + if let Ok(service) = rmcp::serve_server(server, server_transport).await { + let _result = service.waiting().await; + } + }); + Ok(client_transport) + } + .boxed() + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn proactive_refresh_succeeds_before_request() -> anyhow::Result<()> { + run_parent_scenario(Scenario::RefreshSucceeds).await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn provider_rejection_of_proactive_refresh_stops_before_request() -> anyhow::Result<()> { + run_parent_scenario(Scenario::ProviderRejectsRefresh).await +} + +async fn run_parent_scenario(scenario: Scenario) -> 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; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(TokenResponder { scenario }) + .expect(1) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + let status = Command::new(std::env::current_exe()?) + .args([ + "rmcp_client::oauth_tests::oauth_request_time_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(CHILD_SERVER_URL_ENV, server.uri()) + .env(CHILD_SCENARIO_ENV, scenario.as_env()) + .status() + .await?; + assert!( + status.success(), + "request-time refresh child failed: {status}" + ); + + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by request-time refresh parent tests"] +async fn oauth_request_time_child() -> anyhow::Result<()> { + let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; + let scenario = Scenario::from_env(&std::env::var(CHILD_SCENARIO_ENV)?)?; + let list_tools_calls = Arc::new(AtomicUsize::new(0)); + let client = RmcpClient::new_in_process_client(Arc::new(TestTransportFactory { + server: CountingServer { + list_tools_calls: Arc::clone(&list_tools_calls), + }, + })) + .await?; + initialize_in_process_client(&client).await?; + + let mut response = OAuthTokenResponse::new( + AccessToken::new(INITIAL_ACCESS_TOKEN.to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + response.set_refresh_token(Some(RefreshToken::new(INITIAL_REFRESH_TOKEN.to_string()))); + response.set_expires_in(Some(&Duration::from_secs(7200))); + let mcp_url = format!("{server_url}/mcp"); + let initial_tokens = StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: mcp_url.clone(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: Some(0), + }; + let (_, runtime) = create_oauth_transport_and_runtime( + SERVER_NAME, + &mcp_url, + initial_tokens, + OAuthCredentialsStoreMode::File, + HeaderMap::new(), + Environment::default_for_tests().get_http_client(), + ) + .await?; + { + let mut state = client.state.lock().await; + match &mut *state { + ClientState::Ready { oauth, .. } => *oauth = Some(runtime), + ClientState::Connecting { .. } => panic!("client was not initialized"), + ClientState::Closed => panic!("client was unexpectedly closed"), + } + } + + let result = client + .list_tools(/*params*/ None, Some(Duration::from_secs(5))) + .await; + match scenario { + Scenario::RefreshSucceeds => { + assert_eq!(result?.tools, Vec::new()); + assert_eq!(list_tools_calls.load(Ordering::SeqCst), 1); + } + Scenario::ProviderRejectsRefresh => { + let error = match result { + Ok(_) => panic!("provider rejection should fail before tools/list"), + Err(error) => error, + }; + let error_chain = format!("{error:#}"); + assert!( + error_chain.contains( + "failed to refresh OAuth tokens for server request-time-refresh-test" + ), + "unexpected provider rejection error: {error_chain}" + ); + assert!( + error_chain.contains("invalid_grant"), + "provider error was not preserved: {error_chain}" + ); + assert_eq!(list_tools_calls.load(Ordering::SeqCst), 0); + } + } + + client.shutdown().await; + Ok(()) +} + +async fn initialize_in_process_client(client: &RmcpClient) -> anyhow::Result<()> { + let params = InitializeRequestParams::new( + ClientCapabilities::default(), + Implementation::new("codex-test", "0.0.0-test"), + ) + .with_protocol_version(ProtocolVersion::V_2025_06_18); + client + .initialize( + params, + Some(Duration::from_secs(5)), + Box::new(|_, _| { + async { + Ok(ElicitationResponse { + action: ElicitationAction::Accept, + content: Some(json!({})), + meta: None, + }) + } + .boxed() + }), + ) + .await?; + Ok(()) +} diff --git a/codex-rs/rmcp-client/tests/streamable_http_oauth_lifecycle.rs b/codex-rs/rmcp-client/tests/streamable_http_oauth_lifecycle.rs index 6e5c281d89b2..ee3c777c09d4 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_lifecycle.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_lifecycle.rs @@ -2,8 +2,6 @@ mod streamable_http_test_support; use std::sync::Arc; use std::sync::Mutex; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; use std::time::Duration; use codex_config::types::OAuthCredentialsStoreMode; @@ -35,36 +33,10 @@ use streamable_http_test_support::initialize_client; const SERVER_NAME: &str = "test-streamable-http-oauth-lifecycle"; const INITIAL_ACCESS_TOKEN: &str = "initial-expired-access-token"; const INITIAL_REFRESH_TOKEN: &str = "initial-refresh-token"; -const SHORT_LIVED_ACCESS_TOKEN: &str = "short-lived-access-token"; -const SHORT_LIVED_REFRESH_TOKEN: &str = "short-lived-refresh-token"; const LONG_LIVED_ACCESS_TOKEN: &str = "long-lived-access-token"; const LONG_LIVED_REFRESH_TOKEN: &str = "long-lived-refresh-token"; const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_LIFECYCLE_SERVER_URL"; const CHILD_CHECKPOINT_URL_ENV: &str = "MCP_TEST_OAUTH_LIFECYCLE_CHECKPOINT_URL"; -const CHILD_SCENARIO_ENV: &str = "MCP_TEST_OAUTH_LIFECYCLE_SCENARIO"; - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum Scenario { - RefreshSucceeds, - RefreshTokenExpired, -} - -impl Scenario { - fn as_env(self) -> &'static str { - match self { - Self::RefreshSucceeds => "refresh-succeeds", - Self::RefreshTokenExpired => "refresh-token-expired", - } - } - - fn from_env(value: &str) -> anyhow::Result { - match value { - "refresh-succeeds" => Ok(Self::RefreshSucceeds), - "refresh-token-expired" => Ok(Self::RefreshTokenExpired), - _ => anyhow::bail!("unknown OAuth lifecycle scenario: {value}"), - } - } -} #[derive(Clone, Debug, PartialEq, Eq)] enum TimelineEvent { @@ -104,46 +76,27 @@ impl Timeline { } } +#[derive(Clone)] struct TokenResponder { - scenario: Scenario, timeline: Timeline, - calls: AtomicUsize, } impl Respond for TokenResponder { fn respond(&self, request: &Request) -> ResponseTemplate { let body = String::from_utf8_lossy(&request.body); - let refresh_token = if body.contains(INITIAL_REFRESH_TOKEN) { - INITIAL_REFRESH_TOKEN - } else if body.contains(SHORT_LIVED_REFRESH_TOKEN) { - SHORT_LIVED_REFRESH_TOKEN - } else { - panic!("unexpected refresh request body: {body}"); - }; + assert!( + body.contains(INITIAL_REFRESH_TOKEN), + "unexpected refresh request body: {body}" + ); self.timeline - .push(TimelineEvent::Refresh(refresh_token.to_string())); - - match self.calls.fetch_add(1, Ordering::SeqCst) { - 0 => ResponseTemplate::new(200).set_body_json(json!({ - "access_token": SHORT_LIVED_ACCESS_TOKEN, - "token_type": "Bearer", - "expires_in": 32, - "refresh_token": SHORT_LIVED_REFRESH_TOKEN, - })), - 1 if self.scenario == Scenario::RefreshSucceeds => ResponseTemplate::new(200) - .set_body_json(json!({ - "access_token": LONG_LIVED_ACCESS_TOKEN, - "token_type": "Bearer", - "expires_in": 7200, - "refresh_token": LONG_LIVED_REFRESH_TOKEN, - })), - 1 | 2 if self.scenario == Scenario::RefreshTokenExpired => ResponseTemplate::new(400) - .set_body_json(json!({ - "error": "invalid_grant", - "error_description": "refresh token expired", - })), - call => panic!("unexpected OAuth token request #{call}"), - } + .push(TimelineEvent::Refresh(INITIAL_REFRESH_TOKEN.to_string())); + + ResponseTemplate::new(200).set_body_json(json!({ + "access_token": LONG_LIVED_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": LONG_LIVED_REFRESH_TOKEN, + })) } } @@ -206,60 +159,7 @@ impl Respond for McpResponder { } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn refreshes_only_when_the_next_operation_starts_after_idle() -> anyhow::Result<()> { - let timeline = run_scenario(Scenario::RefreshSucceeds).await?; - - assert_eq!( - timeline, - vec![ - TimelineEvent::Refresh(INITIAL_REFRESH_TOKEN.to_string()), - TimelineEvent::Mcp { - method: "initialize".to_string(), - authorization: format!("Bearer {SHORT_LIVED_ACCESS_TOKEN}"), - }, - TimelineEvent::Mcp { - method: "notifications/initialized".to_string(), - authorization: format!("Bearer {SHORT_LIVED_ACCESS_TOKEN}"), - }, - TimelineEvent::IdleStarted, - TimelineEvent::IdleFinished, - TimelineEvent::Refresh(SHORT_LIVED_REFRESH_TOKEN.to_string()), - TimelineEvent::Mcp { - method: "tools/list".to_string(), - authorization: format!("Bearer {LONG_LIVED_ACCESS_TOKEN}"), - }, - ] - ); - Ok(()) -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn expired_refresh_token_before_next_operation_requires_reauthorization() -> anyhow::Result<()> -{ - let timeline = run_scenario(Scenario::RefreshTokenExpired).await?; - - assert_eq!( - timeline, - vec![ - TimelineEvent::Refresh(INITIAL_REFRESH_TOKEN.to_string()), - TimelineEvent::Mcp { - method: "initialize".to_string(), - authorization: format!("Bearer {SHORT_LIVED_ACCESS_TOKEN}"), - }, - TimelineEvent::Mcp { - method: "notifications/initialized".to_string(), - authorization: format!("Bearer {SHORT_LIVED_ACCESS_TOKEN}"), - }, - TimelineEvent::IdleStarted, - TimelineEvent::IdleFinished, - TimelineEvent::Refresh(SHORT_LIVED_REFRESH_TOKEN.to_string()), - TimelineEvent::Refresh(SHORT_LIVED_REFRESH_TOKEN.to_string()), - ] - ); - Ok(()) -} - -async fn run_scenario(scenario: Scenario) -> anyhow::Result> { +async fn does_not_refresh_in_background_while_idle() -> anyhow::Result<()> { let server = MockServer::start().await; let timeline = Timeline::new(); @@ -270,15 +170,15 @@ async fn run_scenario(scenario: Scenario) -> anyhow::Result> "token_endpoint": format!("{}/oauth/token", server.uri()), "scopes_supported": [""], }))) + .expect(1) .mount(&server) .await; Mock::given(method("POST")) .and(path("/oauth/token")) .respond_with(TokenResponder { - scenario, timeline: timeline.clone(), - calls: AtomicUsize::new(0), }) + .expect(1) .mount(&server) .await; Mock::given(method("POST")) @@ -286,20 +186,22 @@ async fn run_scenario(scenario: Scenario) -> anyhow::Result> .respond_with(McpResponder { timeline: timeline.clone(), }) + .expect(3) .mount(&server) .await; - for (path, event) in [ + for (checkpoint_path, event) in [ ("/checkpoint/idle-started", TimelineEvent::IdleStarted), ("/checkpoint/idle-finished", TimelineEvent::IdleFinished), ] { let timeline = timeline.clone(); Mock::given(method("POST")) - .and(wiremock::matchers::path(path)) + .and(path(checkpoint_path)) .respond_with(move |_: &Request| { timeline.push(event.clone()); ResponseTemplate::new(204) }) + .expect(1) .mount(&server) .await; } @@ -315,20 +217,39 @@ async fn run_scenario(scenario: Scenario) -> anyhow::Result> .env("CODEX_HOME", codex_home.path()) .env(CHILD_SERVER_URL_ENV, format!("{}/mcp", server.uri())) .env(CHILD_CHECKPOINT_URL_ENV, server.uri()) - .env(CHILD_SCENARIO_ENV, scenario.as_env()) .status() .await?; assert!(status.success(), "OAuth lifecycle child failed: {status}"); - Ok(timeline.snapshot()) + assert_eq!( + timeline.snapshot(), + vec![ + TimelineEvent::Refresh(INITIAL_REFRESH_TOKEN.to_string()), + TimelineEvent::Mcp { + method: "initialize".to_string(), + authorization: format!("Bearer {LONG_LIVED_ACCESS_TOKEN}"), + }, + TimelineEvent::Mcp { + method: "notifications/initialized".to_string(), + authorization: format!("Bearer {LONG_LIVED_ACCESS_TOKEN}"), + }, + TimelineEvent::IdleStarted, + TimelineEvent::IdleFinished, + TimelineEvent::Mcp { + method: "tools/list".to_string(), + authorization: format!("Bearer {LONG_LIVED_ACCESS_TOKEN}"), + }, + ] + ); + server.verify().await; + Ok(()) } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] -#[ignore = "spawned by OAuth lifecycle parent tests"] +#[ignore = "spawned by OAuth lifecycle parent test"] async fn oauth_lifecycle_child() -> anyhow::Result<()> { let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; let checkpoint_url = std::env::var(CHILD_CHECKPOINT_URL_ENV)?; - let scenario = Scenario::from_env(&std::env::var(CHILD_SCENARIO_ENV)?)?; let mut response = OAuthTokenResponse::new( AccessToken::new(INITIAL_ACCESS_TOKEN.to_string()), @@ -363,25 +284,16 @@ async fn oauth_lifecycle_child() -> anyhow::Result<()> { initialize_client(&client).await?; post_checkpoint(&checkpoint_url, "idle-started").await?; - tokio::time::sleep(Duration::from_secs(3)).await; + tokio::time::sleep(Duration::from_millis(100)).await; post_checkpoint(&checkpoint_url, "idle-finished").await?; - let result = client - .list_tools(/*params*/ None, Some(Duration::from_secs(5))) - .await; - match scenario { - Scenario::RefreshSucceeds => { - assert_eq!(result?.tools, Vec::new()); - } - Scenario::RefreshTokenExpired => { - let error = result.expect_err("expired refresh token should fail the operation"); - assert!( - error.to_string().contains("authorization required"), - "unexpected expired refresh token error: {error:#}" - ); - } - } - + assert_eq!( + client + .list_tools(/*params*/ None, Some(Duration::from_secs(5))) + .await? + .tools, + Vec::new() + ); Ok(()) } From 58f096b30738cc0013412ed45fa3b77a1b787d5b Mon Sep 17 00:00:00 2001 From: Adam Perry Date: Fri, 5 Jun 2026 05:31:02 +0000 Subject: [PATCH 3/3] fix(rmcp): harden request-time OAuth refresh --- codex-rs/rmcp-client/src/oauth.rs | 22 +- .../src/rmcp_client_oauth_tests.rs | 309 +++++++++++++++--- .../tests/streamable_http_oauth_lifecycle.rs | 307 ----------------- 3 files changed, 275 insertions(+), 363 deletions(-) delete mode 100644 codex-rs/rmcp-client/tests/streamable_http_oauth_lifecycle.rs diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index c348460795de..3b21c536f746 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; @@ -52,6 +53,8 @@ use codex_utils_home_dir::find_codex_home; const KEYRING_SERVICE: &str = "Codex MCP Credentials"; const REFRESH_SKEW_MILLIS: u64 = 30_000; +pub(crate) const OAUTH_REQUEST_REFRESH_FAILED_ERROR: &str = + "OAuth access token refresh failed before MCP request"; #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct StoredOAuthTokens { @@ -362,15 +365,20 @@ 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 - ) - })?; + guard + .get_access_token() + .await + .map_err(|_| anyhow!(OAUTH_REQUEST_REFRESH_FAILED_ERROR))?; } - self.persist_if_needed().await + if let Err(error) = self.persist_if_needed().await { + warn!( + "failed to persist refreshed OAuth tokens for server {}; will retry: {error}", + self.inner.server_name + ); + } + + Ok(()) } } diff --git a/codex-rs/rmcp-client/src/rmcp_client_oauth_tests.rs b/codex-rs/rmcp-client/src/rmcp_client_oauth_tests.rs index b2d87902048d..b31d6be79e4d 100644 --- a/codex-rs/rmcp-client/src/rmcp_client_oauth_tests.rs +++ b/codex-rs/rmcp-client/src/rmcp_client_oauth_tests.rs @@ -1,8 +1,13 @@ +use std::fs; use std::io; +use std::path::PathBuf; use std::sync::Arc; +use std::sync::Mutex; 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; @@ -23,6 +28,7 @@ use rmcp::service::RequestContext; use rmcp::service::RoleServer; use rmcp::transport::auth::OAuthTokenResponse; use rmcp::transport::auth::VendorExtraTokenFields; +use serde_json::Value; use serde_json::json; use tempfile::TempDir; use tokio::io::DuplexStream; @@ -36,7 +42,9 @@ use wiremock::matchers::method; use wiremock::matchers::path; use super::*; +use crate::oauth::OAUTH_REQUEST_REFRESH_FAILED_ERROR; use crate::oauth::WrappedOAuthTokenResponse; +use crate::oauth::save_oauth_tokens; const SERVER_NAME: &str = "request-time-refresh-test"; const INITIAL_ACCESS_TOKEN: &str = "initial-access-token"; @@ -44,12 +52,17 @@ const INITIAL_REFRESH_TOKEN: &str = "initial-refresh-token"; const REFRESHED_ACCESS_TOKEN: &str = "refreshed-access-token"; const REFRESHED_REFRESH_TOKEN: &str = "refreshed-refresh-token"; const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_REQUEST_TIME_REFRESH_SERVER_URL"; +const CHILD_CHECKPOINT_URL_ENV: &str = "MCP_TEST_REQUEST_TIME_REFRESH_CHECKPOINT_URL"; const CHILD_SCENARIO_ENV: &str = "MCP_TEST_REQUEST_TIME_REFRESH_SCENARIO"; +const CREDENTIALS_FILE: &str = ".credentials.json"; #[derive(Clone, Copy, Debug, PartialEq, Eq)] enum Scenario { RefreshSucceeds, ProviderRejectsRefresh, + PersistenceFailsThenRetries, + ConcurrentRefreshes, + IdleThenRefresh, } impl Scenario { @@ -57,6 +70,9 @@ impl Scenario { match self { Self::RefreshSucceeds => "refresh-succeeds", Self::ProviderRejectsRefresh => "provider-rejects-refresh", + Self::PersistenceFailsThenRetries => "persistence-fails-then-retries", + Self::ConcurrentRefreshes => "concurrent-refreshes", + Self::IdleThenRefresh => "idle-then-refresh", } } @@ -64,14 +80,52 @@ impl Scenario { match value { "refresh-succeeds" => Ok(Self::RefreshSucceeds), "provider-rejects-refresh" => Ok(Self::ProviderRejectsRefresh), + "persistence-fails-then-retries" => Ok(Self::PersistenceFailsThenRetries), + "concurrent-refreshes" => Ok(Self::ConcurrentRefreshes), + "idle-then-refresh" => Ok(Self::IdleThenRefresh), _ => anyhow::bail!("unknown request-time refresh scenario: {value}"), } } } -#[derive(Clone, Copy)] +#[derive(Clone, Debug, PartialEq, Eq)] +enum TimelineEvent { + IdleStarted, + IdleFinished, + Refresh, +} + +#[derive(Clone)] +struct Timeline { + events: Arc>>, +} + +impl Timeline { + fn new() -> Self { + Self { + events: Arc::new(Mutex::new(Vec::new())), + } + } + + fn push(&self, event: TimelineEvent) { + self.events + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(event); + } + + fn snapshot(&self) -> Vec { + self.events + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() + } +} + +#[derive(Clone)] struct TokenResponder { scenario: Scenario, + timeline: Timeline, } impl Respond for TokenResponder { @@ -81,22 +135,32 @@ impl Respond for TokenResponder { body.contains(INITIAL_REFRESH_TOKEN), "unexpected refresh request body: {body}" ); + self.timeline.push(TimelineEvent::Refresh); match self.scenario { - Scenario::RefreshSucceeds => ResponseTemplate::new(200).set_body_json(json!({ - "access_token": REFRESHED_ACCESS_TOKEN, - "token_type": "Bearer", - "expires_in": 7200, - "refresh_token": REFRESHED_REFRESH_TOKEN, - })), Scenario::ProviderRejectsRefresh => ResponseTemplate::new(400).set_body_json(json!({ "error": "invalid_grant", - "error_description": "provider rejected refresh", + "error_description": "x".repeat(20_000), })), + Scenario::ConcurrentRefreshes => { + refreshed_token_response().set_delay(Duration::from_millis(100)) + } + Scenario::RefreshSucceeds + | Scenario::PersistenceFailsThenRetries + | Scenario::IdleThenRefresh => refreshed_token_response(), } } } +fn refreshed_token_response() -> ResponseTemplate { + ResponseTemplate::new(200).set_body_json(json!({ + "access_token": REFRESHED_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": REFRESHED_REFRESH_TOKEN, + })) +} + #[derive(Clone)] struct CountingServer { list_tools_calls: Arc, @@ -145,17 +209,34 @@ impl InProcessTransportFactory for TestTransportFactory { } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn proactive_refresh_succeeds_before_request() -> anyhow::Result<()> { +async fn proactive_refresh_succeeds_and_persists_before_request() -> anyhow::Result<()> { run_parent_scenario(Scenario::RefreshSucceeds).await } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn provider_rejection_of_proactive_refresh_stops_before_request() -> anyhow::Result<()> { +async fn provider_rejection_is_bounded_and_stops_before_request() -> anyhow::Result<()> { run_parent_scenario(Scenario::ProviderRejectsRefresh).await } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn persistence_failure_does_not_abort_and_retries_without_refresh() -> anyhow::Result<()> { + run_parent_scenario(Scenario::PersistenceFailsThenRetries).await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn concurrent_near_expiry_requests_share_one_refresh() -> anyhow::Result<()> { + run_parent_scenario(Scenario::ConcurrentRefreshes).await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn near_expiry_credentials_do_not_refresh_while_idle_then_refresh_on_request() +-> anyhow::Result<()> { + run_parent_scenario(Scenario::IdleThenRefresh).await +} + async fn run_parent_scenario(scenario: Scenario) -> anyhow::Result<()> { let server = MockServer::start().await; + let timeline = Timeline::new(); Mock::given(method("GET")) .and(path("/.well-known/oauth-authorization-server/mcp")) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ @@ -168,11 +249,42 @@ async fn run_parent_scenario(scenario: Scenario) -> anyhow::Result<()> { .await; Mock::given(method("POST")) .and(path("/oauth/token")) - .respond_with(TokenResponder { scenario }) + .respond_with(TokenResponder { + scenario, + timeline: timeline.clone(), + }) .expect(1) .mount(&server) .await; + if scenario == Scenario::IdleThenRefresh { + let idle_started_timeline = timeline.clone(); + Mock::given(method("POST")) + .and(path("/checkpoint/idle-started")) + .respond_with(move |_: &Request| { + idle_started_timeline.push(TimelineEvent::IdleStarted); + ResponseTemplate::new(204) + }) + .expect(1) + .mount(&server) + .await; + + let idle_finished_timeline = timeline.clone(); + Mock::given(method("POST")) + .and(path("/checkpoint/idle-finished")) + .respond_with(move |_: &Request| { + idle_finished_timeline.push(TimelineEvent::IdleFinished); + assert_eq!( + idle_finished_timeline.snapshot(), + vec![TimelineEvent::IdleStarted, TimelineEvent::IdleFinished] + ); + ResponseTemplate::new(204) + }) + .expect(1) + .mount(&server) + .await; + } + let codex_home = TempDir::new()?; let status = Command::new(std::env::current_exe()?) .args([ @@ -183,6 +295,7 @@ async fn run_parent_scenario(scenario: Scenario) -> anyhow::Result<()> { ]) .env("CODEX_HOME", codex_home.path()) .env(CHILD_SERVER_URL_ENV, server.uri()) + .env(CHILD_CHECKPOINT_URL_ENV, server.uri()) .env(CHILD_SCENARIO_ENV, scenario.as_env()) .status() .await?; @@ -191,6 +304,16 @@ async fn run_parent_scenario(scenario: Scenario) -> anyhow::Result<()> { "request-time refresh child failed: {status}" ); + if scenario == Scenario::IdleThenRefresh { + assert_eq!( + timeline.snapshot(), + vec![ + TimelineEvent::IdleStarted, + TimelineEvent::IdleFinished, + TimelineEvent::Refresh, + ] + ); + } server.verify().await; Ok(()) } @@ -199,40 +322,36 @@ async fn run_parent_scenario(scenario: Scenario) -> anyhow::Result<()> { #[ignore = "spawned by request-time refresh parent tests"] async fn oauth_request_time_child() -> anyhow::Result<()> { let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; + let checkpoint_url = std::env::var(CHILD_CHECKPOINT_URL_ENV)?; let scenario = Scenario::from_env(&std::env::var(CHILD_SCENARIO_ENV)?)?; let list_tools_calls = Arc::new(AtomicUsize::new(0)); - let client = RmcpClient::new_in_process_client(Arc::new(TestTransportFactory { - server: CountingServer { - list_tools_calls: Arc::clone(&list_tools_calls), - }, - })) - .await?; + let client = Arc::new( + RmcpClient::new_in_process_client(Arc::new(TestTransportFactory { + server: CountingServer { + list_tools_calls: Arc::clone(&list_tools_calls), + }, + })) + .await?, + ); initialize_in_process_client(&client).await?; - let mut response = OAuthTokenResponse::new( - AccessToken::new(INITIAL_ACCESS_TOKEN.to_string()), - BasicTokenType::Bearer, - VendorExtraTokenFields::default(), - ); - response.set_refresh_token(Some(RefreshToken::new(INITIAL_REFRESH_TOKEN.to_string()))); - response.set_expires_in(Some(&Duration::from_secs(7200))); - let mcp_url = format!("{server_url}/mcp"); - let initial_tokens = StoredOAuthTokens { - server_name: SERVER_NAME.to_string(), - url: mcp_url.clone(), - client_id: "test-client-id".to_string(), - token_response: WrappedOAuthTokenResponse(response), - expires_at: Some(0), - }; + let initial_tokens = near_expiry_initial_tokens(&server_url)?; + save_oauth_tokens( + SERVER_NAME, + &initial_tokens, + OAuthCredentialsStoreMode::File, + )?; + let initial_persisted = fs::read(credentials_path()?)?; let (_, runtime) = create_oauth_transport_and_runtime( SERVER_NAME, - &mcp_url, - initial_tokens, + &initial_tokens.url, + initial_tokens.clone(), OAuthCredentialsStoreMode::File, HeaderMap::new(), Environment::default_for_tests().get_http_client(), ) .await?; + let runtime_for_assertions = runtime.clone(); { let mut state = client.state.lock().await; match &mut *state { @@ -242,31 +361,57 @@ async fn oauth_request_time_child() -> anyhow::Result<()> { } } - let result = client - .list_tools(/*params*/ None, Some(Duration::from_secs(5))) - .await; match scenario { Scenario::RefreshSucceeds => { - assert_eq!(result?.tools, Vec::new()); + assert_eq!(list_tools(&client).await?.tools, Vec::new()); assert_eq!(list_tools_calls.load(Ordering::SeqCst), 1); + assert_refreshed_credentials_persisted()?; } Scenario::ProviderRejectsRefresh => { - let error = match result { + let error = match list_tools(&client).await { Ok(_) => panic!("provider rejection should fail before tools/list"), Err(error) => error, }; - let error_chain = format!("{error:#}"); - assert!( - error_chain.contains( - "failed to refresh OAuth tokens for server request-time-refresh-test" - ), - "unexpected provider rejection error: {error_chain}" - ); - assert!( - error_chain.contains("invalid_grant"), - "provider error was not preserved: {error_chain}" + assert_eq!(error.to_string(), OAUTH_REQUEST_REFRESH_FAILED_ERROR); + assert_eq!( + format!("{error:#}").chars().count(), + OAUTH_REQUEST_REFRESH_FAILED_ERROR.chars().count() ); assert_eq!(list_tools_calls.load(Ordering::SeqCst), 0); + runtime_for_assertions.persist_if_needed().await?; + assert_eq!(fs::read(credentials_path()?)?, initial_persisted); + } + Scenario::PersistenceFailsThenRetries => { + let credentials_path = credentials_path()?; + let original_permissions = fs::metadata(&credentials_path)?.permissions(); + let mut readonly_permissions = original_permissions.clone(); + readonly_permissions.set_readonly(true); + fs::set_permissions(&credentials_path, readonly_permissions)?; + + assert_eq!(list_tools(&client).await?.tools, Vec::new()); + assert_eq!(list_tools_calls.load(Ordering::SeqCst), 1); + assert_eq!(fs::read(&credentials_path)?, initial_persisted); + + fs::set_permissions(&credentials_path, original_permissions)?; + assert_eq!(list_tools(&client).await?.tools, Vec::new()); + assert_eq!(list_tools_calls.load(Ordering::SeqCst), 2); + assert_refreshed_credentials_persisted()?; + } + Scenario::ConcurrentRefreshes => { + let (first, second) = tokio::join!(list_tools(&client), list_tools(&client)); + assert_eq!(first?.tools, Vec::new()); + assert_eq!(second?.tools, Vec::new()); + assert_eq!(list_tools_calls.load(Ordering::SeqCst), 2); + assert_refreshed_credentials_persisted()?; + } + Scenario::IdleThenRefresh => { + post_checkpoint(&checkpoint_url, "idle-started").await?; + tokio::task::yield_now().await; + post_checkpoint(&checkpoint_url, "idle-finished").await?; + + assert_eq!(list_tools(&client).await?.tools, Vec::new()); + assert_eq!(list_tools_calls.load(Ordering::SeqCst), 1); + assert_refreshed_credentials_persisted()?; } } @@ -274,6 +419,72 @@ async fn oauth_request_time_child() -> anyhow::Result<()> { Ok(()) } +fn near_expiry_initial_tokens(server_url: &str) -> anyhow::Result { + let expires_in = Duration::from_secs(1); + let mut response = OAuthTokenResponse::new( + AccessToken::new(INITIAL_ACCESS_TOKEN.to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + response.set_refresh_token(Some(RefreshToken::new(INITIAL_REFRESH_TOKEN.to_string()))); + response.set_expires_in(Some(&expires_in)); + let now = SystemTime::now().duration_since(UNIX_EPOCH)?; + Ok(StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: format!("{server_url}/mcp"), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: Some(now.saturating_add(expires_in).as_millis() as u64), + }) +} + +async fn list_tools(client: &RmcpClient) -> anyhow::Result { + client + .list_tools(/*params*/ None, Some(Duration::from_secs(5))) + .await +} + +fn credentials_path() -> anyhow::Result { + Ok(PathBuf::from(std::env::var("CODEX_HOME")?).join(CREDENTIALS_FILE)) +} + +fn assert_refreshed_credentials_persisted() -> anyhow::Result<()> { + let persisted: Value = serde_json::from_slice(&fs::read(credentials_path()?)?)?; + let entries = match persisted.as_object() { + Some(entries) => entries, + None => anyhow::bail!("persisted credentials were not an object: {persisted}"), + }; + assert_eq!(entries.len(), 1); + let entry = match entries.values().next() { + Some(entry) => entry, + None => anyhow::bail!("persisted credentials object was empty"), + }; + assert_eq!( + entry.get("access_token").and_then(Value::as_str), + Some(REFRESHED_ACCESS_TOKEN) + ); + assert_eq!( + entry.get("refresh_token").and_then(Value::as_str), + Some(REFRESHED_REFRESH_TOKEN) + ); + let expires_at = match entry.get("expires_at").and_then(Value::as_u64) { + Some(expires_at) => expires_at, + None => anyhow::bail!("persisted refreshed credentials had no expiry: {entry}"), + }; + let now = SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis() as u64; + assert!(expires_at > now); + Ok(()) +} + +async fn post_checkpoint(base_url: &str, checkpoint: &str) -> anyhow::Result<()> { + reqwest::Client::new() + .post(format!("{base_url}/checkpoint/{checkpoint}")) + .send() + .await? + .error_for_status()?; + Ok(()) +} + async fn initialize_in_process_client(client: &RmcpClient) -> anyhow::Result<()> { let params = InitializeRequestParams::new( ClientCapabilities::default(), diff --git a/codex-rs/rmcp-client/tests/streamable_http_oauth_lifecycle.rs b/codex-rs/rmcp-client/tests/streamable_http_oauth_lifecycle.rs deleted file mode 100644 index ee3c777c09d4..000000000000 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_lifecycle.rs +++ /dev/null @@ -1,307 +0,0 @@ -mod streamable_http_test_support; - -use std::sync::Arc; -use std::sync::Mutex; -use std::time::Duration; - -use codex_config::types::OAuthCredentialsStoreMode; -use codex_exec_server::Environment; -use codex_rmcp_client::RmcpClient; -use codex_rmcp_client::StoredOAuthTokens; -use codex_rmcp_client::WrappedOAuthTokenResponse; -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; -use serde_json::json; -use tempfile::TempDir; -use tokio::process::Command; -use wiremock::Mock; -use wiremock::MockServer; -use wiremock::Request; -use wiremock::Respond; -use wiremock::ResponseTemplate; -use wiremock::matchers::method; -use wiremock::matchers::path; - -use streamable_http_test_support::initialize_client; - -const SERVER_NAME: &str = "test-streamable-http-oauth-lifecycle"; -const INITIAL_ACCESS_TOKEN: &str = "initial-expired-access-token"; -const INITIAL_REFRESH_TOKEN: &str = "initial-refresh-token"; -const LONG_LIVED_ACCESS_TOKEN: &str = "long-lived-access-token"; -const LONG_LIVED_REFRESH_TOKEN: &str = "long-lived-refresh-token"; -const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_LIFECYCLE_SERVER_URL"; -const CHILD_CHECKPOINT_URL_ENV: &str = "MCP_TEST_OAUTH_LIFECYCLE_CHECKPOINT_URL"; - -#[derive(Clone, Debug, PartialEq, Eq)] -enum TimelineEvent { - Refresh(String), - Mcp { - method: String, - authorization: String, - }, - IdleStarted, - IdleFinished, -} - -#[derive(Clone)] -struct Timeline { - events: Arc>>, -} - -impl Timeline { - fn new() -> Self { - Self { - events: Arc::new(Mutex::new(Vec::new())), - } - } - - fn push(&self, event: TimelineEvent) { - self.events - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .push(event); - } - - fn snapshot(&self) -> Vec { - self.events - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .clone() - } -} - -#[derive(Clone)] -struct TokenResponder { - timeline: Timeline, -} - -impl Respond for TokenResponder { - fn respond(&self, request: &Request) -> ResponseTemplate { - let body = String::from_utf8_lossy(&request.body); - assert!( - body.contains(INITIAL_REFRESH_TOKEN), - "unexpected refresh request body: {body}" - ); - self.timeline - .push(TimelineEvent::Refresh(INITIAL_REFRESH_TOKEN.to_string())); - - ResponseTemplate::new(200).set_body_json(json!({ - "access_token": LONG_LIVED_ACCESS_TOKEN, - "token_type": "Bearer", - "expires_in": 7200, - "refresh_token": LONG_LIVED_REFRESH_TOKEN, - })) - } -} - -#[derive(Clone)] -struct McpResponder { - timeline: Timeline, -} - -impl Respond for McpResponder { - fn respond(&self, request: &Request) -> ResponseTemplate { - let body: Value = match request.body_json() { - Ok(body) => body, - Err(error) => panic!("invalid JSON-RPC request: {error}"), - }; - let method = match body.get("method").and_then(Value::as_str) { - Some(method) => method, - None => panic!("JSON-RPC request missing method: {body}"), - }; - let authorization = match request - .headers - .get("authorization") - .and_then(|value| value.to_str().ok()) - { - Some(authorization) => authorization.to_string(), - None => panic!("MCP request missing authorization header"), - }; - self.timeline.push(TimelineEvent::Mcp { - method: method.to_string(), - authorization, - }); - - match method { - "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-lifecycle-test", - "version": "0.0.0-test", - }, - }, - })), - "notifications/initialized" => ResponseTemplate::new(202), - "tools/list" => ResponseTemplate::new(200).set_body_json(json!({ - "jsonrpc": "2.0", - "id": body.get("id").cloned().unwrap_or(Value::Null), - "result": { - "tools": [], - }, - })), - _ => ResponseTemplate::new(400) - .set_body_string(format!("unexpected JSON-RPC method: {method}")), - } - } -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn does_not_refresh_in_background_while_idle() -> anyhow::Result<()> { - let server = MockServer::start().await; - let timeline = Timeline::new(); - - 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; - Mock::given(method("POST")) - .and(path("/oauth/token")) - .respond_with(TokenResponder { - timeline: timeline.clone(), - }) - .expect(1) - .mount(&server) - .await; - Mock::given(method("POST")) - .and(path("/mcp")) - .respond_with(McpResponder { - timeline: timeline.clone(), - }) - .expect(3) - .mount(&server) - .await; - - for (checkpoint_path, event) in [ - ("/checkpoint/idle-started", TimelineEvent::IdleStarted), - ("/checkpoint/idle-finished", TimelineEvent::IdleFinished), - ] { - let timeline = timeline.clone(); - Mock::given(method("POST")) - .and(path(checkpoint_path)) - .respond_with(move |_: &Request| { - timeline.push(event.clone()); - ResponseTemplate::new(204) - }) - .expect(1) - .mount(&server) - .await; - } - - let codex_home = TempDir::new()?; - let status = Command::new(std::env::current_exe()?) - .args([ - "oauth_lifecycle_child", - "--exact", - "--ignored", - "--nocapture", - ]) - .env("CODEX_HOME", codex_home.path()) - .env(CHILD_SERVER_URL_ENV, format!("{}/mcp", server.uri())) - .env(CHILD_CHECKPOINT_URL_ENV, server.uri()) - .status() - .await?; - assert!(status.success(), "OAuth lifecycle child failed: {status}"); - - assert_eq!( - timeline.snapshot(), - vec![ - TimelineEvent::Refresh(INITIAL_REFRESH_TOKEN.to_string()), - TimelineEvent::Mcp { - method: "initialize".to_string(), - authorization: format!("Bearer {LONG_LIVED_ACCESS_TOKEN}"), - }, - TimelineEvent::Mcp { - method: "notifications/initialized".to_string(), - authorization: format!("Bearer {LONG_LIVED_ACCESS_TOKEN}"), - }, - TimelineEvent::IdleStarted, - TimelineEvent::IdleFinished, - TimelineEvent::Mcp { - method: "tools/list".to_string(), - authorization: format!("Bearer {LONG_LIVED_ACCESS_TOKEN}"), - }, - ] - ); - server.verify().await; - Ok(()) -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 1)] -#[ignore = "spawned by OAuth lifecycle parent test"] -async fn oauth_lifecycle_child() -> anyhow::Result<()> { - let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; - let checkpoint_url = std::env::var(CHILD_CHECKPOINT_URL_ENV)?; - - let mut response = OAuthTokenResponse::new( - AccessToken::new(INITIAL_ACCESS_TOKEN.to_string()), - BasicTokenType::Bearer, - VendorExtraTokenFields::default(), - ); - response.set_refresh_token(Some(RefreshToken::new(INITIAL_REFRESH_TOKEN.to_string()))); - response.set_expires_in(Some(&Duration::from_secs(7200))); - save_oauth_tokens( - SERVER_NAME, - &StoredOAuthTokens { - server_name: SERVER_NAME.to_string(), - url: server_url.clone(), - client_id: "test-client-id".to_string(), - token_response: WrappedOAuthTokenResponse(response), - expires_at: Some(0), - }, - OAuthCredentialsStoreMode::File, - )?; - - let client = RmcpClient::new_streamable_http_client( - SERVER_NAME, - &server_url, - /*bearer_token*/ None, - /*http_headers*/ None, - /*env_http_headers*/ None, - OAuthCredentialsStoreMode::File, - Environment::default_for_tests().get_http_client(), - /*auth_provider*/ None, - ) - .await?; - initialize_client(&client).await?; - - post_checkpoint(&checkpoint_url, "idle-started").await?; - tokio::time::sleep(Duration::from_millis(100)).await; - post_checkpoint(&checkpoint_url, "idle-finished").await?; - - assert_eq!( - client - .list_tools(/*params*/ None, Some(Duration::from_secs(5))) - .await? - .tools, - Vec::new() - ); - Ok(()) -} - -async fn post_checkpoint(base_url: &str, checkpoint: &str) -> anyhow::Result<()> { - reqwest::Client::new() - .post(format!("{base_url}/checkpoint/{checkpoint}")) - .send() - .await? - .error_for_status()?; - Ok(()) -}