diff --git a/src/wifi/model/block-ack-agreement.h b/src/wifi/model/block-ack-agreement.h
index bd3ce5b47..418b72da3 100644
--- a/src/wifi/model/block-ack-agreement.h
+++ b/src/wifi/model/block-ack-agreement.h
@@ -174,16 +174,16 @@ class BlockAckAgreement
     static std::size_t GetDistance(uint16_t seqNumber, uint16_t startingSeqNumber);
 
   protected:
-    Mac48Address m_peer;       //!< Peer address
-    uint8_t m_amsduSupported;  //!< Flag whether MSDU aggregation is supported
-    uint8_t m_blockAckPolicy;  //!< Type of block ack: immediate or delayed
-    uint8_t m_tid;             //!< Traffic ID
-    uint16_t m_bufferSize;     //!< Buffer size
-    uint16_t m_timeout;        //!< Timeout
-    uint16_t m_startingSeq;    //!< Starting sequence control
-    uint16_t m_winEnd;         //!< Ending sequence number
-    uint8_t m_htSupported;     //!< Flag whether HT is supported
-    EventId m_inactivityEvent; //!< inactivity event
+    Mac48Address m_peer;               //!< Peer address
+    uint8_t m_amsduSupported;          //!< Flag whether MSDU aggregation is supported
+    uint8_t m_blockAckPolicy;          //!< Type of block ack: immediate or delayed
+    uint8_t m_tid;                     //!< Traffic ID
+    uint16_t m_bufferSize;             //!< Buffer size
+    uint16_t m_timeout;                //!< Timeout
+    uint16_t m_startingSeq;            //!< Starting sequence control
+    uint16_t m_winEnd;                 //!< Ending sequence number
+    uint8_t m_htSupported;             //!< Flag whether HT is supported
+    mutable EventId m_inactivityEvent; //!< inactivity event
 };
 
 } // namespace ns3
diff --git a/src/wifi/model/block-ack-manager.cc b/src/wifi/model/block-ack-manager.cc
index 8f2b975ce..6d0c128f6 100644
--- a/src/wifi/model/block-ack-manager.cc
+++ b/src/wifi/model/block-ack-manager.cc
@@ -20,6 +20,7 @@
 #include "block-ack-manager.h"
 
 #include "ctrl-headers.h"
+#include "mac-rx-middle.h"
 #include "mgt-headers.h"
 #include "qos-utils.h"
 #include "wifi-mac-queue.h"
@@ -104,22 +105,34 @@ BlockAckManager::GetAgreementAsOriginator(Mac48Address recipient, uint8_t tid) c
     return std::nullopt;
 }
 
+BlockAckManager::RecipientAgreementOptConstRef
+BlockAckManager::GetAgreementAsRecipient(Mac48Address originator, uint8_t tid) const
+{
+    NS_LOG_FUNCTION(this << originator << +tid);
+    if (auto it = m_recipientAgreements.find({originator, tid}); it != m_recipientAgreements.end())
+    {
+        return std::cref(it->second);
+    }
+
+    return std::nullopt;
+}
+
 void
-BlockAckManager::CreateOriginatorAgreement(const MgtAddBaRequestHeader* reqHdr,
+BlockAckManager::CreateOriginatorAgreement(const MgtAddBaRequestHeader& reqHdr,
                                            Mac48Address recipient,
                                            bool htSupported)
 {
     NS_LOG_FUNCTION(this << reqHdr << recipient << htSupported);
-    const uint8_t tid = reqHdr->GetTid();
+    const uint8_t tid = reqHdr.GetTid();
     OriginatorBlockAckAgreement agreement(recipient, tid);
-    agreement.SetStartingSequence(reqHdr->GetStartingSequence());
+    agreement.SetStartingSequence(reqHdr.GetStartingSequence());
     /* For now we assume that originator doesn't use this field. Use of this field
        is mandatory only for recipient */
-    agreement.SetBufferSize(reqHdr->GetBufferSize());
-    agreement.SetTimeout(reqHdr->GetTimeout());
-    agreement.SetAmsduSupport(reqHdr->IsAmsduSupported());
+    agreement.SetBufferSize(reqHdr.GetBufferSize());
+    agreement.SetTimeout(reqHdr.GetTimeout());
+    agreement.SetAmsduSupport(reqHdr.IsAmsduSupported());
     agreement.SetHtSupported(htSupported);
-    if (reqHdr->IsImmediateBlockAck())
+    if (reqHdr.IsImmediateBlockAck())
     {
         agreement.SetImmediateBlockAck();
     }
@@ -166,22 +179,22 @@ BlockAckManager::DestroyOriginatorAgreement(Mac48Address recipient, uint8_t tid)
 }
 
 void
-BlockAckManager::UpdateOriginatorAgreement(const MgtAddBaResponseHeader* respHdr,
+BlockAckManager::UpdateOriginatorAgreement(const MgtAddBaResponseHeader& respHdr,
                                            Mac48Address recipient,
                                            uint16_t startingSeq)
 {
     NS_LOG_FUNCTION(this << respHdr << recipient << startingSeq);
-    uint8_t tid = respHdr->GetTid();
+    uint8_t tid = respHdr.GetTid();
     auto it = m_originatorAgreements.find({recipient, tid});
     if (it != m_originatorAgreements.end())
     {
         OriginatorBlockAckAgreement& agreement = it->second.first;
-        agreement.SetBufferSize(respHdr->GetBufferSize());
-        agreement.SetTimeout(respHdr->GetTimeout());
-        agreement.SetAmsduSupport(respHdr->IsAmsduSupported());
+        agreement.SetBufferSize(respHdr.GetBufferSize());
+        agreement.SetTimeout(respHdr.GetTimeout());
+        agreement.SetAmsduSupport(respHdr.IsAmsduSupported());
         agreement.SetStartingSequence(startingSeq);
         agreement.InitTxWindow();
-        if (respHdr->IsImmediateBlockAck())
+        if (respHdr.IsImmediateBlockAck())
         {
             agreement.SetImmediateBlockAck();
         }
@@ -210,6 +223,50 @@ BlockAckManager::UpdateOriginatorAgreement(const MgtAddBaResponseHeader* respHdr
     m_unblockPackets(recipient, tid);
 }
 
+void
+BlockAckManager::CreateRecipientAgreement(const MgtAddBaResponseHeader& respHdr,
+                                          Mac48Address originator,
+                                          uint16_t startingSeq,
+                                          bool htSupported,
+                                          Ptr<MacRxMiddle> rxMiddle)
+{
+    NS_LOG_FUNCTION(this << respHdr << originator << startingSeq << htSupported << rxMiddle);
+    uint8_t tid = respHdr.GetTid();
+
+    RecipientBlockAckAgreement agreement(originator,
+                                         respHdr.IsAmsduSupported(),
+                                         tid,
+                                         respHdr.GetBufferSize(),
+                                         respHdr.GetTimeout(),
+                                         startingSeq,
+                                         htSupported);
+    agreement.SetMacRxMiddle(rxMiddle);
+    if (respHdr.IsImmediateBlockAck())
+    {
+        agreement.SetImmediateBlockAck();
+    }
+    else
+    {
+        agreement.SetDelayedBlockAck();
+    }
+
+    m_recipientAgreements.insert_or_assign({originator, tid}, agreement);
+}
+
+void
+BlockAckManager::DestroyRecipientAgreement(Mac48Address originator, uint8_t tid)
+{
+    NS_LOG_FUNCTION(this << originator << tid);
+
+    if (auto agreementIt = m_recipientAgreements.find({originator, tid});
+        agreementIt != m_recipientAgreements.end())
+    {
+        // forward up the buffered MPDUs before destroying the agreement
+        agreementIt->second.Flush();
+        m_recipientAgreements.erase(agreementIt);
+    }
+}
+
 void
 BlockAckManager::StorePacket(Ptr<WifiMpdu> mpdu)
 {
@@ -637,6 +694,36 @@ BlockAckManager::NotifyDiscardedMpdu(Ptr<const WifiMpdu> mpdu)
     ScheduleBar(Create<const WifiMpdu>(bar, hdr));
 }
 
+void
+BlockAckManager::NotifyGotBlockAckRequest(Mac48Address originator,
+                                          uint8_t tid,
+                                          uint16_t startingSeq)
+{
+    NS_LOG_FUNCTION(this << originator << tid << startingSeq);
+    auto it = m_recipientAgreements.find({originator, tid});
+    if (it == m_recipientAgreements.end())
+    {
+        return;
+    }
+    it->second.NotifyReceivedBar(startingSeq);
+}
+
+void
+BlockAckManager::NotifyGotMpdu(Ptr<const WifiMpdu> mpdu)
+{
+    NS_LOG_FUNCTION(this << *mpdu);
+    auto originator = mpdu->GetHeader().GetAddr2();
+    NS_ASSERT(mpdu->GetHeader().IsQosData());
+    auto tid = mpdu->GetHeader().GetQosTid();
+
+    auto it = m_recipientAgreements.find({originator, tid});
+    if (it == m_recipientAgreements.end())
+    {
+        return;
+    }
+    it->second.NotifyReceivedMpdu(mpdu);
+}
+
 CtrlBAckRequestHeader
 BlockAckManager::GetBlockAckReqHeader(Mac48Address recipient, uint8_t tid) const
 {
diff --git a/src/wifi/model/block-ack-manager.h b/src/wifi/model/block-ack-manager.h
index c82c0d6c0..3cdae667d 100644
--- a/src/wifi/model/block-ack-manager.h
+++ b/src/wifi/model/block-ack-manager.h
@@ -22,6 +22,7 @@
 
 #include "block-ack-type.h"
 #include "originator-block-ack-agreement.h"
+#include "recipient-block-ack-agreement.h"
 #include "wifi-mac-header.h"
 #include "wifi-mpdu.h"
 #include "wifi-tx-vector.h"
@@ -40,6 +41,7 @@ class MgtAddBaRequestHeader;
 class CtrlBAckResponseHeader;
 class CtrlBAckRequestHeader;
 class WifiMacQueue;
+class MacRxMiddle;
 
 /**
  * \ingroup wifi
@@ -103,6 +105,9 @@ class BlockAckManager : public Object
     /// optional const reference to OriginatorBlockAckAgreement
     using OriginatorAgreementOptConstRef =
         std::optional<std::reference_wrapper<const OriginatorBlockAckAgreement>>;
+    /// optional const reference to RecipientBlockAckAgreement
+    using RecipientAgreementOptConstRef =
+        std::optional<std::reference_wrapper<const RecipientBlockAckAgreement>>;
 
     /**
      * \param recipient MAC address of the recipient
@@ -114,6 +119,17 @@ class BlockAckManager : public Object
      */
     OriginatorAgreementOptConstRef GetAgreementAsOriginator(Mac48Address recipient,
                                                             uint8_t tid) const;
+    /**
+     * \param originator MAC address of the originator
+     * \param tid Traffic ID
+     *
+     * \return a const reference to the block ack agreement with the given originator, if it exists
+     *
+     * Check if we are the recipient of an existing block ack agreement with the given originator.
+     */
+    RecipientAgreementOptConstRef GetAgreementAsRecipient(Mac48Address originator,
+                                                          uint8_t tid) const;
+
     /**
      * \param reqHdr Relative Add block ack request (action frame).
      * \param recipient Address of peer station involved in block ack mechanism.
@@ -122,7 +138,7 @@ class BlockAckManager : public Object
      * Creates a new originator block ack agreement in pending state. When a ADDBA response
      * with a successful status code is received, the relative agreement becomes established.
      */
-    void CreateOriginatorAgreement(const MgtAddBaRequestHeader* reqHdr,
+    void CreateOriginatorAgreement(const MgtAddBaRequestHeader& reqHdr,
                                    Mac48Address recipient,
                                    bool htSupported = true);
     /**
@@ -140,9 +156,40 @@ class BlockAckManager : public Object
      *
      * Invoked upon receipt of a ADDBA response frame from <i>recipient</i>.
      */
-    void UpdateOriginatorAgreement(const MgtAddBaResponseHeader* respHdr,
+    void UpdateOriginatorAgreement(const MgtAddBaResponseHeader& respHdr,
                                    Mac48Address recipient,
                                    uint16_t startingSeq);
+
+    /**
+     * \param respHdr Add block ack response from originator (action
+     *        frame).
+     * \param originator Address of peer station involved in block ack
+     *        mechanism.
+     * \param startingSeq Sequence number of the first MPDU of all
+     *        packets for which block ack was negotiated.
+     * \param htSupported whether HT support is enabled
+     * \param rxMiddle the MAC RX Middle on this station
+     *
+     * This function is typically invoked only by ns3::WifiMac
+     * when the STA (which may be non-AP in ESS, or in an IBSS) has
+     * received an ADDBA Request frame and is transmitting an ADDBA
+     * Response frame. At this point the frame exchange manager must
+     * allocate buffers to collect all correctly received packets belonging
+     * to the category for which block ack was negotiated.
+     */
+    void CreateRecipientAgreement(const MgtAddBaResponseHeader& respHdr,
+                                  Mac48Address originator,
+                                  uint16_t startingSeq,
+                                  bool htSupported,
+                                  Ptr<MacRxMiddle> rxMiddle);
+    /**
+     * Destroy a recipient Block Ack agreement.
+     *
+     * \param originator the originator MAC address
+     * \param tid the TID associated with the Block Ack agreement
+     */
+    void DestroyRecipientAgreement(Mac48Address originator, uint8_t tid);
+
     /**
      * \param mpdu MPDU to store.
      *
@@ -217,6 +264,20 @@ class BlockAckManager : public Object
      * sent with ack policy set to Block Ack, should be placed in the retransmission queue.
      */
     void NotifyMissedBlockAck(uint8_t linkId, Mac48Address recipient, uint8_t tid);
+    /**
+     * \param originator MAC address of the sender of the Block Ack Request
+     * \param tid Traffic ID
+     * \param startingSeq the starting sequence number in the Block Ack Request
+     *
+     * Perform required actions upon receiving a Block Ack Request frame.
+     */
+    void NotifyGotBlockAckRequest(Mac48Address originator, uint8_t tid, uint16_t startingSeq);
+    /**
+     * \param mpdu the received MPDU
+     *
+     * Perform required actions upon receiving an MPDU.
+     */
+    void NotifyGotMpdu(Ptr<const WifiMpdu> mpdu);
     /**
      * \param recipient Address of peer station involved in block ack mechanism.
      * \param tid Traffic ID.
@@ -427,6 +488,9 @@ class BlockAckManager : public Object
     /// typedef for an iterator for Agreements
     using OriginatorAgreementsI = OriginatorAgreements::iterator;
 
+    /// AgreementKey-indexed map of recipient block ack agreements
+    using RecipientAgreements = std::map<AgreementKey, RecipientBlockAckAgreement>;
+
     /**
      * Handle the given in flight MPDU based on its given status. If the status is
      * ACKNOWLEDGED, the MPDU is removed from both the EDCA queue and the queue of
@@ -455,6 +519,7 @@ class BlockAckManager : public Object
      * erased from this data structure. Pushed back in retransmission queue otherwise.
      */
     OriginatorAgreements m_originatorAgreements;
+    RecipientAgreements m_recipientAgreements; //!< Recipient Block Ack agreements
 
     std::list<Bar> m_bars; ///< list of BARs
 
diff --git a/src/wifi/model/he/he-frame-exchange-manager.cc b/src/wifi/model/he/he-frame-exchange-manager.cc
index d6b4255fc..5a6698e13 100644
--- a/src/wifi/model/he/he-frame-exchange-manager.cc
+++ b/src/wifi/model/he/he-frame-exchange-manager.cc
@@ -126,8 +126,8 @@ HeFrameExchangeManager::StartFrameExchange(Ptr<QosTxop> edca, Time availableTime
     if (m_muScheduler && !edca->GetBaManager()->GetBar(false) &&
         (!(mpdu = edca->PeekNextMpdu(m_linkId)) ||
          (mpdu->GetHeader().IsQosData() && !mpdu->GetHeader().GetAddr1().IsGroup() &&
-          edca->GetBaAgreementEstablished(mpdu->GetHeader().GetAddr1(),
-                                          mpdu->GetHeader().GetQosTid()))))
+          m_mac->GetBaAgreementEstablishedAsOriginator(mpdu->GetHeader().GetAddr1(),
+                                                       mpdu->GetHeader().GetQosTid()))))
     {
         txFormat =
             m_muScheduler->NotifyAccessGranted(edca, availableTime, initialFrame, m_allowedWidth);
@@ -1301,9 +1301,11 @@ HeFrameExchangeManager::SendMultiStaBlockAck(const WifiTxParameters& txParams)
             blockAck.SetAckType(false, index);
 
             auto addressTidPair = staInfo.first;
-            auto agreementIt = m_agreements.find(addressTidPair);
-            NS_ASSERT(agreementIt != m_agreements.end());
-            agreementIt->second.FillBlockAckBitmap(&blockAck, index);
+            auto agreement =
+                GetBaManager(addressTidPair.second)
+                    ->GetAgreementAsRecipient(addressTidPair.first, addressTidPair.second);
+            NS_ASSERT(agreement);
+            agreement->get().FillBlockAckBitmap(&blockAck, index);
             NS_LOG_DEBUG("Multi-STA Block Ack: Sending Block Ack with seq="
                          << blockAck.GetStartingSequence(index) << " to=" << receiver
                          << " tid=" << +tid);
@@ -1397,7 +1399,7 @@ HeFrameExchangeManager::ReceiveBasicTrigger(const CtrlTriggerHeader& trigger,
     {
         Ptr<QosTxop> edca = m_mac->GetQosTxop(tid);
 
-        if (!edca->GetBaAgreementEstablished(hdr.GetAddr2(), tid))
+        if (!m_mac->GetBaAgreementEstablishedAsOriginator(hdr.GetAddr2(), tid))
         {
             // no Block Ack agreement established for this TID
             continue;
@@ -1498,7 +1500,7 @@ HeFrameExchangeManager::SendQosNullFramesInTbPpdu(const CtrlTriggerHeader& trigg
                txParams,
                ppduDuration))
     {
-        if (!m_mac->GetQosTxop(tid)->GetBaAgreementEstablished(hdr.GetAddr2(), tid))
+        if (!m_mac->GetBaAgreementEstablishedAsOriginator(hdr.GetAddr2(), tid))
         {
             NS_LOG_DEBUG("Skipping tid=" << +tid << " because no agreement established");
             header.SetQosTid(++tid);
@@ -1580,14 +1582,14 @@ HeFrameExchangeManager::ReceiveMpdu(Ptr<const WifiMpdu> mpdu,
             mpdu->GetPacket()->PeekHeader(blockAckReq);
             NS_ABORT_MSG_IF(blockAckReq.IsMultiTid(), "Multi-TID BlockAckReq not supported");
             uint8_t tid = blockAckReq.GetTidInfo();
-            auto agreementIt = m_agreements.find({sender, tid});
-            NS_ASSERT(agreementIt != m_agreements.end());
-            agreementIt->second.NotifyReceivedBar(blockAckReq.GetStartingSequence());
+            GetBaManager(tid)->NotifyGotBlockAckRequest(sender,
+                                                        tid,
+                                                        blockAckReq.GetStartingSequence());
 
             // Block Acknowledgment context
             acknowledgment->stationsReceivingMultiStaBa.emplace(std::make_pair(sender, tid), index);
             acknowledgment->baType.m_bitmapLen.push_back(
-                GetBlockAckType(sender, tid).m_bitmapLen.at(0));
+                m_mac->GetBaTypeAsRecipient(sender, tid).m_bitmapLen.at(0));
             uint16_t staId = txVector.GetHeMuUserInfoMap().begin()->first;
             m_muSnrTag.Set(staId, rxSignalInfo.snr);
         }
@@ -1596,9 +1598,7 @@ HeFrameExchangeManager::ReceiveMpdu(Ptr<const WifiMpdu> mpdu,
             NS_LOG_DEBUG("Received an S-MPDU in a TB PPDU from " << sender << " (" << *mpdu << ")");
 
             uint8_t tid = hdr.GetQosTid();
-            auto agreementIt = m_agreements.find({sender, tid});
-            NS_ASSERT(agreementIt != m_agreements.end());
-            agreementIt->second.NotifyReceivedMpdu(mpdu);
+            GetBaManager(tid)->NotifyGotMpdu(mpdu);
 
             // Acknowledgment context of Multi-STA Block Acks
             acknowledgment->stationsReceivingMultiStaBa.emplace(std::make_pair(sender, tid), index);
@@ -1923,21 +1923,23 @@ HeFrameExchangeManager::ReceiveMpdu(Ptr<const WifiMpdu> mpdu,
                 NS_ABORT_MSG_IF(blockAckReq.IsMultiTid(), "Multi-TID BlockAckReq not supported");
                 uint8_t tid = blockAckReq.GetTidInfo();
 
-                auto agreementIt = m_agreements.find({sender, tid});
+                auto agreement = GetBaManager(tid)->GetAgreementAsRecipient(sender, tid);
 
-                if (agreementIt == m_agreements.end())
+                if (!agreement)
                 {
                     NS_LOG_DEBUG("There's not a valid agreement for this BlockAckReq");
                     return;
                 }
 
-                agreementIt->second.NotifyReceivedBar(blockAckReq.GetStartingSequence());
+                GetBaManager(tid)->NotifyGotBlockAckRequest(sender,
+                                                            tid,
+                                                            blockAckReq.GetStartingSequence());
 
                 NS_LOG_DEBUG("Schedule Block Ack in TB PPDU");
                 Simulator::Schedule(m_phy->GetSifs(),
                                     &HeFrameExchangeManager::SendBlockAck,
                                     this,
-                                    agreementIt->second,
+                                    *agreement,
                                     hdr.GetDuration(),
                                     GetHeTbTxVector(trigger, hdr.GetAddr2()),
                                     rxSignalInfo.snr);
@@ -2020,7 +2022,7 @@ HeFrameExchangeManager::EndReceiveAmpdu(Ptr<const WifiPsdu> psdu,
                     acknowledgment->stationsReceivingMultiStaBa.emplace(std::make_pair(sender, tid),
                                                                         index + i++);
                     acknowledgment->baType.m_bitmapLen.push_back(
-                        GetBlockAckType(sender, tid).m_bitmapLen.at(0));
+                        m_mac->GetBaTypeAsRecipient(sender, tid).m_bitmapLen.at(0));
                 }
             }
             uint16_t staId = txVector.GetHeMuUserInfoMap().begin()->first;
diff --git a/src/wifi/model/he/multi-user-scheduler.cc b/src/wifi/model/he/multi-user-scheduler.cc
index a7df1c094..29810ef53 100644
--- a/src/wifi/model/he/multi-user-scheduler.cc
+++ b/src/wifi/model/he/multi-user-scheduler.cc
@@ -292,7 +292,7 @@ MultiUserScheduler::GetMaxSizeOfQosNullAmpdu(const CtrlTriggerHeader& trigger) c
         uint8_t staNTids = 0;
         for (uint8_t tid = 0; tid < 8; tid++)
         {
-            if (m_heFem->GetBaAgreementEstablished(staIt->second, tid))
+            if (m_apMac->GetBaAgreementEstablishedAsRecipient(staIt->second, tid))
             {
                 staNTids++;
             }
diff --git a/src/wifi/model/he/rr-multi-user-scheduler.cc b/src/wifi/model/he/rr-multi-user-scheduler.cc
index d7815f91c..6c56e8897 100644
--- a/src/wifi/model/he/rr-multi-user-scheduler.cc
+++ b/src/wifi/model/he/rr-multi-user-scheduler.cc
@@ -235,7 +235,7 @@ RrMultiUserScheduler::GetTxVectorForUlMu(Func canbeSolicited)
         {
             // check that a BA agreement is established with the receiver for the
             // considered TID, since ack sequences for UL MU require block ack
-            if (m_heFem->GetBaAgreementEstablished(staIt->address, tid))
+            if (m_apMac->GetBaAgreementEstablishedAsRecipient(staIt->address, tid))
             {
                 break;
             }
@@ -653,7 +653,7 @@ RrMultiUserScheduler::TrySendingDlMuPpdu()
             NS_ASSERT(ac >= primaryAc);
             // check that a BA agreement is established with the receiver for the
             // considered TID, since ack sequences for DL MU PPDUs require block ack
-            if (m_apMac->GetQosTxop(ac)->GetBaAgreementEstablished(staIt->address, tid))
+            if (m_apMac->GetBaAgreementEstablishedAsOriginator(staIt->address, tid))
             {
                 mpdu =
                     m_apMac->GetQosTxop(ac)->PeekNextMpdu(SINGLE_LINK_OP_ID, tid, staIt->address);
diff --git a/src/wifi/model/ht/ht-frame-exchange-manager.cc b/src/wifi/model/ht/ht-frame-exchange-manager.cc
index 8677b731b..eaddac602 100644
--- a/src/wifi/model/ht/ht-frame-exchange-manager.cc
+++ b/src/wifi/model/ht/ht-frame-exchange-manager.cc
@@ -67,7 +67,6 @@ void
 HtFrameExchangeManager::DoDispose()
 {
     NS_LOG_FUNCTION(this);
-    m_agreements.clear();
     m_pendingAddBaResp.clear();
     m_msduAggregator = nullptr;
     m_mpduAggregator = nullptr;
@@ -178,7 +177,7 @@ HtFrameExchangeManager::SendAddBaRequest(Mac48Address dest,
     // set the starting sequence number for the BA agreement
     reqHdr.SetStartingSequence(startingSeq);
 
-    GetBaManager(tid)->CreateOriginatorAgreement(&reqHdr, dest);
+    GetBaManager(tid)->CreateOriginatorAgreement(reqHdr, dest);
 
     packet->AddHeader(reqHdr);
     packet->AddHeader(actionHdr);
@@ -228,7 +227,8 @@ HtFrameExchangeManager::SendAddBaResponse(const MgtAddBaRequestHeader* reqHdr,
     {
         respHdr.SetDelayedBlockAck();
     }
-    respHdr.SetTid(reqHdr->GetTid());
+    auto tid = reqHdr->GetTid();
+    respHdr.SetTid(tid);
 
     respHdr.SetBufferSize(GetSupportedBaBufferSize());
     respHdr.SetTimeout(reqHdr->GetTimeout());
@@ -242,7 +242,28 @@ HtFrameExchangeManager::SendAddBaResponse(const MgtAddBaRequestHeader* reqHdr,
     packet->AddHeader(respHdr);
     packet->AddHeader(actionHdr);
 
-    CreateBlockAckAgreement(&respHdr, originator, reqHdr->GetStartingSequence());
+    bool htSupported = GetWifiRemoteStationManager()->GetHtSupported() &&
+                       GetWifiRemoteStationManager()->GetHtSupported(originator);
+    GetBaManager(tid)->CreateRecipientAgreement(respHdr,
+                                                originator,
+                                                reqHdr->GetStartingSequence(),
+                                                htSupported,
+                                                m_rxMiddle);
+
+    auto agreement = GetBaManager(tid)->GetAgreementAsRecipient(originator, tid);
+    NS_ASSERT(agreement);
+    if (respHdr.GetTimeout() != 0)
+    {
+        Time timeout = MicroSeconds(1024 * agreement->get().GetTimeout());
+
+        agreement->get().m_inactivityEvent =
+            Simulator::Schedule(timeout,
+                                &HtFrameExchangeManager::SendDelbaFrame,
+                                this,
+                                originator,
+                                tid,
+                                false);
+    }
 
     auto mpdu = Create<WifiMpdu>(packet, hdr);
 
@@ -257,7 +278,7 @@ HtFrameExchangeManager::SendAddBaResponse(const MgtAddBaRequestHeader* reqHdr,
      */
 
     // remove any pending ADDBA_RESPONSE frame
-    AgreementKey key(originator, reqHdr->GetTid());
+    AgreementKey key(originator, tid);
     if (auto it = m_pendingAddBaResp.find(key); it != m_pendingAddBaResp.end())
     {
         NS_ASSERT_MSG(it->second, "The pointer to the pending ADDBA_RESPONSE cannot be null");
@@ -271,7 +292,7 @@ HtFrameExchangeManager::SendAddBaResponse(const MgtAddBaRequestHeader* reqHdr,
     // bung it into the queue corresponding to the TID for which we are
     // establishing an agreement, and push it to the head.
     //  Wifi MAC queue scheduler is expected to prioritize management frames
-    m_mac->GetQosTxop(reqHdr->GetTid())->Queue(mpdu);
+    m_mac->GetQosTxop(tid)->Queue(mpdu);
 }
 
 uint16_t
@@ -302,7 +323,7 @@ HtFrameExchangeManager::SendDelbaFrame(Mac48Address addr, uint8_t tid, bool byOr
     else
     {
         delbaHdr.SetByRecipient();
-        DestroyBlockAckAgreement(addr, tid);
+        GetBaManager(tid)->DestroyRecipientAgreement(addr, tid);
     }
 
     WifiActionHeader actionHdr;
@@ -317,62 +338,6 @@ HtFrameExchangeManager::SendDelbaFrame(Mac48Address addr, uint8_t tid, bool byOr
     m_mac->GetQosTxop(tid)->GetWifiMacQueue()->Enqueue(Create<WifiMpdu>(packet, hdr));
 }
 
-void
-HtFrameExchangeManager::CreateBlockAckAgreement(const MgtAddBaResponseHeader* respHdr,
-                                                Mac48Address originator,
-                                                uint16_t startingSeq)
-{
-    NS_LOG_FUNCTION(this << *respHdr << originator << startingSeq);
-    uint8_t tid = respHdr->GetTid();
-
-    RecipientBlockAckAgreement agreement(
-        originator,
-        respHdr->IsAmsduSupported(),
-        tid,
-        respHdr->GetBufferSize(),
-        respHdr->GetTimeout(),
-        startingSeq,
-        GetWifiRemoteStationManager()->GetHtSupported() &&
-            GetWifiRemoteStationManager()->GetHtSupported(originator));
-    agreement.SetMacRxMiddle(m_rxMiddle);
-    if (respHdr->IsImmediateBlockAck())
-    {
-        agreement.SetImmediateBlockAck();
-    }
-    else
-    {
-        agreement.SetDelayedBlockAck();
-    }
-
-    if (respHdr->GetTimeout() != 0)
-    {
-        Time timeout = MicroSeconds(1024 * agreement.GetTimeout());
-
-        agreement.m_inactivityEvent = Simulator::Schedule(timeout,
-                                                          &HtFrameExchangeManager::SendDelbaFrame,
-                                                          this,
-                                                          originator,
-                                                          tid,
-                                                          false);
-    }
-
-    m_agreements.insert_or_assign({originator, tid}, agreement);
-}
-
-void
-HtFrameExchangeManager::DestroyBlockAckAgreement(Mac48Address originator, uint8_t tid)
-{
-    NS_LOG_FUNCTION(this << originator << +tid);
-
-    auto agreementIt = m_agreements.find({originator, tid});
-    if (agreementIt != m_agreements.end())
-    {
-        // forward up the buffered MPDUs before destroying the agreement
-        agreementIt->second.Flush();
-        m_agreements.erase(agreementIt);
-    }
-}
-
 bool
 HtFrameExchangeManager::StartFrameExchange(Ptr<QosTxop> edca, Time availableTime, bool initialFrame)
 {
@@ -588,7 +553,7 @@ HtFrameExchangeManager::NotifyReceivedNormalAck(Ptr<WifiMpdu> mpdu)
         uint8_t tid = mpdu->GetHeader().GetQosTid();
         Ptr<QosTxop> edca = m_mac->GetQosTxop(tid);
 
-        if (edca->GetBaAgreementEstablished(mpdu->GetHeader().GetAddr1(), tid))
+        if (m_mac->GetBaAgreementEstablishedAsOriginator(mpdu->GetHeader().GetAddr1(), tid))
         {
             // notify the BA manager that the MPDU was acknowledged
             edca->GetBaManager()->NotifyGotAck(m_linkId, mpdu);
@@ -605,14 +570,15 @@ HtFrameExchangeManager::NotifyReceivedNormalAck(Ptr<WifiMpdu> mpdu)
             {
                 MgtDelBaHeader delBa;
                 p->PeekHeader(delBa);
+                auto tid = delBa.GetTid();
                 if (delBa.IsByOriginator())
                 {
-                    GetBaManager(delBa.GetTid())
-                        ->DestroyOriginatorAgreement(mpdu->GetHeader().GetAddr1(), delBa.GetTid());
+                    GetBaManager(tid)->DestroyOriginatorAgreement(mpdu->GetHeader().GetAddr1(),
+                                                                  tid);
                 }
                 else
                 {
-                    DestroyBlockAckAgreement(mpdu->GetHeader().GetAddr1(), delBa.GetTid());
+                    GetBaManager(tid)->DestroyRecipientAgreement(mpdu->GetHeader().GetAddr1(), tid);
                 }
             }
             else if (actionHdr.GetAction().blockAck == WifiActionHeader::BLOCK_ACK_ADDBA_REQUEST)
@@ -632,12 +598,12 @@ HtFrameExchangeManager::NotifyReceivedNormalAck(Ptr<WifiMpdu> mpdu)
                 // A recipient Block Ack agreement must exist
                 MgtAddBaResponseHeader addBa;
                 p->PeekHeader(addBa);
-                AgreementKey key(mpdu->GetHeader().GetAddr1(), addBa.GetTid());
-                auto agreementIt = m_agreements.find(key);
-                NS_ASSERT_MSG(agreementIt != m_agreements.end(),
-                              "Recipient BA agreement {" << key.first << ", " << +key.second
+                auto originator = mpdu->GetHeader().GetAddr1();
+                auto tid = addBa.GetTid();
+                NS_ASSERT_MSG(GetBaManager(tid)->GetAgreementAsRecipient(originator, tid),
+                              "Recipient BA agreement {" << originator << ", " << +tid
                                                          << "} not found");
-                m_pendingAddBaResp.erase(key);
+                m_pendingAddBaResp.erase({originator, tid});
             }
         }
     }
@@ -718,7 +684,7 @@ HtFrameExchangeManager::RetransmitMpduAfterMissedAck(Ptr<WifiMpdu> mpdu) const
         uint8_t tid = mpdu->GetHeader().GetQosTid();
         Ptr<QosTxop> edca = m_mac->GetQosTxop(tid);
 
-        if (edca->GetBaAgreementEstablished(mpdu->GetHeader().GetAddr1(), tid))
+        if (m_mac->GetBaAgreementEstablishedAsOriginator(mpdu->GetHeader().GetAddr1(), tid))
         {
             // notify the BA manager that the MPDU was not acknowledged
             edca->GetBaManager()->NotifyMissedAck(m_linkId, mpdu);
@@ -740,7 +706,7 @@ HtFrameExchangeManager::ReleaseSequenceNumber(Ptr<WifiMpdu> mpdu) const
         uint8_t tid = hdr.GetQosTid();
         Ptr<QosTxop> edca = m_mac->GetQosTxop(tid);
 
-        if (edca->GetBaAgreementEstablished(hdr.GetAddr1(), tid) && !hdr.IsRetry())
+        if (m_mac->GetBaAgreementEstablishedAsOriginator(hdr.GetAddr1(), tid) && !hdr.IsRetry())
         {
             // The MPDU has never been transmitted, so we can make its sequence
             // number available again if it is lower than the sequence number
@@ -1269,7 +1235,7 @@ HtFrameExchangeManager::MissedBlockAck(Ptr<WifiPsdu> psdu,
             // if a BA agreement exists, we can get here if there is no outstanding
             // MPDU whose lifetime has not expired yet.
             GetWifiRemoteStationManager()->ReportFinalDataFailed(*psdu->begin());
-            if (edca->GetBaAgreementEstablished(recipient, tid))
+            if (m_mac->GetBaAgreementEstablishedAsOriginator(recipient, tid))
             {
                 // schedule a BlockAckRequest with skipIfNoDataQueued set to true, so that the
                 // BlockAckRequest is only sent if there are data frames queued for this recipient.
@@ -1347,20 +1313,6 @@ HtFrameExchangeManager::SendBlockAck(const RecipientBlockAckAgreement& agreement
     ForwardPsduDown(psdu, blockAckTxVector);
 }
 
-bool
-HtFrameExchangeManager::GetBaAgreementEstablished(Mac48Address originator, uint8_t tid) const
-{
-    return (m_agreements.find({originator, tid}) != m_agreements.end());
-}
-
-BlockAckType
-HtFrameExchangeManager::GetBlockAckType(Mac48Address originator, uint8_t tid) const
-{
-    auto it = m_agreements.find({originator, tid});
-    NS_ABORT_MSG_IF(it == m_agreements.end(), "No established Block Ack agreement");
-    return it->second.GetBlockAckType();
-}
-
 void
 HtFrameExchangeManager::ReceiveMpdu(Ptr<const WifiMpdu> mpdu,
                                     RxSignalInfo rxSignalInfo,
@@ -1441,22 +1393,24 @@ HtFrameExchangeManager::ReceiveMpdu(Ptr<const WifiMpdu> mpdu,
             NS_ABORT_MSG_IF(blockAckReq.IsMultiTid(), "Multi-TID BlockAckReq not supported");
             uint8_t tid = blockAckReq.GetTidInfo();
 
-            auto agreementIt = m_agreements.find({sender, tid});
+            auto agreement = GetBaManager(tid)->GetAgreementAsRecipient(sender, tid);
 
-            if (agreementIt == m_agreements.end())
+            if (!agreement)
             {
                 NS_LOG_DEBUG("There's not a valid agreement for this BlockAckReq");
                 return;
             }
 
-            agreementIt->second.NotifyReceivedBar(blockAckReq.GetStartingSequence());
+            GetBaManager(tid)->NotifyGotBlockAckRequest(sender,
+                                                        tid,
+                                                        blockAckReq.GetStartingSequence());
 
             NS_LOG_DEBUG("Schedule Block Ack");
             Simulator::Schedule(
                 m_phy->GetSifs(),
                 &HtFrameExchangeManager::SendBlockAck,
                 this,
-                agreementIt->second,
+                *agreement,
                 hdr.GetDuration(),
                 GetWifiRemoteStationManager()->GetBlockAckTxVector(sender, txVector),
                 rxSnr);
@@ -1473,13 +1427,12 @@ HtFrameExchangeManager::ReceiveMpdu(Ptr<const WifiMpdu> mpdu,
     {
         uint8_t tid = hdr.GetQosTid();
 
-        auto agreementIt = m_agreements.find({hdr.GetAddr2(), tid});
-        if (agreementIt != m_agreements.end())
+        if (m_mac->GetBaAgreementEstablishedAsRecipient(hdr.GetAddr2(), tid))
         {
             // a Block Ack agreement has been established
             NS_LOG_DEBUG("Received from=" << hdr.GetAddr2() << " (" << *mpdu << ")");
 
-            agreementIt->second.NotifyReceivedMpdu(mpdu);
+            GetBaManager(tid)->NotifyGotMpdu(mpdu);
 
             if (!inAmpdu && hdr.GetQosAckPolicy() == WifiMacHeader::NORMAL_ACK)
             {
@@ -1519,14 +1472,14 @@ HtFrameExchangeManager::EndReceiveAmpdu(Ptr<const WifiPsdu> psdu,
         {
             // Normal Ack or Implicit Block Ack Request
             NS_LOG_DEBUG("Schedule Block Ack");
-            auto agreementIt = m_agreements.find({psdu->GetAddr2(), tid});
-            NS_ASSERT(agreementIt != m_agreements.end());
+            auto agreement = GetBaManager(tid)->GetAgreementAsRecipient(psdu->GetAddr2(), tid);
+            NS_ASSERT(agreement);
 
             Simulator::Schedule(
                 m_phy->GetSifs(),
                 &HtFrameExchangeManager::SendBlockAck,
                 this,
-                agreementIt->second,
+                *agreement,
                 psdu->GetDuration(),
                 GetWifiRemoteStationManager()->GetBlockAckTxVector(psdu->GetAddr2(), txVector),
                 rxSignalInfo.snr);
diff --git a/src/wifi/model/ht/ht-frame-exchange-manager.h b/src/wifi/model/ht/ht-frame-exchange-manager.h
index cc6421afd..56cc3a23d 100644
--- a/src/wifi/model/ht/ht-frame-exchange-manager.h
+++ b/src/wifi/model/ht/ht-frame-exchange-manager.h
@@ -140,31 +140,6 @@ class HtFrameExchangeManager : public QosFrameExchangeManager
                                                const WifiTxParameters& txParams,
                                                Time ppduDurationLimit) const;
 
-    /**
-     * \param respHdr Add block ack response from originator (action
-     *        frame).
-     * \param originator Address of peer station involved in block ack
-     *        mechanism.
-     * \param startingSeq Sequence number of the first MPDU of all
-     *        packets for which block ack was negotiated.
-     *
-     * This function is typically invoked only by ns3::WifiMac
-     * when the STA (which may be non-AP in ESS, or in an IBSS) has
-     * received an ADDBA Request frame and is transmitting an ADDBA
-     * Response frame. At this point the frame exchange manager must
-     * allocate buffers to collect all correctly received packets belonging
-     * to the category for which block ack was negotiated.
-     */
-    void CreateBlockAckAgreement(const MgtAddBaResponseHeader* respHdr,
-                                 Mac48Address originator,
-                                 uint16_t startingSeq);
-    /**
-     * Destroy a Block Ack agreement.
-     *
-     * \param originator the originator MAC address
-     * \param tid the TID associated with the Block Ack agreement
-     */
-    void DestroyBlockAckAgreement(Mac48Address originator, uint8_t tid);
     /**
      * This method can be called to accept a received ADDBA Request. An
      * ADDBA Response will be constructed and queued for transmission.
@@ -181,26 +156,6 @@ class HtFrameExchangeManager : public QosFrameExchangeManager
      */
     virtual uint16_t GetSupportedBaBufferSize() const;
 
-    /**
-     * Return true if a Block Ack agreement has been established with the given
-     * originator for the given TID.
-     *
-     * \param originator the MAC address of the given originator
-     * \param tid the Traffic ID
-     * \return true if a Block Ack agreement has been established with the given
-     *         originator for the given TID
-     */
-    bool GetBaAgreementEstablished(Mac48Address originator, uint8_t tid) const;
-
-    /**
-     * Get the type of BlockAck frames sent to the given originator.
-     *
-     * \param originator the MAC address of the given originator
-     * \param tid the Traffic ID
-     * \return the type of BlockAck frames sent to the given originator
-     */
-    BlockAckType GetBlockAckType(Mac48Address originator, uint8_t tid) const;
-
     /**
      * Sends DELBA frame to cancel a block ack agreement with STA
      * addressed by <i>addr</i> for TID <i>tid</i>.
@@ -387,12 +342,8 @@ class HtFrameExchangeManager : public QosFrameExchangeManager
     /// agreement key typedef (MAC address and TID)
     typedef std::pair<Mac48Address, uint8_t> AgreementKey;
 
-    /// typedef for map of recipient Block Ack agreements
-    using RecipientBlockAckAgreementMap = std::map<AgreementKey, RecipientBlockAckAgreement>;
-
-    RecipientBlockAckAgreementMap m_agreements; //!< Block Ack agreements
-    Ptr<MsduAggregator> m_msduAggregator;       //!< A-MSDU aggregator
-    Ptr<MpduAggregator> m_mpduAggregator;       //!< A-MPDU aggregator
+    Ptr<MsduAggregator> m_msduAggregator; //!< A-MSDU aggregator
+    Ptr<MpduAggregator> m_mpduAggregator; //!< A-MPDU aggregator
 
     /// pending ADDBA_RESPONSE frames indexed by agreement key
     std::map<AgreementKey, Ptr<WifiMpdu>> m_pendingAddBaResp;
diff --git a/src/wifi/model/mpdu-aggregator.cc b/src/wifi/model/mpdu-aggregator.cc
index c9bbf6440..f31ccf055 100644
--- a/src/wifi/model/mpdu-aggregator.cc
+++ b/src/wifi/model/mpdu-aggregator.cc
@@ -210,7 +210,7 @@ MpduAggregator::GetNextAmpdu(Ptr<WifiMpdu> mpdu,
     NS_ASSERT(qosTxop);
 
     // Have to make sure that the block ack agreement is established and A-MPDU is enabled
-    if (qosTxop->GetBaAgreementEstablished(recipient, tid) &&
+    if (m_mac->GetBaAgreementEstablishedAsOriginator(recipient, tid) &&
         GetMaxAmpduSize(recipient, tid, txParams.m_txVector.GetModulationClass()) > 0)
     {
         /* here is performed MPDU aggregation */
diff --git a/src/wifi/model/qos-txop.cc b/src/wifi/model/qos-txop.cc
index 19ef5a4fe..c95352f65 100644
--- a/src/wifi/model/qos-txop.cc
+++ b/src/wifi/model/qos-txop.cc
@@ -263,13 +263,6 @@ QosTxop::GetBaManager()
     return m_baManager;
 }
 
-bool
-QosTxop::GetBaAgreementEstablished(Mac48Address address, uint8_t tid) const
-{
-    auto agreement = m_baManager->GetAgreementAsOriginator(address, tid);
-    return agreement && agreement->get().IsEstablished();
-}
-
 uint16_t
 QosTxop::GetBaBufferSize(Mac48Address address, uint8_t tid) const
 {
@@ -356,7 +349,7 @@ QosTxop::IsQosOldPacket(Ptr<const WifiMpdu> mpdu)
     Mac48Address recipient = mpdu->GetHeader().GetAddr1();
     uint8_t tid = mpdu->GetHeader().GetQosTid();
 
-    if (!GetBaAgreementEstablished(recipient, tid))
+    if (!m_mac->GetBaAgreementEstablishedAsOriginator(recipient, tid))
     {
         return false;
     }
@@ -417,8 +410,8 @@ QosTxop::PeekNextMpdu(uint8_t linkId, uint8_t tid, Mac48Address recipient, Ptr<W
 
             // if no BA agreement, we cannot have multiple MPDUs in-flight
             if (item->GetHeader().IsQosData() &&
-                !GetBaAgreementEstablished(item->GetHeader().GetAddr1(),
-                                           item->GetHeader().GetQosTid()))
+                !m_mac->GetBaAgreementEstablishedAsOriginator(item->GetHeader().GetAddr1(),
+                                                              item->GetHeader().GetQosTid()))
             {
                 NS_LOG_DEBUG("No BA agreement and an MPDU is already in-flight");
                 return nullptr;
@@ -456,7 +449,7 @@ QosTxop::PeekNextMpdu(uint8_t linkId, uint8_t tid, Mac48Address recipient, Ptr<W
         Mac48Address recipient = hdr.GetAddr1();
         uint8_t tid = hdr.GetQosTid();
 
-        if (GetBaAgreementEstablished(recipient, tid) &&
+        if (m_mac->GetBaAgreementEstablishedAsOriginator(recipient, tid) &&
             !IsInWindow(sequence,
                         GetBaStartingSequence(recipient, tid),
                         GetBaBufferSize(recipient, tid)))
@@ -511,7 +504,7 @@ QosTxop::GetNextMpdu(uint8_t linkId,
         // we should not be asked to dequeue an MPDU that is beyond the transmit window.
         // Note that PeekNextMpdu() temporarily assigns the next available sequence number
         // to the peeked frame
-        NS_ASSERT(!GetBaAgreementEstablished(recipient, tid) ||
+        NS_ASSERT(!m_mac->GetBaAgreementEstablishedAsOriginator(recipient, tid) ||
                   IsInWindow(peekedItem->GetHeader().GetSequenceNumber(),
                              GetBaStartingSequence(recipient, tid),
                              GetBaBufferSize(recipient, tid)));
@@ -615,11 +608,11 @@ QosTxop::GetRemainingTxop(uint8_t linkId) const
 }
 
 void
-QosTxop::GotAddBaResponse(const MgtAddBaResponseHeader* respHdr, Mac48Address recipient)
+QosTxop::GotAddBaResponse(const MgtAddBaResponseHeader& respHdr, Mac48Address recipient)
 {
     NS_LOG_FUNCTION(this << respHdr << recipient);
-    uint8_t tid = respHdr->GetTid();
-    if (respHdr->GetStatusCode().IsSuccess())
+    uint8_t tid = respHdr.GetTid();
+    if (respHdr.GetStatusCode().IsSuccess())
     {
         NS_LOG_DEBUG("block ack agreement established with " << recipient << " tid " << +tid);
         // A (destination, TID) pair is "blocked" (i.e., no more packets are sent)
@@ -676,7 +669,8 @@ QosTxop::CompleteMpduTx(Ptr<WifiMpdu> mpdu)
 {
     NS_ASSERT(mpdu->GetHeader().IsQosData());
     // If there is an established BA agreement, store the packet in the queue of outstanding packets
-    if (GetBaAgreementEstablished(mpdu->GetHeader().GetAddr1(), mpdu->GetHeader().GetQosTid()))
+    if (m_mac->GetBaAgreementEstablishedAsOriginator(mpdu->GetHeader().GetAddr1(),
+                                                     mpdu->GetHeader().GetQosTid()))
     {
         m_baManager->StorePacket(mpdu);
     }
diff --git a/src/wifi/model/qos-txop.h b/src/wifi/model/qos-txop.h
index 970484167..2f5a686af 100644
--- a/src/wifi/model/qos-txop.h
+++ b/src/wifi/model/qos-txop.h
@@ -113,16 +113,6 @@ class QosTxop : public Txop
      * \returns the Block Ack Manager
      */
     Ptr<BlockAckManager> GetBaManager();
-    /**
-     * \param address recipient address of the peer station
-     * \param tid traffic ID.
-     *
-     * \return true if a block ack agreement is established, false otherwise.
-     *
-     * Checks if a block ack agreement is established with station addressed by
-     * <i>recipient</i> for TID <i>tid</i>.
-     */
-    bool GetBaAgreementEstablished(Mac48Address address, uint8_t tid) const;
     /**
      * \param address recipient address of the peer station
      * \param tid traffic ID.
@@ -171,7 +161,7 @@ class QosTxop : public Txop
      * \param respHdr ADDBA response header.
      * \param recipient address of the recipient.
      */
-    void GotAddBaResponse(const MgtAddBaResponseHeader* respHdr, Mac48Address recipient);
+    void GotAddBaResponse(const MgtAddBaResponseHeader& respHdr, Mac48Address recipient);
     /**
      * Event handler when a DELBA frame is received.
      *
diff --git a/src/wifi/model/wifi-default-ack-manager.cc b/src/wifi/model/wifi-default-ack-manager.cc
index 90a28a122..165672df7 100644
--- a/src/wifi/model/wifi-default-ack-manager.cc
+++ b/src/wifi/model/wifi-default-ack-manager.cc
@@ -101,7 +101,7 @@ WifiDefaultAckManager::GetMaxDistFromStartingSeq(Ptr<const WifiMpdu> mpdu,
 
     uint8_t tid = hdr.GetQosTid();
     Ptr<QosTxop> edca = m_mac->GetQosTxop(tid);
-    NS_ABORT_MSG_IF(!edca->GetBaAgreementEstablished(receiver, tid),
+    NS_ABORT_MSG_IF(!m_mac->GetBaAgreementEstablishedAsOriginator(receiver, tid),
                     "An established Block Ack agreement is required");
 
     uint16_t startingSeq = edca->GetBaStartingSequence(receiver, tid);
@@ -257,8 +257,8 @@ WifiDefaultAckManager::TryAddMpdu(Ptr<const WifiMpdu> mpdu, const WifiTxParamete
         return std::unique_ptr<WifiAcknowledgment>(acknowledgment);
     }
 
-    if ((!hdr.IsQosData() || !m_mac->GetQosTxop(hdr.GetQosTid())
-                                  ->GetBaAgreementEstablished(receiver, hdr.GetQosTid())) &&
+    if ((!hdr.IsQosData() ||
+         !m_mac->GetBaAgreementEstablishedAsOriginator(receiver, hdr.GetQosTid())) &&
         !hdr.IsBlockAckReq())
     {
         NS_LOG_DEBUG(
@@ -675,7 +675,7 @@ WifiDefaultAckManager::TryUlMuTransmission(Ptr<const WifiMpdu> mpdu,
 
             // find a TID for which a BA agreement exists with the given originator
             uint8_t tid = 0;
-            while (tid < 8 && !heFem->GetBaAgreementEstablished(staAddress, tid))
+            while (tid < 8 && !m_mac->GetBaAgreementEstablishedAsRecipient(staAddress, tid))
             {
                 tid++;
             }
@@ -689,7 +689,7 @@ WifiDefaultAckManager::TryUlMuTransmission(Ptr<const WifiMpdu> mpdu,
             // we assume the Block Acknowledgment context is used for the multi-STA BlockAck frame
             // (since it requires the longest TX time due to the presence of a bitmap)
             acknowledgment->baType.m_bitmapLen.push_back(
-                heFem->GetBlockAckType(staAddress, tid).m_bitmapLen.at(0));
+                m_mac->GetBaTypeAsRecipient(staAddress, tid).m_bitmapLen.at(0));
         }
 
         uint16_t staId = trigger.begin()->GetAid12();
diff --git a/src/wifi/model/wifi-mac.cc b/src/wifi/model/wifi-mac.cc
index 00049afc0..4882af2e7 100644
--- a/src/wifi/model/wifi-mac.cc
+++ b/src/wifi/model/wifi-mac.cc
@@ -1175,7 +1175,7 @@ WifiMac::Receive(Ptr<const WifiMpdu> mpdu, uint8_t linkId)
                 // seems a waste given the level of the current model)
                 // and act by locally establishing the agreement on
                 // the appropriate queue.
-                GetQosTxop(respHdr.GetTid())->GotAddBaResponse(&respHdr, from);
+                GetQosTxop(respHdr.GetTid())->GotAddBaResponse(respHdr, from);
                 auto htFem = DynamicCast<HtFrameExchangeManager>(link.feManager);
                 if (htFem)
                 {
@@ -1195,14 +1195,11 @@ WifiMac::Receive(Ptr<const WifiMpdu> mpdu, uint8_t linkId)
                 {
                     // This DELBA frame was sent by the originator, so
                     // this means that an ingoing established
-                    // agreement exists in HtFrameExchangeManager and we need to
+                    // agreement exists in BlockAckManager and we need to
                     // destroy it.
-                    NS_ASSERT(link.feManager);
-                    auto htFem = DynamicCast<HtFrameExchangeManager>(link.feManager);
-                    if (htFem)
-                    {
-                        htFem->DestroyBlockAckAgreement(from, delBaHdr.GetTid());
-                    }
+                    GetQosTxop(delBaHdr.GetTid())
+                        ->GetBaManager()
+                        ->DestroyRecipientAgreement(from, delBaHdr.GetTid());
                 }
                 else
                 {
@@ -1238,10 +1235,45 @@ WifiMac::DeaggregateAmsduAndForward(Ptr<const WifiMpdu> mpdu)
     }
 }
 
+std::optional<Mac48Address>
+WifiMac::GetMldAddress(const Mac48Address& remoteAddr) const
+{
+    for (std::size_t linkId = 0; linkId < m_links.size(); linkId++)
+    {
+        if (auto mldAddress = m_links[linkId]->stationManager->GetMldAddress(remoteAddr))
+        {
+            return *mldAddress;
+        }
+    }
+    return std::nullopt;
+}
+
+WifiMac::OriginatorAgreementOptConstRef
+WifiMac::GetBaAgreementEstablishedAsOriginator(Mac48Address recipient, uint8_t tid) const
+{
+    // BA agreements are indexed by the MLD address if ML setup was performed
+    recipient = GetMldAddress(recipient).value_or(recipient);
+
+    auto agreement = GetQosTxop(tid)->GetBaManager()->GetAgreementAsOriginator(recipient, tid);
+    if (!agreement || !agreement->get().IsEstablished())
+    {
+        return std::nullopt;
+    }
+    return agreement;
+}
+
+WifiMac::RecipientAgreementOptConstRef
+WifiMac::GetBaAgreementEstablishedAsRecipient(Mac48Address originator, uint8_t tid) const
+{
+    // BA agreements are indexed by the MLD address if ML setup was performed
+    originator = GetMldAddress(originator).value_or(originator);
+    return GetQosTxop(tid)->GetBaManager()->GetAgreementAsRecipient(originator, tid);
+}
+
 BlockAckType
 WifiMac::GetBaTypeAsOriginator(const Mac48Address& recipient, uint8_t tid) const
 {
-    auto agreement = GetQosTxop(tid)->GetBaManager()->GetAgreementAsOriginator(recipient, tid);
+    auto agreement = GetBaAgreementEstablishedAsOriginator(recipient, tid);
     NS_ABORT_MSG_IF(!agreement,
                     "No existing Block Ack agreement with " << recipient << " TID: " << +tid);
     return agreement->get().GetBlockAckType();
@@ -1250,12 +1282,30 @@ WifiMac::GetBaTypeAsOriginator(const Mac48Address& recipient, uint8_t tid) const
 BlockAckReqType
 WifiMac::GetBarTypeAsOriginator(const Mac48Address& recipient, uint8_t tid) const
 {
-    auto agreement = GetQosTxop(tid)->GetBaManager()->GetAgreementAsOriginator(recipient, tid);
+    auto agreement = GetBaAgreementEstablishedAsOriginator(recipient, tid);
     NS_ABORT_MSG_IF(!agreement,
                     "No existing Block Ack agreement with " << recipient << " TID: " << +tid);
     return agreement->get().GetBlockAckReqType();
 }
 
+BlockAckType
+WifiMac::GetBaTypeAsRecipient(Mac48Address originator, uint8_t tid) const
+{
+    auto agreement = GetBaAgreementEstablishedAsRecipient(originator, tid);
+    NS_ABORT_MSG_IF(!agreement,
+                    "No existing Block Ack agreement with " << originator << " TID: " << +tid);
+    return agreement->get().GetBlockAckType();
+}
+
+BlockAckReqType
+WifiMac::GetBarTypeAsRecipient(Mac48Address originator, uint8_t tid) const
+{
+    auto agreement = GetBaAgreementEstablishedAsRecipient(originator, tid);
+    NS_ABORT_MSG_IF(!agreement,
+                    "No existing Block Ack agreement with " << originator << " TID: " << +tid);
+    return agreement->get().GetBlockAckReqType();
+}
+
 Ptr<HtConfiguration>
 WifiMac::GetHtConfiguration() const
 {
diff --git a/src/wifi/model/wifi-mac.h b/src/wifi/model/wifi-mac.h
index d42787f6f..496fa27ee 100644
--- a/src/wifi/model/wifi-mac.h
+++ b/src/wifi/model/wifi-mac.h
@@ -51,6 +51,8 @@ class FrameExchangeManager;
 class ChannelAccessManager;
 class ExtendedCapabilities;
 class WifiMacQueueScheduler;
+class OriginatorBlockAckAgreement;
+class RecipientBlockAckAgreement;
 
 /**
  * Enumeration for type of station
@@ -147,6 +149,13 @@ class WifiMac : public Object
      */
     virtual std::optional<uint8_t> GetLinkIdByAddress(const Mac48Address& address) const;
 
+    /**
+     * \param remoteAddr the (MLD or link) address of a remote device
+     * \return the MLD address of the remote device having the given (MLD or link) address, if
+     *         the remote device is an MLD.
+     */
+    std::optional<Mac48Address> GetMldAddress(const Mac48Address& remoteAddr) const;
+
     /**
      * Accessor for the Txop object
      *
@@ -529,6 +538,36 @@ class WifiMac : public Object
      */
     uint16_t GetMaxAmsduSize(AcIndex ac) const;
 
+    /// optional const reference to OriginatorBlockAckAgreement
+    using OriginatorAgreementOptConstRef =
+        std::optional<std::reference_wrapper<const OriginatorBlockAckAgreement>>;
+    /// optional const reference to RecipientBlockAckAgreement
+    using RecipientAgreementOptConstRef =
+        std::optional<std::reference_wrapper<const RecipientBlockAckAgreement>>;
+
+    /**
+     * \param recipient (link or device) MAC address of the recipient
+     * \param tid traffic ID.
+     *
+     * \return the originator block ack agreement, if one has been established
+     *
+     * Checks if an originator block ack agreement is established with station addressed by
+     * <i>recipient</i> for TID <i>tid</i>.
+     */
+    OriginatorAgreementOptConstRef GetBaAgreementEstablishedAsOriginator(Mac48Address recipient,
+                                                                         uint8_t tid) const;
+    /**
+     * \param originator (link or device) MAC address of the originator
+     * \param tid traffic ID.
+     *
+     * \return the recipient block ack agreement, if one has been established
+     *
+     * Checks if a recipient block ack agreement is established with station addressed by
+     * <i>originator</i> for TID <i>tid</i>.
+     */
+    RecipientAgreementOptConstRef GetBaAgreementEstablishedAsRecipient(Mac48Address originator,
+                                                                       uint8_t tid) const;
+
     /**
      * \param recipient MAC address
      * \param tid traffic ID
@@ -547,6 +586,24 @@ class WifiMac : public Object
      * This function returns the type of Block Ack Requests sent to the recipient.
      */
     BlockAckReqType GetBarTypeAsOriginator(const Mac48Address& recipient, uint8_t tid) const;
+    /**
+     * \param originator MAC address of originator
+     * \param tid traffic ID
+     *
+     * \return the type of Block Acks sent to the originator
+     *
+     * This function returns the type of Block Acks sent to the originator.
+     */
+    BlockAckType GetBaTypeAsRecipient(Mac48Address originator, uint8_t tid) const;
+    /**
+     * \param originator MAC address of originator
+     * \param tid traffic ID
+     *
+     * \return the type of Block Ack Requests sent by the originator
+     *
+     * This function returns the type of Block Ack Requests sent by the originator.
+     */
+    BlockAckReqType GetBarTypeAsRecipient(Mac48Address originator, uint8_t tid) const;
 
   protected:
     void DoInitialize() override;
diff --git a/src/wifi/test/wifi-aggregation-test.cc b/src/wifi/test/wifi-aggregation-test.cc
index ee253dee2..40e93ca97 100644
--- a/src/wifi/test/wifi-aggregation-test.cc
+++ b/src/wifi/test/wifi-aggregation-test.cc
@@ -173,7 +173,7 @@ AmpduAggregationTest::DoRun()
     reqHdr.SetBufferSize(64);
     reqHdr.SetTimeout(0);
     reqHdr.SetStartingSequence(0);
-    m_mac->GetBEQueue()->GetBaManager()->CreateOriginatorAgreement(&reqHdr, hdr.GetAddr1());
+    m_mac->GetBEQueue()->GetBaManager()->CreateOriginatorAgreement(reqHdr, hdr.GetAddr1());
 
     MgtAddBaResponseHeader respHdr;
     StatusCode code;
@@ -184,7 +184,7 @@ AmpduAggregationTest::DoRun()
     respHdr.SetTid(reqHdr.GetTid());
     respHdr.SetBufferSize(64);
     respHdr.SetTimeout(reqHdr.GetTimeout());
-    m_mac->GetBEQueue()->GetBaManager()->UpdateOriginatorAgreement(&respHdr, hdr.GetAddr1(), 0);
+    m_mac->GetBEQueue()->GetBaManager()->UpdateOriginatorAgreement(respHdr, hdr.GetAddr1(), 0);
 
     //-----------------------------------------------------------------------------------------------------
 
@@ -501,7 +501,7 @@ TwoLevelAggregationTest::DoRun()
     reqHdr.SetBufferSize(64);
     reqHdr.SetTimeout(0);
     reqHdr.SetStartingSequence(0);
-    m_mac->GetVIQueue()->GetBaManager()->CreateOriginatorAgreement(&reqHdr, hdr.GetAddr1());
+    m_mac->GetVIQueue()->GetBaManager()->CreateOriginatorAgreement(reqHdr, hdr.GetAddr1());
 
     MgtAddBaResponseHeader respHdr;
     StatusCode code;
@@ -512,7 +512,7 @@ TwoLevelAggregationTest::DoRun()
     respHdr.SetTid(reqHdr.GetTid());
     respHdr.SetBufferSize(64);
     respHdr.SetTimeout(reqHdr.GetTimeout());
-    m_mac->GetVIQueue()->GetBaManager()->UpdateOriginatorAgreement(&respHdr, hdr.GetAddr1(), 0);
+    m_mac->GetVIQueue()->GetBaManager()->UpdateOriginatorAgreement(respHdr, hdr.GetAddr1(), 0);
 
     m_mac->SetAttribute("VI_MaxAmsduSize", UintegerValue(3050)); // max 2 MSDUs per A-MSDU
     m_mac->SetAttribute("VI_MaxAmpduSize", UintegerValue(65535));
@@ -694,7 +694,7 @@ HeAggregationTest::DoRunSubTest(uint16_t bufferSize)
     reqHdr.SetBufferSize(bufferSize);
     reqHdr.SetTimeout(0);
     reqHdr.SetStartingSequence(0);
-    m_mac->GetBEQueue()->GetBaManager()->CreateOriginatorAgreement(&reqHdr, hdr.GetAddr1());
+    m_mac->GetBEQueue()->GetBaManager()->CreateOriginatorAgreement(reqHdr, hdr.GetAddr1());
 
     MgtAddBaResponseHeader respHdr;
     StatusCode code;
@@ -705,7 +705,7 @@ HeAggregationTest::DoRunSubTest(uint16_t bufferSize)
     respHdr.SetTid(reqHdr.GetTid());
     respHdr.SetBufferSize(bufferSize);
     respHdr.SetTimeout(reqHdr.GetTimeout());
-    m_mac->GetBEQueue()->GetBaManager()->UpdateOriginatorAgreement(&respHdr, hdr.GetAddr1(), 0);
+    m_mac->GetBEQueue()->GetBaManager()->UpdateOriginatorAgreement(respHdr, hdr.GetAddr1(), 0);
 
     /*
      * Test behavior when 300 packets are ready for transmission but negociated buffer size is 64