diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index bf98513415..c526af6bb1 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -611,10 +611,6 @@ impl OAuthPersistor { return Ok(()); } - let snapshot = { - let guard = self.inner.last_credentials.lock().await; - guard.clone() - }; let key = compute_store_key(&self.inner.server_name, &self.inner.url)?; let _lock = RefreshCredentialLock::acquire(&key).await?; let latest = load_oauth_tokens_with_keyring_store( @@ -625,7 +621,7 @@ impl OAuthPersistor { self.inner.keyring_backend_kind, )?; - if latest.is_none() && snapshot.is_some() { + let Some(latest) = latest else { self.clear_manager_credentials().await; let mut last_credentials = self.inner.last_credentials.lock().await; *last_credentials = None; @@ -633,15 +629,15 @@ impl OAuthPersistor { "OAuth tokens for server {} were removed before refresh; authorization required", self.inner.server_name ); - } + }; - if latest != snapshot { - if let Some(latest) = latest { - self.adopt_credentials(latest).await?; - } + if !token_needs_refresh(latest.expires_at) { + self.adopt_credentials(latest).await?; return Ok(()); } + self.adopt_credentials(latest).await?; + { let manager = self.inner.authorization_manager.clone(); let guard = manager.lock().await; @@ -1261,29 +1257,14 @@ mod tests { -> Result<()> { let _env = TempCodexHome::new(); 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": ["scope-a", "scope-b"], - }))) - .mount(&server) - .await; - Mock::given(method("POST")) - .and(path("/oauth/token")) - .and(body_string_contains("grant_type=refresh_token")) - .and(body_string_contains("refresh_token=refresh-token")) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "access_token": "refreshed-access-token", - "token_type": "Bearer", - "expires_in": 3600, - "refresh_token": "rotated-refresh-token", - "scope": "scope-a scope-b", - }))) - .expect(1) - .mount(&server) - .await; + mount_oauth_metadata(&server).await; + mount_refresh_response( + &server, + "refresh-token", + "refreshed-access-token", + "rotated-refresh-token", + ) + .await; let store = MockKeyringStore::default(); let initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri())); @@ -1341,6 +1322,195 @@ mod tests { Ok(()) } + #[tokio::test] + async fn refresh_transaction_adopts_valid_reread_without_provider_refresh() -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + let store = MockKeyringStore::default(); + let mut initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri())); + initial_tokens + .token_response + .0 + .set_refresh_token(Some(RefreshToken::new("stale-refresh-token".to_string()))); + + let mut latest_tokens = sample_tokens(); + latest_tokens.url.clone_from(&initial_tokens.url); + latest_tokens + .token_response + .0 + .set_access_token(AccessToken::new( + "already-refreshed-access-token".to_string(), + )); + latest_tokens + .token_response + .0 + .set_refresh_token(Some(RefreshToken::new( + "already-rotated-refresh-token".to_string(), + ))); + + super::save_oauth_tokens_with_keyring_store( + &store, + &latest_tokens.server_name, + &latest_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + + let manager = authorization_manager_for(&initial_tokens).await?; + let persistor = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + manager.clone(), + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + Some(initial_tokens), + ); + + persistor + .refresh_if_needed_with_keyring_store(&store) + .await?; + + let manager_tokens = tokens_from_manager(&manager).await?; + assert_eq!( + access_token(&manager_tokens), + "already-refreshed-access-token" + ); + assert_eq!( + refresh_token(&manager_tokens), + Some("already-rotated-refresh-token".to_string()) + ); + Ok(()) + } + + #[tokio::test] + async fn refresh_transaction_refreshes_when_only_derived_expires_in_drifted() -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + mount_refresh_response( + &server, + "refresh-token", + "refreshed-after-expiry-drift", + "rotated-after-expiry-drift", + ) + .await; + + let store = MockKeyringStore::default(); + let mut initial_tokens = sample_tokens(); + initial_tokens.url = format!("{}/mcp", server.uri()); + initial_tokens.expires_at = Some(now_millis().saturating_add(5_000)); + initial_tokens + .token_response + .0 + .set_expires_in(Some(&Duration::from_secs(3600))); + super::save_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + + let manager = authorization_manager_for(&initial_tokens).await?; + let persistor = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + manager, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + Some(initial_tokens.clone()), + ); + + persistor + .refresh_if_needed_with_keyring_store(&store) + .await?; + + server.verify().await; + let stored = super::load_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens.url, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )? + .expect("refreshed tokens should be persisted"); + assert_eq!(access_token(&stored), "refreshed-after-expiry-drift"); + assert_eq!( + refresh_token(&stored), + Some("rotated-after-expiry-drift".to_string()) + ); + Ok(()) + } + + #[tokio::test] + async fn refresh_transaction_uses_latest_refresh_token_when_reread_is_expired() -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + mount_refresh_response( + &server, + "latest-refresh-token", + "refreshed-from-latest-token", + "rotated-from-latest-token", + ) + .await; + + let store = MockKeyringStore::default(); + let mut initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri())); + initial_tokens + .token_response + .0 + .set_refresh_token(Some(RefreshToken::new("stale-refresh-token".to_string()))); + + let mut latest_tokens = initial_tokens.clone(); + latest_tokens + .token_response + .0 + .set_access_token(AccessToken::new("latest-expired-access-token".to_string())); + latest_tokens + .token_response + .0 + .set_refresh_token(Some(RefreshToken::new("latest-refresh-token".to_string()))); + super::save_oauth_tokens_with_keyring_store( + &store, + &latest_tokens.server_name, + &latest_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + + let manager = authorization_manager_for(&initial_tokens).await?; + let persistor = OAuthPersistor::new( + initial_tokens.server_name.clone(), + initial_tokens.url.clone(), + manager, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + Some(initial_tokens.clone()), + ); + + persistor + .refresh_if_needed_with_keyring_store(&store) + .await?; + + server.verify().await; + let stored = super::load_oauth_tokens_with_keyring_store( + &store, + &initial_tokens.server_name, + &initial_tokens.url, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )? + .expect("refreshed tokens should be persisted"); + assert_eq!(access_token(&stored), "refreshed-from-latest-token"); + assert_eq!( + refresh_token(&stored), + Some("rotated-from-latest-token".to_string()) + ); + Ok(()) + } + #[test] fn save_oauth_tokens_with_secrets_backend_falls_back_to_file_when_keyring_fails() -> Result<()> { @@ -1644,6 +1814,42 @@ mod tests { Ok(Arc::new(TokioMutex::new(manager))) } + 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": ["scope-a", "scope-b"], + }))) + .mount(server) + .await; + } + + async fn mount_refresh_response( + server: &MockServer, + request_refresh_token: &str, + response_access_token: &str, + response_refresh_token: &str, + ) { + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={request_refresh_token}" + ))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": response_access_token, + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": response_refresh_token, + "scope": "scope-a scope-b", + }))) + .expect(1) + .mount(server) + .await; + } + #[expect( clippy::await_holding_invalid_type, reason = "AuthorizationManager async access must be serialized through its mutex" @@ -1686,6 +1892,13 @@ mod tests { tokens } + fn now_millis() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_else(|_| Duration::from_secs(0)) + .as_millis() as u64 + } + fn sample_tokens() -> StoredOAuthTokens { let mut response = OAuthTokenResponse::new( AccessToken::new("access-token".to_string()), diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 862b59c5d4..fad6fb0200 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -491,7 +491,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(); @@ -507,7 +507,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(); @@ -552,7 +552,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(); @@ -568,7 +568,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(); @@ -584,7 +584,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(); @@ -602,7 +602,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) => { @@ -658,7 +658,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, @@ -688,7 +688,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(); @@ -752,12 +752,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( 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 713a0b4781..894ab000b3 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs @@ -39,6 +39,7 @@ const EXPIRED_ACCESS_TOKEN: &str = "expired-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 PREFLIGHT_REFRESH_ERROR: &str = "preflight refresh failed distinctly"; const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_STARTUP_SERVER_URL"; const UNREFRESHABLE_SERVER_URL: &str = "https://unrefreshable.example/mcp"; const UNEXPIRED_SERVER_URL: &str = "https://unexpired.example/mcp"; @@ -233,6 +234,86 @@ async fn concurrent_file_mode_startup_refreshes_once() -> anyhow::Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn operation_preflight_refresh_failure_blocks_rmcp_request() -> 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(&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(400).set_body_json(json!({ + "error": "invalid_grant", + "error_description": PREFLIGHT_REFRESH_ERROR, + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {EXPIRED_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-preflight-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()); + let status = Command::new(std::env::current_exe()?) + .args([ + "operation_preflight_refresh_failure_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(CHILD_SERVER_URL_ENV, &server_url) + .status() + .await?; + assert!( + status.success(), + "OAuth preflight failure child failed: {status}" + ); + + server.verify().await; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn reports_auth_status_for_persisted_credentials() -> anyhow::Result<()> { let codex_home = TempDir::new()?; @@ -392,6 +473,44 @@ async fn oauth_startup_child() -> anyhow::Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by operation_preflight_refresh_failure_blocks_rmcp_request"] +async fn operation_preflight_refresh_failure_child() -> anyhow::Result<()> { + let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; + save_near_expiry_file_mode_tokens(&server_url)?; + + let client = RmcpClient::new_streamable_http_client( + SERVER_NAME, + &server_url, + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await?; + + initialize_client(&client).await?; + + tokio::time::sleep(Duration::from_millis(1_200)).await; + let error = client + .list_tools(/*params*/ None, Some(Duration::from_secs(5))) + .await + .expect_err("preflight refresh failure should abort the operation"); + let message = format!("{error:#}"); + assert!( + message.contains("failed to refresh OAuth tokens for server"), + "unexpected error: {message}" + ); + assert!( + message.contains(PREFLIGHT_REFRESH_ERROR), + "unexpected error: {message}" + ); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] #[ignore = "spawned by concurrent_file_mode_startup_refreshes_once"] async fn oauth_concurrency_seed_child() -> anyhow::Result<()> { @@ -444,3 +563,31 @@ fn save_expired_file_mode_tokens(server_url: &str) -> anyhow::Result<()> { )?; Ok(()) } + +fn save_near_expiry_file_mode_tokens(server_url: &str) -> anyhow::Result<()> { + 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(31))); + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_else(|_| Duration::from_secs(0)) + .as_millis() as u64; + let tokens = StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: server_url.to_string(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: Some(now.saturating_add(/*rhs*/ 31_000)), + }; + save_oauth_tokens( + SERVER_NAME, + &tokens, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + Ok(()) +}