diff --git a/slixmpp_omemo/__init__.py b/slixmpp_omemo/__init__.py
index f7320cb..0b81463 100644
--- a/slixmpp_omemo/__init__.py
+++ b/slixmpp_omemo/__init__.py
@@ -62,11 +62,6 @@ PUBLISH_OPTIONS_NODE = 'http://jabber.org/protocol/pubsub#publish-options'
 PUBSUB_ERRORS = 'http://jabber.org/protocol/pubsub#errors'
 
 
-class ChainLengths(TypedDict):
-    receiving: List[Tuple[int, int]]
-    sending: List[Tuple[int, int]]
-
-
 def b64enc(data: bytes) -> str:
     return base64.b64encode(bytes(bytearray(data))).decode('ASCII')
 
@@ -524,23 +519,15 @@ class XEP_0384(BasePlugin):
         devices = await self._omemo().getDevices(jid.bare)
         return devices.get('active', [])
 
-    async def _chain_lengths(self, jid: JID) -> ChainLengths:
+    async def _chain_lengths(self, jid: JID) -> Set[Tuple[int, int]]:
         """
-            Gather receiving and sending chain lengths for all devices (active
-            / inactive) of a JID.
+            Gather receiving chain lengths for all devices (active / inactive)
+            of a JID.
 
             Receiving chain length is used to know when to send a heartbeat to
             signal recipients our device is still active and listening. See:
             https://xmpp.org/extensions/xep-0384.html#rules
-
-            Sending chain length is used on the other side when a device
-            hasn't been sending us messages and seems inactive.
-
-            # XXX: Only the receiving part is used in this library for the
-            # moment.
         """
-        # XXX: This method uses APIs that haven't been made public yet in the
-        # OMEMO library as of 0.12 (9fd7123).
 
         bare = jid.bare
         devices = await self._omemo().getDevices(bare)
@@ -548,16 +535,10 @@ class XEP_0384(BasePlugin):
         inactive = devices.get('inactive', set())
         devices = active.union(inactive)
 
-        lengths: ChainLengths = {'sending': [], 'receiving': []}
+        lengths: Set[Tuple[int, int]] = set()
         for did in devices:
-            session = self._omemo()._SessionManager__loadSession(bare, did)
-            if session is None:
-                continue
-            skr = session._DoubleRatchet__skr
-            sending = skr.sending_chain_length or 0
-            receiving = skr.receiving_chain_length or 0
-            lengths['sending'].append((did, sending))
-            lengths['receiving'].append((did, receiving))
+            receiving = await self._omemo().receiving_chain_length(bare, did)
+            lengths.add((did, receiving or 0))
 
         return lengths
 
@@ -575,8 +556,7 @@ class XEP_0384(BasePlugin):
             still active.
         """
 
-        chain_lengths = await self._chain_lengths(jid)
-        receiving_chain_lengths = chain_lengths.get('receiving', [])
+        receiving_chain_lengths = await self._chain_lengths(jid)
         lengths = map(lambda d_l: d_l[1], receiving_chain_lengths)
         inactive_session = max(lengths, default=0) > self.heartbeat_after