fix: optimize LID migration performance and introduce cache (#1782)

* chore: lid-map cache, bulk migration, fixes

* fix: type error

* fix: signalrepository with lid type, remove check from assertsessions, lid-mapping store level deduplication

* fix: migrations logs

* fix: remove migrationsops lenght check

* fix: skip deletesession if cached or existing migrations

* chore: remove check migrated session

* fix: remove return null loadsession

* fix: getLIDforPN call and jidsrequiringfetch

* fix: longer lid map cache TTL

* fix: lint

* fix: lid type

* fix: lid socket type

src/Signal/libsignal.ts

@@ -1,9 +1,9 @@
+/* @ts-ignore */
 import * as libsignal from 'libsignal'
-import { LRUCache } from 'lru-cache'
 import type { SignalAuthState, SignalKeyStoreWithTransaction } from '../Types'
-import type { SignalRepository } from '../Types/Signal'
+import type { SignalRepositoryWithLIDStore } from '../Types/Signal'
 import { generateSignalPubKey } from '../Utils'
-import { jidDecode } from '../WABinary'
+import { jidDecode, transferDevice } from '../WABinary'
 import type { SenderKeyStore } from './Group/group_cipher'
 import { SenderKeyName } from './Group/sender-key-name'
 import { SenderKeyRecord } from './Group/sender-key-record'
@@ -20,7 +20,7 @@ export function makeLibSignalRepository(
 		  }[]
 		| undefined
 	>
-): SignalRepository {
+): SignalRepositoryWithLIDStore {
 	const lidMapping = new LIDMappingStore(auth.keys as SignalKeyStoreWithTransaction, onWhatsAppFunc)
 	const storage = signalStorage(auth, lidMapping)
 
@@ -38,13 +38,7 @@ export function makeLibSignalRepository(
 		)
 	}
 
-	// Simple operation-level deduplication (5 minutes)
-	const recentMigrations = new LRUCache<string, boolean>({
-		max: 500,
-		ttl: 5 * 60 * 1000
-	})
-
-	const repository: SignalRepository = {
+	const repository: SignalRepositoryWithLIDStore = {
 		decryptGroupMessage({ group, authorJid, msg }) {
 			const senderName = jidToSignalSenderKeyName(group, authorJid)
 			const cipher = new GroupCipher(storage, senderName)
@@ -136,7 +130,7 @@ export function makeLibSignalRepository(
 
 						if (pnSession) {
 							// Migrate PN to LID
-							await repository.migrateSession(jid, lidForPN)
+							await repository.migrateSession([jid], lidForPN)
 							encryptionJid = lidForPN
 						}
 					}
@@ -185,13 +179,8 @@ export function makeLibSignalRepository(
 			return jidToSignalProtocolAddress(jid).toString()
 		},
 
-		async storeLIDPNMapping(lid: string, pn: string) {
-			await lidMapping.storeLIDPNMapping(lid, pn)
-		},
-
-		getLIDMappingStore() {
-			return lidMapping
-		},
+		// Optimized direct access to LID mapping store
+		lidMapping,
 
 		async validateSession(jid: string) {
 			try {
@@ -212,71 +201,97 @@ export function makeLibSignalRepository(
 			}
 		},
 
-		async deleteSession(jid: string) {
-			const addr = jidToSignalProtocolAddress(jid)
+		async deleteSession(jids: string[]) {
+			if (!jids.length) return
 
+			// Convert JIDs to signal addresses and prepare for bulk deletion
+			const sessionUpdates: { [key: string]: null } = {}
+			jids.forEach(jid => {
+				const addr = jidToSignalProtocolAddress(jid)
+				sessionUpdates[addr.toString()] = null
+			})
+
+			// Single transaction for all deletions
 			return parsedKeys.transaction(async () => {
-				await auth.keys.set({ session: { [addr.toString()]: null } })
-			}, jid)
+				await auth.keys.set({ session: sessionUpdates })
+			}, `delete-${jids.length}-sessions`)
 		},
 
-		async migrateSession(fromJid: string, toJid: string) {
-			// Only migrate PN → LID
-			if (!fromJid.includes('@s.whatsapp.net') || !toJid.includes('@lid')) {
-				return
-			}
+		async migrateSession(
+			fromJids: string[],
+			toJid: string
+		): Promise<{ migrated: number; skipped: number; total: number }> {
+			if (!fromJids.length || !toJid.includes('@lid')) return { migrated: 0, skipped: 0, total: 0 }
 
-			const fromDecoded = jidDecode(fromJid)
-			const toDecoded = jidDecode(toJid)
-			if (!fromDecoded || !toDecoded) return
+			// Filter valid PN JIDs
+			const validJids = fromJids.filter(jid => jid.includes('@s.whatsapp.net'))
+			if (!validJids.length) return { migrated: 0, skipped: 0, total: fromJids.length }
 
-			const deviceId = fromDecoded.device || 0
-			const migrationKey = `${fromDecoded.user}.${deviceId}→${toDecoded.user}.${deviceId}`
+			// Single optimized transaction for all migrations
+			return parsedKeys.transaction(
+				async (): Promise<{ migrated: number; skipped: number; total: number }> => {
+					// 1. Batch store all LID mappings
+					const mappings = validJids.map(jid => ({
+						lid: transferDevice(jid, toJid),
+						pn: jid
+					}))
+					await lidMapping.storeLIDPNMappings(mappings)
 
-			// Check if recently migrated (5 min window)
-			if (recentMigrations.has(migrationKey)) {
-				return
-			}
+					// 2. Prepare migration operations
+					const migrationOps = validJids.map(jid => {
+						const lidWithDevice = transferDevice(jid, toJid)
+						const fromDecoded = jidDecode(jid)!
+						const toDecoded = jidDecode(lidWithDevice)!
 
-			// Check if LID session already exists
-			const lidAddr = jidToSignalProtocolAddress(toJid)
-			const { [lidAddr.toString()]: lidExists } = await auth.keys.get('session', [lidAddr.toString()])
-			if (lidExists) {
-				recentMigrations.set(migrationKey, true)
-				return
-			}
+						return {
+							fromJid: jid,
+							toJid: lidWithDevice,
+							pnUser: fromDecoded.user,
+							lidUser: toDecoded.user,
+							deviceId: fromDecoded.device || 0,
+							fromAddr: jidToSignalProtocolAddress(jid),
+							toAddr: jidToSignalProtocolAddress(lidWithDevice)
+						}
+					})
 
-			return parsedKeys.transaction(async () => {
-				// Store mapping
-				await lidMapping.storeLIDPNMapping(toJid, fromJid)
+					// 3. Batch check which LID sessions already exist
+					const lidAddrs = migrationOps.map(op => op.toAddr.toString())
+					const existingSessions = await auth.keys.get('session', lidAddrs)
 
-				// Load and copy session
-				const fromAddr = jidToSignalProtocolAddress(fromJid)
-				const fromSession = await storage.loadSession(fromAddr.toString())
+					// 4. Filter out sessions that already have LID sessions
+					const opsToMigrate = migrationOps.filter(op => !existingSessions[op.toAddr.toString()])
+					const skippedCount = migrationOps.length - opsToMigrate.length
 
-				if (fromSession?.haveOpenSession()) {
-					// Deep copy session to prevent reference issues
-					const sessionBytes = fromSession.serialize()
-					const copiedSession = libsignal.SessionRecord.deserialize(sessionBytes)
+					if (!opsToMigrate.length) {
+						return { migrated: 0, skipped: skippedCount, total: validJids.length }
+					}
 
-					// Store at LID address
-					await storage.storeSession(lidAddr.toString(), copiedSession)
+					// 5. Execute all migrations in parallel
+					await Promise.all(
+						opsToMigrate.map(async op => {
+							const fromSession = await storage.loadSession(op.fromAddr.toString())
 
-					// Delete PN session - maintain single encryption layer
-					await auth.keys.set({ session: { [fromAddr.toString()]: null } })
-				}
+							if (fromSession?.haveOpenSession()) {
+								// Copy session to LID address
+								const sessionBytes = fromSession.serialize()
+								const copiedSession = libsignal.SessionRecord.deserialize(sessionBytes)
+								await storage.storeSession(op.toAddr.toString(), copiedSession)
 
-				recentMigrations.set(migrationKey, true)
-			}, fromJid)
+								// Delete PN session
+								await auth.keys.set({ session: { [op.fromAddr.toString()]: null } })
+							}
+						})
+					)
+
+					return { migrated: opsToMigrate.length, skipped: skippedCount, total: validJids.length }
+				},
+				`migrate-${validJids.length}-sessions-${jidDecode(toJid)?.user}`
+			)
 		},
 
 		async encryptMessageWithWire({ encryptionJid, wireJid, data }) {
 			const result = await repository.encryptMessage({ jid: encryptionJid, data })
 			return { ...result, wireJid }
-		},
-
-		destroy() {
-			recentMigrations.clear()
 		}
 	}
 
@@ -287,6 +302,12 @@ const jidToSignalProtocolAddress = (jid: string): libsignal.ProtocolAddress => {
 	const decoded = jidDecode(jid)!
 	const { user, device, server } = decoded
 
+	if (!user) {
+		throw new Error(
+			`JID decoded but user is empty: "${jid}" -> user: "${user}", server: "${server}", device: ${device}`
+		)
+	}
+
 	// LID addresses get _1 suffix for Signal protocol
 	const signalUser = server === 'lid' ? `${user}_1` : user
 	const finalDevice = device || 0

src/Signal/lid-mapping.ts

@@ -1,9 +1,14 @@
+import { LRUCache } from 'lru-cache'
 import type { SignalKeyStoreWithTransaction } from '../Types'
 import logger from '../Utils/logger'
 import { isLidUser, isPnUser, jidDecode } from '../WABinary'
 
-//TODO: Caching
 export class LIDMappingStore {
+	private readonly mappingCache = new LRUCache<string, string>({
+		ttl: 7 * 24 * 60 * 60 * 1000, // 7 days
+		ttlAutopurge: true,
+		updateAgeOnGet: true
+	})
 	private readonly keys: SignalKeyStoreWithTransaction
 	private onWhatsAppFunc?: (...jids: string[]) => Promise<
 		| {
@@ -29,13 +34,6 @@ export class LIDMappingStore {
 		this.onWhatsAppFunc = onWhatsAppFunc // needed to get LID from PN if not found
 	}
 
-	/**
-	 * Store LID-PN mapping - USER LEVEL
-	 */
-	async storeLIDPNMapping(lid: string, pn: string): Promise<void> {
-		return this.storeLIDPNMappings([{ lid, pn }])
-	}
-
 	/**
 	 * Store LID-PN mapping - USER LEVEL
 	 */
@@ -57,6 +55,25 @@ export class LIDMappingStore {
 
 			const pnUser = pnDecoded.user
 			const lidUser = lidDecoded.user
+
+			// Check if mapping already exists (cache first, then database)
+			let existingLidUser = this.mappingCache.get(`pn:${pnUser}`)
+			if (!existingLidUser) {
+				// Cache miss - check database
+				const stored = await this.keys.get('lid-mapping', [pnUser])
+				existingLidUser = stored[pnUser]
+				if (existingLidUser) {
+					// Update cache with database value
+					this.mappingCache.set(`pn:${pnUser}`, existingLidUser)
+					this.mappingCache.set(`lid:${existingLidUser}`, pnUser)
+				}
+			}
+
+			if (existingLidUser === lidUser) {
+				logger.debug({ pnUser, lidUser }, 'LID mapping already exists, skipping')
+				continue
+			}
+
 			pairMap[pnUser] = lidUser
 		}
 
@@ -70,6 +87,10 @@ export class LIDMappingStore {
 						[`${lidUser}_reverse`]: pnUser // "102765716062358_reverse" -> "554396160286"
 					}
 				})
+
+				// Update cache with both directions
+				this.mappingCache.set(`pn:${pnUser}`, lidUser)
+				this.mappingCache.set(`lid:${lidUser}`, pnUser)
 			}
 		}, 'lid-mapping')
 	}
@@ -83,22 +104,38 @@ export class LIDMappingStore {
 		const decoded = jidDecode(pn)
 		if (!decoded) return null
 
-		// Look up user-level mapping (whatsmeow approach)
+		// Check cache first for PN → LID mapping
 		const pnUser = decoded.user
-		const stored = await this.keys.get('lid-mapping', [pnUser])
-		let lidUser = stored[pnUser]
+		let lidUser = this.mappingCache.get(`pn:${pnUser}`)
 
 		if (!lidUser) {
-			logger.trace(`No LID mapping found for PN user ${pnUser}; getting from USync`)
-			const { exists, lid } = (await this.onWhatsAppFunc?.(pn))?.[0]! // this function already adds LIDs to mapping
-			if (exists) {
-				lidUser = jidDecode(lid)?.user
+			// Cache miss - check database
+			const stored = await this.keys.get('lid-mapping', [pnUser])
+			lidUser = stored[pnUser]
+
+			if (lidUser) {
+				// Cache the database result
+				this.mappingCache.set(`pn:${pnUser}`, lidUser)
 			} else {
-				return null
+				// Not in database - try USync
+				logger.trace(`No LID mapping found for PN user ${pnUser}; getting from USync`)
+				const { exists, lid } = (await this.onWhatsAppFunc?.(pn))?.[0]! // this function already adds LIDs to mapping
+				if (exists && lid) {
+					lidUser = jidDecode(lid)?.user
+					if (lidUser) {
+						// Cache the USync result
+						this.mappingCache.set(`pn:${pnUser}`, lidUser)
+					}
+				} else {
+					return null
+				}
 			}
 		}
 
-		if (typeof lidUser !== 'string') return null
+		if (typeof lidUser !== 'string' || !lidUser) {
+			logger.warn(`Invalid or empty LID user for PN ${pn}: lidUser = "${lidUser}"`)
+			return null
+		}
 
 		// Push the PN device ID to the LID to maintain device separation
 		const pnDevice = decoded.device !== undefined ? decoded.device : 0
@@ -117,15 +154,21 @@ export class LIDMappingStore {
 		const decoded = jidDecode(lid)
 		if (!decoded) return null
 
-		// Look up reverse user mapping
+		// Check cache first for LID → PN mapping
 		const lidUser = decoded.user
-		// TODO: remove this style and instead load all mappings somehow, and then assign them in the map
-		const stored = await this.keys.get('lid-mapping', [`${lidUser}_reverse`])
-		const pnUser = stored[`${lidUser}_reverse`]
+		let pnUser = this.mappingCache.get(`lid:${lidUser}`)
 
 		if (!pnUser || typeof pnUser !== 'string') {
-			logger.trace(`No reverse mapping found for LID user: ${lidUser}`)
-			return null
+			// Cache miss - check database
+			const stored = await this.keys.get('lid-mapping', [`${lidUser}_reverse`])
+			pnUser = stored[`${lidUser}_reverse`]
+
+			if (!pnUser || typeof pnUser !== 'string') {
+				logger.trace(`No reverse mapping found for LID user: ${lidUser}`)
+				return null
+			}
+
+			this.mappingCache.set(`lid:${lidUser}`, pnUser)
 		}
 
 		// Construct device-specific PN JID

src/Socket/messages-recv.ts

@@ -1150,21 +1150,24 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
 			const lid = jidNormalizedUser(node.attrs.from),
 				pn = jidNormalizedUser(node.attrs.sender_pn)
 			ev.emit('lid-mapping.update', { lid, pn })
-			await signalRepository.storeLIDPNMapping(lid, pn)
+			await signalRepository.lidMapping.storeLIDPNMappings([{ lid, pn }])
 		}
 
 		const alt = msg.key.participantAlt || msg.key.remoteJidAlt
 		// store new mappings we didn't have before
 		if (!!alt) {
 			const altServer = jidDecode(alt)?.server
-			const lidMapping = signalRepository.getLIDMappingStore()
 			if (altServer === 'lid') {
-				if (typeof (await lidMapping.getPNForLID(alt)) === 'string') {
-					await lidMapping.storeLIDPNMapping(alt, msg.key.participant || msg.key.remoteJid!)
+				if (typeof (await signalRepository.lidMapping.getPNForLID(alt)) === 'string') {
+					await signalRepository.lidMapping.storeLIDPNMappings([
+						{ lid: alt, pn: msg.key.participant || msg.key.remoteJid! }
+					])
 				}
 			} else {
-				if (typeof (await lidMapping.getLIDForPN(alt)) === 'string') {
-					await lidMapping.storeLIDPNMapping(msg.key.participant || msg.key.remoteJid!, alt)
+				if (typeof (await signalRepository.lidMapping.getLIDForPN(alt)) === 'string') {
+					await signalRepository.lidMapping.storeLIDPNMappings([
+						{ lid: msg.key.participant || msg.key.remoteJid!, pn: alt }
+					])
 				}
 			}
 		}

src/Socket/messages-send.ts

@@ -47,8 +47,7 @@ import {
 	jidEncode,
 	jidNormalizedUser,
 	type JidWithDevice,
-	S_WHATSAPP_NET,
-	transferDevice
+	S_WHATSAPP_NET
 } from '../WABinary'
 import { USyncQuery, USyncUser } from '../WAUSync'
 import { makeNewsletterSocket } from './newsletter'
@@ -366,19 +365,6 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 		return deviceResults
 	}
 
-	// Helper to check if JID has migrated LID session
-	const checkForMigratedLidSession = async (jid: string): Promise<boolean> => {
-		if (!jid.includes('@s.whatsapp.net')) return false
-
-		const lidMapping = signalRepository.getLIDMappingStore()
-		const lidForPN = await lidMapping.getLIDForPN(jid)
-		if (!lidForPN?.includes('@lid')) return false
-
-		const lidSignalId = signalRepository.jidToSignalProtocolAddress(lidForPN)
-		const lidSessions = await authState.keys.get('session', [lidSignalId])
-		return !!lidSessions[lidSignalId]
-	}
-
 	const assertSessions = async (jids: string[], force: boolean) => {
 		let didFetchNewSession = false
 		const jidsRequiringFetch: string[] = []
@@ -391,20 +377,13 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 			const addrs = jids.map(jid => signalRepository.jidToSignalProtocolAddress(jid))
 			const sessions = await authState.keys.get('session', addrs)
 
-			// Helper to check session for a JID
-			const checkJidSession = async (jid: string) => {
+			// Simplified: Check session existence directly
+			const checkJidSession = (jid: string) => {
 				const signalId = signalRepository.jidToSignalProtocolAddress(jid)
-				let hasSession = !!sessions[signalId]
-
-				// Check for migrated LID session if PN session missing
-				if (!hasSession) {
-					hasSession = await checkForMigratedLidSession(jid)
-					if (hasSession) {
-						logger.debug({ jid }, 'Found migrated LID session during force assert, skipping PN fetch')
-					}
-				}
+				const hasSession = !!sessions[signalId]
 
 				// Add to fetch list if no session exists
+				// Session type selection (LID vs PN) is handled in encryptMessage
 				if (!hasSession) {
 					if (jid.includes('@lid')) {
 						logger.debug({ jid }, 'No LID session found, will create new LID session')
@@ -416,10 +395,9 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 
 			// Process all JIDs
 			for (const jid of jids) {
-				await checkJidSession(jid)
+				checkJidSession(jid)
 			}
 		} else {
-			const lidMapping = signalRepository.getLIDMappingStore()
 			const addrs = jids.map(jid => signalRepository.jidToSignalProtocolAddress(jid))
 			const sessions = await authState.keys.get('session', addrs)
 
@@ -441,7 +419,9 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 				}
 
 				try {
-					const mapping = await lidMapping.getLIDForPN(user)
+					// Convert user to proper PN JID format for getLIDForPN
+					const pnJid = `${user}@s.whatsapp.net`
+					const mapping = await signalRepository.lidMapping.getLIDForPN(pnJid)
 					if (mapping?.includes('@lid')) {
 						logger.debug(
 							{ user, lidForPN: mapping, deviceCount: userJids.length },
@@ -456,27 +436,6 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 				return { shouldMigrate: false, lidForPN: undefined }
 			}
 
-			// Helper to migrate a single device
-			const migrateDeviceToLid = async (jid: string, lidForPN: string) => {
-				if (!jid.includes('@s.whatsapp.net')) return
-
-				try {
-					const lidWithDevice = transferDevice(jid, lidForPN)
-					await signalRepository.migrateSession(jid, lidWithDevice)
-					logger.debug({ fromJid: jid, toJid: lidWithDevice }, 'Migrated device session to LID')
-
-					// Delete PN session after successful migration
-					try {
-						await signalRepository.deleteSession(jid)
-						logger.debug({ deletedPNSession: jid }, 'Deleted PN session after migration')
-					} catch (deleteError) {
-						logger.warn({ jid, error: deleteError }, 'Failed to delete PN session')
-					}
-				} catch (migrationError) {
-					logger.warn({ jid, error: migrationError }, 'Failed to migrate device session')
-				}
-			}
-
 			// Process each user group for potential bulk LID migration
 			for (const [user, userJids] of userGroups) {
 				const mappingResult = await checkUserLidMapping(user, userJids)
@@ -485,59 +444,52 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 
 				// Migrate all devices for this user if LID mapping exists
 				if (shouldMigrateUser && lidForPN) {
-					// Migrate each device individually
-					for (const jid of userJids) {
-						await migrateDeviceToLid(jid, lidForPN)
-					}
+					// Bulk migrate all user devices in single transaction
+					const migrationResult = await signalRepository.migrateSession(userJids, lidForPN)
 
-					logger.info(
-						{
-							user,
-							lidMapping: lidForPN,
-							deviceCount: userJids.length
-						},
-						'Completed migration attempt for user devices'
-					)
-				}
-
-				// Helper to check session for migrated user
-				const checkMigratedSession = async (jid: string) => {
-					const signalId = signalRepository.jidToSignalProtocolAddress(jid)
-					let hasSession = !!sessions[signalId]
-					let jidToFetch = jid
-
-					// Check if we should use migrated LID session instead
-					if (shouldMigrateUser && lidForPN && jid.includes('@s.whatsapp.net')) {
-						const originalDecoded = jidDecode(jid)
-						const deviceId = originalDecoded?.device || 0
-						const lidDecoded = jidDecode(lidForPN)
-						const lidWithDevice = jidEncode(lidDecoded?.user!, 'lid', deviceId)
-
-						// Check if LID session exists
-						const lidSignalId = signalRepository.jidToSignalProtocolAddress(lidWithDevice)
-						const lidSessions = await authState.keys.get('session', [lidSignalId])
-						hasSession = !!lidSessions[lidSignalId]
-						jidToFetch = lidWithDevice
-
-						if (hasSession) {
-							logger.debug({ originalJid: jid, lidJid: lidWithDevice }, '✅ Found bulk-migrated LID session')
-						}
-					}
-
-					// Add to fetch list if no session exists
-					if (!hasSession) {
-						jidsRequiringFetch.push(jidToFetch)
+					if (migrationResult.migrated > 0) {
+						logger.info(
+							{
+								user,
+								lidMapping: lidForPN,
+								migrated: migrationResult.migrated,
+								skipped: migrationResult.skipped,
+								total: migrationResult.total
+							},
+							'Completed bulk migration for user devices'
+						)
+					} else {
 						logger.debug(
-							{ jid: jidToFetch, originalJid: jid !== jidToFetch ? jid : undefined },
-							'Adding to session fetch list'
+							{
+								user,
+								lidMapping: lidForPN,
+								skipped: migrationResult.skipped,
+								total: migrationResult.total
+							},
+							'All user device sessions already migrated'
 						)
 					}
 				}
 
-				// Now check which sessions need to be fetched for this user
-				for (const jid of userJids) {
-					await checkMigratedSession(jid)
+				// Direct bulk session check with LID single source of truth
+				const addMissingSessionsToFetchList = (jid: string) => {
+					const signalId = signalRepository.jidToSignalProtocolAddress(jid)
+					if (sessions[signalId]) return
+
+					// Determine correct JID to fetch (LID if mapping exists, otherwise original)
+					if (jid.includes('@s.whatsapp.net') && shouldMigrateUser && lidForPN) {
+						const decoded = jidDecode(jid)!
+						const lidDeviceJid =
+							decoded.device !== undefined ? `${jidDecode(lidForPN)!.user}:${decoded.device}@lid` : lidForPN
+						jidsRequiringFetch.push(lidDeviceJid)
+						logger.debug({ pnJid: jid, lidJid: lidDeviceJid }, 'Adding LID JID to fetch list (conversion)')
+					} else {
+						jidsRequiringFetch.push(jid)
+						logger.debug({ jid }, 'Adding JID to fetch list')
+					}
 				}
+
+				userJids.forEach(addMissingSessionsToFetchList)
 			}
 		}
 
@@ -674,8 +626,7 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 					if (!wireJid.includes('@s.whatsapp.net')) return wireJid
 
 					try {
-						const lidMapping = signalRepository.getLIDMappingStore()
-						const lidForPN = await lidMapping.getLIDForPN(wireJid)
+						const lidForPN = await signalRepository.lidMapping.getLIDForPN(wireJid)
 
 						if (!lidForPN?.includes('@lid')) return wireJid
 
@@ -687,7 +638,7 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 
 						// Migrate session to LID for unified encryption layer
 						try {
-							await signalRepository.migrateSession(wireJid, lidWithDevice)
+							const migrationResult = await signalRepository.migrateSession([wireJid], lidWithDevice)
 							const recipientUser = jidNormalizedUser(wireJid)
 							const ownPnUser = jidNormalizedUser(meId)
 							const isOwnDevice = recipientUser === ownPnUser
@@ -695,8 +646,10 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 
 							// Delete PN session after successful migration
 							try {
-								await signalRepository.deleteSession(wireJid)
-								logger.debug({ deletedPNSession: wireJid }, 'Deleted PN session')
+								if (migrationResult.migrated) {
+									await signalRepository.deleteSession([wireJid])
+									logger.debug({ deletedPNSession: wireJid }, 'Deleted PN session')
+								}
 							} catch (deleteError) {
 								logger.warn({ wireJid, error: deleteError }, 'Failed to delete PN session')
 							}

src/Socket/socket.ts

@@ -271,10 +271,9 @@ export const makeSocket = (config: SocketConfig) => {
 		const results = await executeUSyncQuery(usyncQuery)
 
 		if (results) {
-			if (results.list.filter(a => a.lid).length > 0) {
-				const lidMapping = signalRepository.getLIDMappingStore()
-				const lidOnly = results.list.filter(a => a.lid)
-				await lidMapping.storeLIDPNMappings(lidOnly.map(a => ({ pn: a.id, lid: a.lid as string })))
+			if (results.list.filter(a => !!a.lid).length > 0) {
+				const lidOnly = results.list.filter(a => !!a.lid)
+				await signalRepository.lidMapping.storeLIDPNMappings(lidOnly.map(a => ({ pn: a.id, lid: a.lid as string })))
 			}
 
 			return results.list
@@ -851,10 +850,10 @@ export const makeSocket = (config: SocketConfig) => {
 					const myPN = authState.creds.me!.id
 
 					// Store our own LID-PN mapping
-					await signalRepository.storeLIDPNMapping(myLID, myPN)
+					await signalRepository.lidMapping.storeLIDPNMappings([{ lid: myLID, pn: myPN }])
 
 					// Create LID session for ourselves (whatsmeow pattern)
-					await signalRepository.migrateSession(myPN, myLID)
+					await signalRepository.migrateSession([myPN], myLID)
 
 					logger.info({ myPN, myLID }, 'Own LID session created successfully')
 				} catch (error) {

src/Types/Signal.ts

@@ -76,10 +76,12 @@ export type SignalRepository = {
 	injectE2ESession(opts: E2ESessionOpts): Promise<void>
 	validateSession(jid: string): Promise<{ exists: boolean; reason?: string }>
 	jidToSignalProtocolAddress(jid: string): string
-	storeLIDPNMapping(lid: string, pn: string): Promise<void>
-	getLIDMappingStore(): LIDMappingStore
-	migrateSession(fromJid: string, toJid: string): Promise<void>
+	migrateSession(fromJids: string[], toJid: string): Promise<{ migrated: number; skipped: number; total: number }>
 	validateSession(jid: string): Promise<{ exists: boolean; reason?: string }>
-	deleteSession(jid: string): Promise<void>
-	destroy(): void
+	deleteSession(jids: string[]): Promise<void>
+}
+
+// Optimized repository with pre-loaded LID mapping store
+export interface SignalRepositoryWithLIDStore extends SignalRepository {
+	lidMapping: LIDMappingStore
 }

src/Types/Socket.ts

@@ -6,7 +6,7 @@ import type { ILogger } from '../Utils/logger'
 import type { AuthenticationState, SignalAuthState, TransactionCapabilityOptions } from './Auth'
 import type { GroupMetadata } from './GroupMetadata'
 import { type MediaConnInfo } from './Message'
-import type { SignalRepository } from './Signal'
+import type { SignalRepositoryWithLIDStore } from './Signal'
 
 export type WAVersion = [number, number, number]
 export type WABrowserDescription = [string, string, string]
@@ -153,5 +153,5 @@ export type SocketConfig = {
 			  }[]
 			| undefined
 		>
-	) => SignalRepository
+	) => SignalRepositoryWithLIDStore
 }

src/Utils/decode-wa-message.ts

@@ -1,6 +1,7 @@
 import { Boom } from '@hapi/boom'
 import { proto } from '../../WAProto/index.js'
-import type { SignalRepository, WAMessage, WAMessageKey } from '../Types'
+import type { WAMessage, WAMessageKey } from '../Types'
+import type { SignalRepositoryWithLIDStore } from '../Types/Signal'
 import {
 	areJidsSameUser,
 	type BinaryNode,
@@ -16,26 +17,26 @@ import {
 import { unpadRandomMax16 } from './generics'
 import type { ILogger } from './logger'
 
-const getDecryptionJid = async (sender: string, repository: SignalRepository): Promise<string> => {
+const getDecryptionJid = async (sender: string, repository: SignalRepositoryWithLIDStore): Promise<string> => {
 	if (!sender.includes('@s.whatsapp.net')) {
 		return sender
 	}
 
-	return (await repository.getLIDMappingStore().getLIDForPN(sender))!
+	return (await repository.lidMapping.getLIDForPN(sender))!
 }
 
 const storeMappingFromEnvelope = async (
 	stanza: BinaryNode,
 	sender: string,
 	decryptionJid: string,
-	repository: SignalRepository,
+	repository: SignalRepositoryWithLIDStore,
 	logger: ILogger
 ): Promise<void> => {
 	const { senderAlt } = extractAddressingContext(stanza)
 
 	if (senderAlt && isLidUser(senderAlt) && isPnUser(sender) && decryptionJid === sender) {
 		try {
-			await repository.storeLIDPNMapping(senderAlt, sender)
+			await repository.lidMapping.storeLIDPNMappings([{ lid: senderAlt, pn: sender }])
 			logger.debug({ sender, senderAlt }, 'Stored LID mapping from envelope')
 		} catch (error) {
 			logger.warn({ sender, senderAlt, error }, 'Failed to store LID mapping')
@@ -206,7 +207,7 @@ export const decryptMessageNode = (
 	stanza: BinaryNode,
 	meId: string,
 	meLid: string,
-	repository: SignalRepository,
+	repository: SignalRepositoryWithLIDStore,
 	logger: ILogger
 ) => {
 	const { fullMessage, author, sender } = decodeMessageNode(stanza, meId, meLid)

src/Utils/process-message.ts

@@ -10,7 +10,7 @@ import type {
 	RequestJoinAction,
 	RequestJoinMethod,
 	SignalKeyStoreWithTransaction,
-	SignalRepository
+	SignalRepositoryWithLIDStore
 } from '../Types'
 import { WAMessageStubType } from '../Types'
 import { getContentType, normalizeMessageContent } from '../Utils/messages'
@@ -28,7 +28,7 @@ type ProcessMessageContext = {
 	ev: BaileysEventEmitter
 	logger?: ILogger
 	options: AxiosRequestConfig<{}>
-	signalRepository: SignalRepository
+	signalRepository: SignalRepositoryWithLIDStore
 }
 
 const REAL_MSG_STUB_TYPES = new Set([
@@ -304,7 +304,6 @@ const processMessage = async (
 				])
 				break
 			case proto.Message.ProtocolMessage.Type.LID_MIGRATION_MAPPING_SYNC:
-				const lidMappingStore = signalRepository.getLIDMappingStore()
 				const encodedPayload = protocolMsg.lidMigrationMappingSyncMessage?.encodedMappingPayload!
 				const { pnToLidMappings, chatDbMigrationTimestamp } =
 					proto.LIDMigrationMappingSyncPayload.decode(encodedPayload)
@@ -315,7 +314,7 @@ const processMessage = async (
 					pairs.push({ lid: `${lid}@lid`, pn: `${pn}@s.whatsapp.net` })
 				}
 
-				await lidMappingStore.storeLIDPNMappings(pairs)
+				await signalRepository.lidMapping.storeLIDPNMappings(pairs)
 		}
 	} else if (content?.reactionMessage) {
 		const reaction: proto.IReaction = {

src/Utils/signal.ts

@@ -1,5 +1,5 @@
 import { KEY_BUNDLE_TYPE } from '../Defaults'
-import type { SignalRepository } from '../Types'
+import type { SignalRepositoryWithLIDStore } from '../Types'
 import type {
 	AuthenticationCreds,
 	AuthenticationState,
@@ -85,7 +85,7 @@ export const xmppPreKey = (pair: KeyPair, id: number): BinaryNode => ({
 	]
 })
 
-export const parseAndInjectE2ESessions = async (node: BinaryNode, repository: SignalRepository) => {
+export const parseAndInjectE2ESessions = async (node: BinaryNode, repository: SignalRepositoryWithLIDStore) => {
 	const extractKey = (key: BinaryNode) =>
 		key
 			? {