@nexustechpro/baileys 2.0.5 → 2.1.0

package/lib/Socket/socket.js

@@ -1,16 +1,60 @@
-import { Boom } from "@hapi/boom"
-import { randomBytes } from "crypto"
-import { URL } from "url"
-import { promisify } from "util"
-import { proto } from "../../WAProto/index.js"
-import { DEF_CALLBACK_PREFIX, DEF_TAG_PREFIX, INITIAL_PREKEY_COUNT, MIN_PREKEY_COUNT, MIN_UPLOAD_INTERVAL, NOISE_WA_HEADER, UPLOAD_TIMEOUT, BATCH_SIZE, TimeMs } from "../Defaults/index.js"
-import { DisconnectReason } from "../Types/index.js"
-import { addTransactionCapability, aesEncryptCTR, bindWaitForConnectionUpdate, bytesToCrockford, configureSuccessfulPairing, Curve, derivePairingCodeKey, generateLoginNode, generateMdTagPrefix, generateRegistrationNode, getCodeFromWSError, getErrorCodeFromStreamError, getNextPreKeysNode, makeEventBuffer, makeNoiseHandler, promiseTimeout, signedKeyPair, xmppSignedPreKey } from "../Utils/index.js"
-import { getPlatformId, migrateIndexKey } from "../Utils/index.js"
-import { assertNodeErrorFree, binaryNodeToString, encodeBinaryNode, getBinaryNodeChild, getBinaryNodeChildren, isLidUser, jidDecode, jidEncode, S_WHATSAPP_NET } from "../WABinary/index.js"
-import { BinaryInfo } from "../WAM/BinaryInfo.js"
-import { USyncQuery, USyncUser } from "../WAUSync/index.js"
-import { WebSocketClient } from "./Client/index.js"
+import { Boom } from '@hapi/boom'
+import { randomBytes } from 'crypto'
+import { URL } from 'url'
+import { promisify } from 'util'
+import { proto } from '../../WAProto/index.js'
+import {
+  DEF_CALLBACK_PREFIX,
+  DEF_TAG_PREFIX,
+  INITIAL_PREKEY_COUNT,
+  MIN_PREKEY_COUNT,
+  MIN_UPLOAD_INTERVAL,
+  NOISE_WA_HEADER,
+  UPLOAD_TIMEOUT,
+  BATCH_SIZE,
+  TimeMs
+} from '../Defaults/index.js'
+import { QueryIds, ReachoutTimelockEnforcementType } from '../Types/index.js'
+import { DisconnectReason, XWAPaths } from '../Types/index.js'
+import {
+  addTransactionCapability,
+  aesEncryptCTR,
+  bindWaitForConnectionUpdate,
+  buildPairingQRData,
+  bytesToCrockford,
+  configureSuccessfulPairing,
+  Curve,
+  derivePairingCodeKey,
+  generateLoginNode,
+  generateMdTagPrefix,
+  generateRegistrationNode,
+  getCodeFromWSError,
+  getCompanionPlatformId,
+  getErrorCodeFromStreamError,
+  getNextPreKeysNode,
+  makeEventBuffer,
+  makeNoiseHandler,
+  promiseTimeout,
+  signedKeyPair,
+  xmppSignedPreKey
+} from '../Utils/index.js'
+import { getPlatformId, migrateIndexKey } from '../Utils/index.js'
+import {
+  assertNodeErrorFree,
+  binaryNodeToString,
+  encodeBinaryNode,
+  getAllBinaryNodeChildren,
+  getBinaryNodeChild,
+  getBinaryNodeChildren,
+  isLidUser,
+  jidDecode,
+  jidEncode,
+  S_WHATSAPP_NET
+} from '../WABinary/index.js'
+import { BinaryInfo } from '../WAM/BinaryInfo.js'
+import { USyncQuery, USyncUser } from '../WAUSync/index.js'
+import { WebSocketClient } from './Client/index.js'
+import { executeWMexQuery } from './mex.js'
 
 // ─── Module-scope helpers ──────────────────────────────────────────────────────
 
@@ -43,7 +87,7 @@ export const makeSocket = (config) => {
     makeSignalRepository
   } = config
 
-  if (printQRInTerminal) logger?.warn("printQRInTerminal deprecated")
+
 
   const url = typeof waWebSocketUrl === 'string' ? new URL(waWebSocketUrl) : waWebSocketUrl
 
@@ -65,6 +109,14 @@ export const makeSocket = (config) => {
   logger.info({ url: url.toString() }, 'Initiating WebSocket connection')
   ws.connect()
 
+  if (printQRInTerminal) {
+    import('qrcode-terminal').then(({ default: qrcode }) => {
+      ev.on('connection.update', ({ qr }) => {
+        if (qr) qrcode.generate(qr, { small: true })
+      })
+    })
+  }
+
   const ev = makeEventBuffer(logger)
   const { creds } = authState
   const keys = addTransactionCapability(authState.keys, logger, transactionOpts)
@@ -73,6 +125,7 @@ export const makeSocket = (config) => {
   const uqTagId = generateMdTagPrefix()
   const sendPromise = promisify(ws.send)
 
+  // State
   let epoch = 1
   let lastDateRecv
   let lastUploadTime = 0
@@ -81,6 +134,10 @@ export const makeSocket = (config) => {
   let keepAliveReq
   let qrTimer
   let serverTimeOffsetMs = 0
+  let didStartBuffer = false
+
+  /** Socket end handlers — registered via registerSocketEndHandler() */
+  const socketEndHandlers = []
 
   const generateMessageTag = () => `${uqTagId}${epoch++}`
 
@@ -153,7 +210,13 @@ export const makeSocket = (config) => {
       attrs: { to: S_WHATSAPP_NET, type: 'get', xmlns: 'usync' },
       content: [{
         tag: 'usync',
-        attrs: { context: usyncQuery.context, mode: usyncQuery.mode, sid: generateMessageTag(), last: 'true', index: '0' },
+        attrs: {
+          context: usyncQuery.context,
+          mode: usyncQuery.mode,
+          sid: generateMessageTag(),
+          last: 'true',
+          index: '0'
+        },
         content: [
           { tag: 'query', attrs: {}, content: usyncQuery.protocols.map((a) => a.getQueryElement()) },
           { tag: 'list', attrs: {}, content: userNodes }
@@ -166,14 +229,39 @@ export const makeSocket = (config) => {
   async function pnFromLIDUSync(jids) {
     const usyncQuery = new USyncQuery().withLIDProtocol().withContext('background')
     for (const jid of jids) {
-      if (!isLidUser(jid)) usyncQuery.withUser(new USyncUser().withId(jid))
-      else logger?.warn('LID user found in LID fetch call')
+      if (isLidUser(jid)) { logger?.warn('LID user found in LID fetch call'); continue }
+      usyncQuery.withUser(new USyncUser().withId(jid))
     }
     if (usyncQuery.users.length === 0) return []
     const results = await executeUSyncQuery(usyncQuery)
     return results ? results.list.filter((a) => !!a.lid).map(({ lid, id }) => ({ pn: id, lid })) : []
   }
 
+  const onWhatsApp = async (...phoneNumbers) => {
+    let usyncQuery = new USyncQuery()
+    let contactEnabled = false
+    for (let jid of phoneNumbers) {
+      if (isLidUser(jid)) {
+        const pn = await signalRepository.lidMapping.getPNForLID(jid)
+        if (pn) {
+          jid = pn
+        } else {
+          if (!contactEnabled) { contactEnabled = true; usyncQuery = usyncQuery.withContactProtocol() }
+          usyncQuery.withUser(new USyncUser().withId(jid))
+          continue
+        }
+      }
+      if (!contactEnabled) { contactEnabled = true; usyncQuery = usyncQuery.withContactProtocol() }
+      const phone = `+${jid.replace('+', '').split('@')[0]?.split(':')[0]}`
+      usyncQuery.withUser(new USyncUser().withPhone(phone))
+    }
+    if (usyncQuery.users.length === 0) return []
+    const results = await executeUSyncQuery(usyncQuery)
+    return results
+      ? results.list.filter((a) => !!a.contact && a.id && a.id !== 'undefined').map(({ contact, id }) => ({ jid: id, exists: contact }))
+      : []
+  }
+
   // ─── Pre-keys ───────────────────────────────────────────────────────────────
 
   const getAvailablePreKeysOnServer = async () => {
@@ -185,11 +273,24 @@ export const makeSocket = (config) => {
     return +getBinaryNodeChild(result, 'count').attrs.value
   }
 
+  /**
+   * Verify that our current pre-key actually exists in local key storage.
+   * Catches the case where the server still has keys but our local store is missing them.
+   */
+  const verifyCurrentPreKeyExists = async () => {
+    const currentPreKeyId = creds.nextPreKeyId - 1
+    if (currentPreKeyId <= 0) return { exists: false, currentPreKeyId: 0 }
+    const preKeys = await keys.get('pre-key', [currentPreKeyId.toString()])
+    return { exists: !!preKeys[currentPreKeyId.toString()], currentPreKeyId }
+  }
+
   const uploadPreKeys = async (count = MIN_PREKEY_COUNT, retryCount = 0) => {
+    // Rate-limit guard — only on the first attempt, not retries
     if (retryCount === 0 && Date.now() - lastUploadTime < MIN_UPLOAD_INTERVAL) {
-      logger.debug(`Skipping upload, only ${Date.now() - lastUploadTime}ms since last`)
+      logger.debug(`Skipping upload — only ${Date.now() - lastUploadTime}ms since last upload`)
       return
     }
+    // Dedup: if an upload is already in-flight, wait for it
     if (uploadPreKeysPromise) {
       logger.debug('Pre-key upload in progress, waiting...')
       await uploadPreKeysPromise
@@ -197,6 +298,7 @@ export const makeSocket = (config) => {
     }
     const uploadLogic = async () => {
       logger.info({ count, retryCount }, 'Uploading pre-keys')
+      // Generate keys inside a transaction to prevent ID collisions on retry
       const node = await keys.transaction(async () => {
         const { update, node } = await getNextPreKeysNode({ creds, keys }, count)
         ev.emit('creds.update', update)
@@ -230,15 +332,23 @@ export const makeSocket = (config) => {
     try {
       const preKeyCount = await getAvailablePreKeysOnServer()
       logger.info(`${preKeyCount} pre-keys found on server`)
-      if (preKeyCount < MIN_PREKEY_COUNT) {
-        const uploadCount = INITIAL_PREKEY_COUNT - preKeyCount
-        logger.info(`Server pre-key count low (${preKeyCount}), uploading ${uploadCount}`)
+      const { exists: currentPreKeyExists, currentPreKeyId } = await verifyCurrentPreKeyExists()
+      logger.info(`Current prekey ID: ${currentPreKeyId}, exists in storage: ${currentPreKeyExists}`)
+      const lowServerCount = preKeyCount <= MIN_PREKEY_COUNT
+      const missingCurrentPreKey = !currentPreKeyExists && currentPreKeyId > 0
+      if (lowServerCount || missingCurrentPreKey) {
+        const reasons = []
+        if (lowServerCount) reasons.push(`server count low (${preKeyCount})`)
+        if (missingCurrentPreKey) reasons.push(`current prekey ${currentPreKeyId} missing from storage`)
+        logger.info(`Uploading PreKeys due to: ${reasons.join(', ')}`)
+        const uploadCount = preKeyCount === 0 ? INITIAL_PREKEY_COUNT : MIN_PREKEY_COUNT
         await uploadPreKeys(uploadCount)
       } else {
-        logger.info(`✅ PreKey validation passed - Server: ${preKeyCount} pre-keys`)
+        logger.info(`✅ PreKey validation passed — Server: ${preKeyCount}, prekey ${currentPreKeyId} exists`)
       }
     } catch (error) {
       logger.error({ error }, 'Failed to check/upload pre-keys during init')
+      // Non-fatal — allow connection to continue
     }
   }
 
@@ -285,7 +395,7 @@ export const makeSocket = (config) => {
     const parsed = Number(tValue)
     if (Number.isNaN(parsed) || parsed <= 0) return
     serverTimeOffsetMs = parsed * 1000 - Date.now()
-    logger.debug({ offset: serverTimeOffsetMs }, 'calculated server time offset')
+    logger.debug({ offset: serverTimeOffsetMs }, 'Calculated server time offset')
   }
 
   // ─── Unified session telemetry ───────────────────────────────────────────────
@@ -305,14 +415,65 @@ export const makeSocket = (config) => {
         content: [{ tag: 'unified_session', attrs: { id: getUnifiedSessionId() } }]
       })
     } catch (error) {
-      logger.debug({ error }, 'failed to send unified_session telemetry')
+      logger.debug({ error }, 'Failed to send unified_session telemetry')
     }
   }
 
+  // ─── WAM buffer ─────────────────────────────────────────────────────────────
+
+  const sendWAMBuffer = (wamBuffer) =>
+    query({
+      tag: 'iq',
+      attrs: { to: S_WHATSAPP_NET, id: generateMessageTag(), xmlns: 'w:stats' },
+      content: [{ tag: 'add', attrs: { t: Math.round(Date.now() / 1000) + '' }, content: wamBuffer }]
+    })
+
+  // ─── WMex queries ────────────────────────────────────────────────────────────
+
+  /**
+   * Fetches account restriction / reachout timelock status.
+   */
+  const fetchAccountReachoutTimelock = async () => {
+    const queryResult = await executeWMexQuery(
+      {},
+      QueryIds.REACHOUT_TIMELOCK,
+      XWAPaths.xwa2_fetch_account_reachout_timelock,
+      query,
+      generateMessageTag
+    )
+    const result = {
+      isActive: !!queryResult?.is_active,
+      timeEnforcementEnds:
+        queryResult?.time_enforcement_ends && queryResult.time_enforcement_ends !== '0'
+          ? new Date(parseInt(queryResult.time_enforcement_ends, 10) * 1000)
+          : undefined,
+      enforcementType: queryResult?.enforcement_type ?? ReachoutTimelockEnforcementType.DEFAULT
+    }
+    ev.emit('connection.update', { reachoutTimeLock: result })
+    return result
+  }
+
+  /**
+   * Fetches new-chat message cap quota and usage.
+   */
+  const fetchNewChatMessageCap = async () =>
+    executeWMexQuery(
+      { input: { type: 'INDIVIDUAL_NEW_CHAT_MSG' } },
+      QueryIds.MESSAGE_CAPPING_INFO,
+      XWAPaths.xwa2_message_capping_info,
+      query,
+      generateMessageTag
+    )
+
   // ─── Connection lifecycle ────────────────────────────────────────────────────
 
   const onUnexpectedError = (err, msg) => {
-    logger.error({ err }, `unexpected error in '${msg}'`)
+    const isClosed = err?.message?.includes('Connection Closed') || err?.output?.statusCode === 428
+    if (isClosed) {
+      logger.debug({ msg: err?.message }, `Connection closed during '${msg}'`)
+    } else {
+      logger.error({ err }, `unexpected error in '${msg}'`)
+    }
   }
 
   const awaitNextMessage = async (sendMsg) => {
@@ -335,7 +496,7 @@ export const makeSocket = (config) => {
   }
 
   const validateConnection = async () => {
-    let helloMsg = proto.HandshakeMessage.fromObject({ clientHello: { ephemeral: ephemeralKeyPair.public } })
+    const helloMsg = proto.HandshakeMessage.fromObject({ clientHello: { ephemeral: ephemeralKeyPair.public } })
     logger.info({ browser, helloMsg }, 'Connected to WhatsApp')
     const init = proto.HandshakeMessage.encode(helloMsg).finish()
     const result = await awaitNextMessage(init)
@@ -346,9 +507,7 @@ export const makeSocket = (config) => {
     logger.info({ node }, !creds.me ? 'Attempting registration...' : 'Logging in...')
     const payloadEnc = noise.encrypt(proto.ClientPayload.encode(node).finish())
     await sendRawMessage(
-      proto.HandshakeMessage.encode({
-        clientFinish: { static: keyEnc, payload: payloadEnc }
-      }).finish()
+      proto.HandshakeMessage.encode({ clientFinish: { static: keyEnc, payload: payloadEnc } }).finish()
     )
     await noise.finishInit()
     startKeepAliveRequest()
@@ -376,27 +535,32 @@ export const makeSocket = (config) => {
    * Keep-alive: ping WA every keepAliveIntervalMs.
    * If the server stops responding (diff > interval + 5s) the connection
    * is considered lost and we call end() — the consumer handles reconnection.
-   * No internal reconnect loop — clean separation of concerns.
    */
   const startKeepAliveRequest = () => {
     keepAliveReq = setInterval(() => {
       if (!lastDateRecv) lastDateRecv = new Date()
       const diff = Date.now() - lastDateRecv.getTime()
       if (diff > keepAliveIntervalMs + 5000) {
-        end(new Boom('Connection was lost', { statusCode: DisconnectReason.connectionLost }))
+        void end(new Boom('Connection was lost', { statusCode: DisconnectReason.connectionLost }))
       } else if (ws.isOpen) {
         query({
           tag: 'iq',
           attrs: { id: generateMessageTag(), to: S_WHATSAPP_NET, type: 'get', xmlns: 'w:p' },
           content: [{ tag: 'ping', attrs: {} }]
-        }).catch((err) => logger.error({ trace: err.stack }, 'error in sending keep alive'))
+        }).catch((err) => logger.error({ trace: err.stack }, 'Error in sending keep alive'))
       } else {
-        logger.warn('keep alive called when WS not open')
+        logger.warn('Keep alive called when WS not open')
       }
     }, keepAliveIntervalMs)
   }
 
-  const end = (error) => {
+  /**
+   * Tear down the socket.
+   * Awaits ws.close() so cleanup is deterministic before emitting connection.update.
+   * Runs all registered socketEndHandlers in order.
+   * Calls signalRepository.close() and ev.destroy() to prevent leaks.
+   */
+  const end = async (error) => {
     if (closed) { logger.trace({ trace: error?.stack }, 'Connection already closed'); return }
     closed = true
     logger.info({ trace: error?.stack }, error ? 'connection errored' : 'connection closed')
@@ -405,9 +569,18 @@ export const makeSocket = (config) => {
     ws.removeAllListeners('close')
     ws.removeAllListeners('open')
     ws.removeAllListeners('message')
-    if (!ws.isClosed && !ws.isClosing) { try { ws.close() } catch { } }
+    signalRepository.close?.()
+    if (!ws.isClosed && !ws.isClosing) {
+      try { await ws.close() } catch { }
+    }
+    for (const handler of socketEndHandlers) {
+      try { await handler(error) }
+      catch (err) { logger.error({ err }, 'error in socket end handler') }
+    }
     ev.emit('connection.update', { connection: 'close', lastDisconnect: { error, date: new Date() } })
     ev.removeAllListeners('connection.update')
+    ev.removeAllListeners()
+    ev.destroy?.()
   }
 
   const sendPassiveIq = (tag) =>
@@ -426,14 +599,21 @@ export const makeSocket = (config) => {
         content: [{ tag: 'remove-companion-device', attrs: { jid, reason: 'user_initiated' } }]
       })
     }
-    end(new Boom(msg || 'Intentional Logout', { statusCode: DisconnectReason.loggedOut }))
+    void end(new Boom(msg || 'Intentional Logout', { statusCode: DisconnectReason.loggedOut }))
   }
 
+  /**
+   * Register a cleanup handler that will be awaited during end().
+   * Use for releasing external resources tied to this socket's lifetime.
+   */
+  const registerSocketEndHandler = (handler) => socketEndHandlers.push(handler)
+
   // ─── Pairing ─────────────────────────────────────────────────────────────────
 
   const requestPairingCode = async (phoneNumber, customPairingCode) => {
     await waitForSocketOpen()
-    await new Promise(resolve => setTimeout(resolve, 500))
+    // Brief stabilisation delay — ensures the WS open event has fully propagated
+    await new Promise((resolve) => setTimeout(resolve, 500))
     const pairingCode = customPairingCode ?? bytesToCrockford(randomBytes(5))
     if (customPairingCode && customPairingCode?.length !== 8)
       throw new Error('Custom pairing code must be exactly 8 chars')
@@ -449,7 +629,7 @@ export const makeSocket = (config) => {
         content: [
           { tag: 'link_code_pairing_wrapped_companion_ephemeral_pub', attrs: {}, content: await generatePairingKey() },
           { tag: 'companion_server_auth_key_pub', attrs: {}, content: authState.creds.noiseKey.public },
-          { tag: 'companion_platform_id', attrs: {}, content: getPlatformId(browser[1]) },
+          { tag: 'companion_platform_id', attrs: {}, content: getCompanionPlatformId(browser) },
           { tag: 'companion_platform_display', attrs: {}, content: `${browser[1]} (${browser[0]})` },
           { tag: 'link_code_pairing_nonce', attrs: {}, content: '0' }
         ]
@@ -492,31 +672,22 @@ export const makeSocket = (config) => {
     })
   }
 
-  const sendWAMBuffer = (wamBuffer) =>
-    query({
-      tag: 'iq',
-      attrs: { to: S_WHATSAPP_NET, id: generateMessageTag(), xmlns: 'w:stats' },
-      content: [{ tag: 'add', attrs: { t: Math.round(Date.now() / 1000) + '' }, content: wamBuffer }]
-    })
-
   // ─── WebSocket event bindings ────────────────────────────────────────────────
 
   ws.on('message', onMessageReceived)
 
   ws.on('open', async () => {
     try { await validateConnection() }
-    catch (err) { logger.error({ err }, 'error in validating connection'); end(err) }
+    catch (err) { logger.error({ err }, 'error in validating connection'); void end(err) }
   })
 
-  // Let mapWebSocketError convert the raw error then call end()
   ws.on('error', mapWebSocketError(end))
 
-  // Any close → end(), consumer decides whether to reconnect
-  ws.on('close', () => end(new Boom('Connection Terminated', { statusCode: DisconnectReason.connectionClosed })))
+  ws.on('close', () => void end(new Boom('Connection Terminated', { statusCode: DisconnectReason.connectionClosed })))
 
   ws.on('CB:xmlstreamend', () => {
     logger.info('Stream ended by server')
-    if (!closed) end(new Boom('Connection Terminated by Server', { statusCode: DisconnectReason.connectionClosed }))
+    if (!closed) void end(new Boom('Connection Terminated by Server', { statusCode: DisconnectReason.connectionClosed }))
   })
 
   // ─── QR pairing ─────────────────────────────────────────────────────────────
@@ -532,9 +703,10 @@ export const makeSocket = (config) => {
     const genPairQR = () => {
       if (!ws.isOpen) return
       const refNode = refNodes.shift()
-      if (!refNode) { end(new Boom('QR refs attempts ended', { statusCode: DisconnectReason.timedOut })); return }
+      if (!refNode) { void end(new Boom('QR refs attempts ended', { statusCode: DisconnectReason.timedOut })); return }
       const ref = refNode.content.toString('utf-8')
-      const qr = [ref, noiseKeyB64, identityKeyB64, advB64].join(',')
+      // Use buildPairingQRData so the browser tuple is included in the QR payload
+      const qr = buildPairingQRData(ref, noiseKeyB64, identityKeyB64, advB64, browser)
       ev.emit('connection.update', { qr })
       qrTimer = setTimeout(genPairQR, qrMs)
       qrMs = qrTimeout || 20000
@@ -554,7 +726,7 @@ export const makeSocket = (config) => {
       void sendUnifiedSession()
     } catch (error) {
       logger.info({ trace: error.stack }, 'Error in pairing')
-      end(error)
+      void end(error)
     }
   })
 
@@ -565,15 +737,11 @@ export const makeSocket = (config) => {
       updateServerTimeOffset(node)
       await uploadPreKeysToServerIfRequired()
       await sendPassiveIq('active')
-      try {
-        await digestKeyBundle()
-      } catch (e) {
-        logger.warn({ e }, 'failed to run digest after login')
-      }
+      try { await digestKeyBundle() }
+      catch (e) { logger.warn({ e }, 'failed to run digest after login') }
     } catch (err) {
       logger.warn({ err }, 'Failed to send initial passive IQ')
     }
-
     logger.info('✅ Opened connection to WhatsApp')
     clearTimeout(qrTimer)
     ev.emit('creds.update', { me: { ...authState.creds.me, lid: node.attrs.lid } })
@@ -585,18 +753,22 @@ export const makeSocket = (config) => {
       process.nextTick(async () => {
         try {
           const myPN = authState.creds.me.id
+          // Store own LID-PN mapping
           await signalRepository.lidMapping.storeLIDPNMappings([{ lid: myLID, pn: myPN }])
+          // Build device-list using index-based batching to avoid unbounded key growth
           const { user, device } = jidDecode(myPN)
           const currentBatch = await migrateIndexKey(authState.keys, 'device-list')
           currentBatch[user] = [device?.toString() || '0']
           const deviceKeys = Object.keys(currentBatch)
           if (deviceKeys.length > BATCH_SIZE) {
             deviceKeys.sort()
-            deviceKeys.slice(0, deviceKeys.length - BATCH_SIZE).forEach(k => delete currentBatch[k])
+            deviceKeys.slice(0, deviceKeys.length - BATCH_SIZE).forEach((k) => delete currentBatch[k])
           }
-          await authState.keys.set({ 'device-list': { 'index': currentBatch } })
+          await authState.keys.set({ 'device-list': { index: currentBatch } })
+          // Migrate own session from PN → LID
           await signalRepository.migrateSession(myPN, myLID)
           logger.info({ myPN, myLID }, 'Own LID session created successfully')
+          // Batch-migrate any remaining PN sessions to LID
           if (signalRepository.migrateAllPNSessionsToLID) {
             try {
               const migrated = await signalRepository.migrateAllPNSessionsToLID()
@@ -615,37 +787,37 @@ export const makeSocket = (config) => {
   // ─── Stream / connection error handlers ─────────────────────────────────────
 
   ws.on('CB:stream:error', (node) => {
-    logger.error({ node }, 'Stream errored out')
+    const [reasonNode] = getAllBinaryNodeChildren(node)
+    logger.error({ reasonNode, fullErrorNode: node }, 'Stream errored out')
     const { reason, statusCode } = getErrorCodeFromStreamError(node)
-    end(new Boom(`Stream Errored (${reason})`, { statusCode, data: node }))
+    void end(new Boom(`Stream Errored (${reason})`, { statusCode, data: reasonNode || node }))
   })
 
   ws.on('CB:failure', (node) => {
     const reason = +(node.attrs.reason || 500)
-    end(new Boom('Connection Failure', { statusCode: reason, data: node.attrs }))
+    void end(new Boom('Connection Failure', { statusCode: reason, data: node.attrs }))
   })
 
   ws.on('CB:ib,,downgrade_webclient', () =>
-    end(new Boom('Multi-device beta not joined', { statusCode: DisconnectReason.multideviceMismatch }))
+    void end(new Boom('Multi-device beta not joined', { statusCode: DisconnectReason.multideviceMismatch }))
   )
 
-  ws.on('CB:ib,,offline_preview', (node) => {
+  ws.on('CB:ib,,offline_preview', async (node) => {
     logger.info('Offline preview received', JSON.stringify(node))
-    sendNode({ tag: 'ib', attrs: {}, content: [{ tag: 'offline_batch', attrs: { count: '100' } }] })
+    await sendNode({ tag: 'ib', attrs: {}, content: [{ tag: 'offline_batch', attrs: { count: '100' } }] })
   })
 
   ws.on('CB:ib,,edge_routing', (node) => {
     const edgeRoutingNode = getBinaryNodeChild(node, 'edge_routing')
     const routingInfo = getBinaryNodeChild(edgeRoutingNode, 'routing_info')
     if (routingInfo?.content) {
-      authState.creds.routingInfo = Buffer.from(routingInfo?.content)
+      authState.creds.routingInfo = Buffer.from(routingInfo.content)
       ev.emit('creds.update', authState.creds)
     }
   })
 
   // ─── Buffering & offline notifications ──────────────────────────────────────
 
-  let didStartBuffer = false
   process.nextTick(() => {
     if (creds.me?.id) { ev.buffer(); didStartBuffer = true }
     ev.emit('connection.update', { connection: 'connecting', receivedPendingNotifications: false, qr: undefined })
@@ -689,6 +861,7 @@ export const makeSocket = (config) => {
     sendNode,
     logout,
     end,
+    registerSocketEndHandler,
     onUnexpectedError,
     uploadPreKeys,
     uploadPreKeysToServerIfRequired,
@@ -701,6 +874,9 @@ export const makeSocket = (config) => {
     waitForConnectionUpdate: bindWaitForConnectionUpdate(ev),
     sendWAMBuffer,
     executeUSyncQuery,
+    onWhatsApp,
+    fetchAccountReachoutTimelock,
+    fetchNewChatMessageCap,
     listener: (eventName) => {
       if (typeof ev.listenerCount === 'function') return ev.listenerCount(eventName)
       if (typeof ev.listener === 'function') return ev.listener(eventName)?.length || 0