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()),