diff --git a/components/wpa_supplicant/src/rsn_supp/wpa.c b/components/wpa_supplicant/src/rsn_supp/wpa.c
index cc79358ceb..5e3e363bb2 100644
--- a/components/wpa_supplicant/src/rsn_supp/wpa.c
+++ b/components/wpa_supplicant/src/rsn_supp/wpa.c
@@ -780,6 +780,8 @@ void   wpa_supplicant_key_neg_complete(struct wpa_sm *sm,
     wpa_sm_set_state(WPA_COMPLETED);
 
     sm->wpa_neg_complete();
+    sm->eapol1_count = 0;
+    sm->use_ext_key_id = 0;
 
     if (secure) {
         wpa_sm_mlme_setprotection(
@@ -1271,7 +1273,7 @@ int   ieee80211w_set_keys(struct wpa_sm *sm,
         goto failed;
     }
 
-    if (sm->key_install && sm->key_info & WPA_KEY_INFO_INSTALL) {
+    if (sm->key_install && sm->key_info & WPA_KEY_INFO_INSTALL && sm->use_ext_key_id) {
         wpa_supplicant_install_ptk(sm, 0);
     }
 
@@ -1293,8 +1295,13 @@ failed:
        u16 key_info=sm->key_info;
 
     if (sm->key_install && key_info & WPA_KEY_INFO_INSTALL) {
-        if (wpa_supplicant_install_ptk(sm, 2))
-            goto failed;
+        if (sm->use_ext_key_id) {
+            if (wpa_supplicant_install_ptk(sm, 2))
+                goto failed;
+        } else {
+            if (wpa_supplicant_install_ptk(sm, 1))
+                goto failed;
+        }
     }
     else if (sm->key_install == false) {
         wpa_printf(MSG_DEBUG, "PTK has been installed, it may be an attack, ignor it.");
@@ -2154,6 +2161,7 @@ int wpa_set_bss(char *macddr, char * bssid, u8 pairwise_cipher, u8 group_cipher,
     memcpy(sm->own_addr, macddr, ETH_ALEN);
     memcpy(sm->bssid, bssid, ETH_ALEN);
     sm->ap_notify_completed_rsne = esp_wifi_sta_is_ap_notify_completed_rsne_internal();
+    sm->use_ext_key_id = (sm->proto == WPA_PROTO_WPA);
 
     if (sm->key_mgmt == WPA_KEY_MGMT_SAE ||
 	    is_wpa2_enterprise_connection()) {
diff --git a/components/wpa_supplicant/src/rsn_supp/wpa_i.h b/components/wpa_supplicant/src/rsn_supp/wpa_i.h
index e2911fbc92..88a5c84f4b 100644
--- a/components/wpa_supplicant/src/rsn_supp/wpa_i.h
+++ b/components/wpa_supplicant/src/rsn_supp/wpa_i.h
@@ -78,7 +78,7 @@ struct wpa_sm {
     struct install_key install_gtk;
     int mic_errors_seen; /* Michael MIC errors with the current PTK */
     int  key_entry_valid;   //present current avaliable entry for bssid, for pairkey:0,5,10,15,20, gtk: pairkey_no+i (i:1~4)
-
+    int use_ext_key_id; /* Enabled only for WPA PSK first key exchange */
     void (* sendto) (void *buffer, uint16_t len);
     void (*config_assoc_ie) (u8 proto, u8 *assoc_buf, u32 assoc_wpa_ie_len);
     void (*install_ppkey) (enum wpa_alg alg, u8 *addr, int key_idx, int set_tx,