From 4de7a2b9d8eae19e00ca7f744647fa1aabdc204f Mon Sep 17 00:00:00 2001 From: "Adam Perry @ OpenAI" Date: Thu, 4 Jun 2026 19:31:06 -0700 Subject: [PATCH] fix(rmcp): refresh expired OAuth tokens before startup (#26482) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Why Codex persists OAuth expiry as an absolute `expires_at`, then reconstructs RMCP’s relative `expires_in` when credentials are loaded. For an already-expired token, Codex reconstructed `expires_in` as missing. [RMCP 0.15 treated a missing `expires_in` as zero when a refresh token was present](https://github.com/modelcontextprotocol/rust-sdk/blob/9cfc905a9ef17c8bba6748dc0a9bdd2452681733/crates/rmcp/src/transport/auth.rs#L704-L723), so this still triggered a refresh. [RMCP 1.7 treats missing expiry information as unknown and uses the access token as-is](https://github.com/modelcontextprotocol/rust-sdk/blob/3529c3675ff64db805bd947ca6ece6090809e43d/crates/rmcp/src/transport/auth.rs#L1233-L1265), causing the stale token to be sent during `initialize`. ## What changed - Represent a known-expired persisted token as `expires_in = 0`, preserving `None` for genuinely unknown expiry. - Add Streamable HTTP coverage requiring the token to refresh before the startup handshake. ## Validation - The new regression test fails on RMCP 1.7 before the fix and passes afterward. - The same scenario passes on the commit immediately before the RMCP 1.7 update, using RMCP 0.15. - `just test -p codex-rmcp-client` (63 passed). --- codex-rs/Cargo.lock | 1 + codex-rs/rmcp-client/Cargo.toml | 1 + codex-rs/rmcp-client/src/oauth.rs | 12 +- .../tests/streamable_http_oauth_startup.rs | 155 ++++++++++++++++++ .../tests/streamable_http_test_support.rs | 9 +- 5 files changed, 173 insertions(+), 5 deletions(-) create mode 100644 codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 1e79865d4..6da768843 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3597,6 +3597,7 @@ dependencies = [ "urlencoding", "webbrowser", "which 8.0.0", + "wiremock", ] [[package]] diff --git a/codex-rs/rmcp-client/Cargo.toml b/codex-rs/rmcp-client/Cargo.toml index e3417de70..1006ad74e 100644 --- a/codex-rs/rmcp-client/Cargo.toml +++ b/codex-rs/rmcp-client/Cargo.toml @@ -70,6 +70,7 @@ codex-utils-cargo-bin = { workspace = true } pretty_assertions = { workspace = true } serial_test = { workspace = true } tempfile = { workspace = true } +wiremock = { workspace = true } [target.'cfg(target_os = "linux")'.dependencies] keyring = { workspace = true, features = ["linux-native-async-persistent"] } diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index e23eee84b..c34846079 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -113,7 +113,13 @@ fn refresh_expires_in_from_timestamp(tokens: &mut StoredOAuthTokens) { tokens.token_response.0.set_expires_in(Some(&duration)); } None => { - tokens.token_response.0.set_expires_in(None); + // RMCP treats a missing expiry as unknown and uses the access token + // as-is. Treat a known-expired timestamp as an explicit zero so + // startup refreshes the token before the first request. + tokens + .token_response + .0 + .set_expires_in(Some(&Duration::ZERO)); } } } @@ -830,7 +836,7 @@ mod tests { } #[test] - fn refresh_expires_in_from_timestamp_clears_expired_tokens() { + fn refresh_expires_in_from_timestamp_marks_expired_tokens() { let mut tokens = sample_tokens(); let now = SystemTime::now() .duration_since(UNIX_EPOCH) @@ -843,7 +849,7 @@ mod tests { super::refresh_expires_in_from_timestamp(&mut tokens); - assert!(tokens.token_response.0.expires_in().is_none()); + assert_eq!(tokens.token_response.0.expires_in(), Some(Duration::ZERO)); } fn assert_tokens_match_without_expiry( diff --git a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs new file mode 100644 index 000000000..1c18f2c98 --- /dev/null +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs @@ -0,0 +1,155 @@ +mod streamable_http_test_support; + +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 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::ResponseTemplate; +use wiremock::matchers::body_string_contains; +use wiremock::matchers::header; +use wiremock::matchers::method; +use wiremock::matchers::path; + +use streamable_http_test_support::initialize_client; + +const SERVER_NAME: &str = "test-streamable-http-oauth-startup"; +const EXPIRED_ACCESS_TOKEN: &str = "expired-access-token"; +const REFRESH_TOKEN: &str = "valid-refresh-token"; +const REFRESHED_ACCESS_TOKEN: &str = "refreshed-access-token"; +const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_STARTUP_SERVER_URL"; + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": [""], + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={REFRESH_TOKEN}" + ))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": REFRESHED_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": REFRESH_TOKEN, + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {REFRESHED_ACCESS_TOKEN}"), + )) + .respond_with(|request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body.get("method").and_then(Value::as_str) { + Some("initialize") => ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { + "protocolVersion": body + .pointer("/params/protocolVersion") + .cloned() + .unwrap_or_else(|| json!("2025-06-18")), + "capabilities": {}, + "serverInfo": { + "name": "oauth-startup-test", + "version": "0.0.0-test", + }, + }, + })), + Some("notifications/initialized") => ResponseTemplate::new(202), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + }) + .expect(2) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + let server_url = format!("{}/mcp", server.uri()); + + // Credential storage resolves CODEX_HOME from the process environment. + // Run the client half of the test in an ignored helper test so it can use + // an isolated home without mutating the parent test runner's environment. + let status = Command::new(std::env::current_exe()?) + .args(["oauth_startup_child", "--exact", "--ignored", "--nocapture"]) + .env("CODEX_HOME", codex_home.path()) + .env(CHILD_SERVER_URL_ENV, server_url) + .status() + .await?; + assert!(status.success(), "OAuth startup child failed: {status}"); + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by refreshes_expired_persisted_token_before_initialize"] +async fn oauth_startup_child() -> anyhow::Result<()> { + let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; + + // Save an expired access token with a valid refresh token so startup must + // refresh before sending the initialize request. + let mut response = OAuthTokenResponse::new( + AccessToken::new(EXPIRED_ACCESS_TOKEN.to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + response.set_refresh_token(Some(RefreshToken::new(REFRESH_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(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: Some(0), + }; + save_oauth_tokens(SERVER_NAME, &tokens, OAuthCredentialsStoreMode::File)?; + + // 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( + 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?; + Ok(()) +} 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 822acef1a..03ee79392 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_test_support.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_test_support.rs @@ -86,6 +86,12 @@ pub(crate) async fn create_client(base_url: &str) -> anyhow::Result ) .await?; + initialize_client(&client).await?; + + Ok(client) +} + +pub(crate) async fn initialize_client(client: &RmcpClient) -> anyhow::Result<()> { client .initialize( init_params(), @@ -102,8 +108,7 @@ pub(crate) async fn create_client(base_url: &str) -> anyhow::Result }), ) .await?; - - Ok(client) + Ok(()) } /// Creates a Streamable HTTP RMCP client that sends traffic through the remote