diff --git a/drivers/wifi/nrf700x/osal/fw_if/umac_if/src/fmac_ap.c b/drivers/wifi/nrf700x/osal/fw_if/umac_if/src/fmac_ap.c
index 2bfa8fd65d8e..51479a02538a 100644
--- a/drivers/wifi/nrf700x/osal/fw_if/umac_if/src/fmac_ap.c
+++ b/drivers/wifi/nrf700x/osal/fw_if/umac_if/src/fmac_ap.c
@@ -13,6 +13,7 @@
 #include "fmac_peer.h"
 #include "queue.h"
 #include "fmac_tx.h"
+#include "fmac_util.h"
 
 enum wifi_nrf_status sap_client_ps_get_frames(struct wifi_nrf_fmac_dev_ctx *fmac_dev_ctx,
 					      struct nrf_wifi_sap_ps_get_frames *config)
@@ -23,6 +24,8 @@ enum wifi_nrf_status sap_client_ps_get_frames(struct wifi_nrf_fmac_dev_ctx *fmac
 	int id = -1;
 	int ac = 0;
 	int desc = 0;
+	struct wifi_nrf_fmac_dev_ctx_def *def_dev_ctx = NULL;
+	struct wifi_nrf_fmac_priv_def *def_priv = NULL;
 
 	if (!fmac_dev_ctx || !config) {
 		wifi_nrf_osal_log_err(fmac_dev_ctx->fpriv->opriv,
@@ -31,8 +34,11 @@ enum wifi_nrf_status sap_client_ps_get_frames(struct wifi_nrf_fmac_dev_ctx *fmac
 		goto out;
 	}
 
+	def_dev_ctx = wifi_dev_priv(fmac_dev_ctx);
+	def_priv = wifi_fmac_priv(fmac_dev_ctx->fpriv);
+
 	wifi_nrf_osal_spinlock_take(fmac_dev_ctx->fpriv->opriv,
-				    fmac_dev_ctx->tx_config.tx_lock);
+				    def_dev_ctx->tx_config.tx_lock);
 
 	id = wifi_nrf_fmac_peer_get_id(fmac_dev_ctx, config->mac_addr);
 
@@ -43,15 +49,15 @@ enum wifi_nrf_status sap_client_ps_get_frames(struct wifi_nrf_fmac_dev_ctx *fmac
 				      config->mac_addr);
 
 		wifi_nrf_osal_spinlock_rel(fmac_dev_ctx->fpriv->opriv,
-					   fmac_dev_ctx->tx_config.tx_lock);
+					   def_dev_ctx->tx_config.tx_lock);
 		goto out;
 	}
 
 
-	peer = &fmac_dev_ctx->tx_config.peers[id];
+	peer = &def_dev_ctx->tx_config.peers[id];
 	peer->ps_token_count = config->num_frames;
 
-	wakeup_client_q = fmac_dev_ctx->tx_config.wakeup_client_q;
+	wakeup_client_q = def_dev_ctx->tx_config.wakeup_client_q;
 
 	if (wakeup_client_q) {
 		wifi_nrf_utils_q_enqueue(fmac_dev_ctx->fpriv->opriv,
@@ -62,13 +68,13 @@ enum wifi_nrf_status sap_client_ps_get_frames(struct wifi_nrf_fmac_dev_ctx *fmac
 	for (ac = WIFI_NRF_FMAC_AC_VO; ac >= 0; --ac) {
 		desc = tx_desc_get(fmac_dev_ctx, ac);
 
-		if (desc < fmac_dev_ctx->fpriv->num_tx_tokens) {
+		if (desc < def_priv->num_tx_tokens) {
 			tx_pending_process(fmac_dev_ctx, desc, ac);
 		}
 	}
 
 	wifi_nrf_osal_spinlock_rel(fmac_dev_ctx->fpriv->opriv,
-				   fmac_dev_ctx->tx_config.tx_lock);
+				   def_dev_ctx->tx_config.tx_lock);
 
 	status = WIFI_NRF_STATUS_SUCCESS;
 out:
@@ -88,9 +94,6 @@ enum wifi_nrf_status sap_client_update_pmmode(struct wifi_nrf_fmac_dev_ctx *fmac
 	struct wifi_nrf_fmac_dev_ctx_def *def_dev_ctx = NULL;
 	struct wifi_nrf_fmac_priv_def *def_priv = NULL;
 
-	def_dev_ctx = wifi_dev_priv(fmac_dev_ctx);
-	def_priv = wifi_fmac_priv(fmac_dev_ctx->fpriv);
-
 	if (!fmac_dev_ctx || !config) {
 		wifi_nrf_osal_log_err(fmac_dev_ctx->fpriv->opriv,
 				      "%s: Invalid params\n",
@@ -98,6 +101,9 @@ enum wifi_nrf_status sap_client_update_pmmode(struct wifi_nrf_fmac_dev_ctx *fmac
 		goto out;
 	}
 
+	def_dev_ctx = wifi_dev_priv(fmac_dev_ctx);
+	def_priv = wifi_fmac_priv(fmac_dev_ctx->fpriv);
+
 	wifi_nrf_osal_spinlock_take(fmac_dev_ctx->fpriv->opriv,
 				    def_dev_ctx->tx_config.tx_lock);
 
@@ -139,7 +145,7 @@ enum wifi_nrf_status sap_client_update_pmmode(struct wifi_nrf_fmac_dev_ctx *fmac
 	}
 
 	wifi_nrf_osal_spinlock_rel(fmac_dev_ctx->fpriv->opriv,
-				   fmac_dev_ctx->tx_config.tx_lock);
+				   def_dev_ctx->tx_config.tx_lock);
 
 	status = WIFI_NRF_STATUS_SUCCESS;