feat: add onQr/onStatus callbacks to session; auto-restart on loggedOut

Adds OnQrCallback and OnStatusCallback parameters to createWhatsAppSession,
broadcasts QR strings and connection status changes to callers, and clears
session files then restarts automatically when the socket is logged out.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-05-29 11:04:39 +05:30
parent 43f4133f1d
commit 18edce7552
2 changed files with 133 additions and 9 deletions
+106
View File
@@ -0,0 +1,106 @@
import { Boom } from '@hapi/boom';
import { createWhatsAppSession } from './session';
import * as fs from 'fs/promises';
// Capture connection.update handler so tests can trigger it
let connectionUpdateHandler: (update: any) => Promise<void>;
// These are referenced inside jest.mock factories, so they must be declared
// before the mock calls. jest.mock is hoisted but the factory runs lazily.
const mockGroupFetch = jest.fn();
const mockEnd = jest.fn();
const mockEvOn = jest.fn();
// The baileys default export is a factory function; we need __esModule: true
// so that the `import makeWASocket from '...'` interop works correctly.
jest.mock('@whiskeysockets/baileys', () => {
const mockEvOnInner = jest.fn().mockImplementation((event: string, handler: any) => {
if (event === 'connection.update') connectionUpdateHandler = handler;
});
return {
__esModule: true,
default: jest.fn().mockReturnValue({
ev: { on: mockEvOnInner },
end: jest.fn(),
groupFetchAllParticipating: jest.fn().mockResolvedValue({}),
}),
useMultiFileAuthState: jest.fn().mockResolvedValue({ state: {}, saveCreds: jest.fn() }),
fetchLatestBaileysVersion: jest.fn().mockResolvedValue({ version: [2, 0, 0] }),
DisconnectReason: { loggedOut: 401 },
};
});
jest.mock('qrcode-terminal', () => ({ generate: jest.fn() }));
jest.mock('fs/promises', () => ({
rm: jest.fn().mockResolvedValue(undefined),
mkdir: jest.fn().mockResolvedValue(undefined),
}));
// eslint-disable-next-line @typescript-eslint/no-require-imports
const baileys = require('@whiskeysockets/baileys');
describe('createWhatsAppSession', () => {
beforeEach(() => {
// Clear all mock state (calls, instances, results) between tests
jest.clearAllMocks();
// Re-apply implementations after clearAllMocks wipes them
const freshEvOn = jest.fn().mockImplementation((event: string, handler: any) => {
if (event === 'connection.update') connectionUpdateHandler = handler;
});
baileys.default.mockReturnValue({
ev: { on: freshEvOn },
end: jest.fn(),
groupFetchAllParticipating: jest.fn().mockResolvedValue({}),
});
baileys.useMultiFileAuthState.mockResolvedValue({ state: {}, saveCreds: jest.fn() });
baileys.fetchLatestBaileysVersion.mockResolvedValue({ version: [2, 0, 0] });
(fs.rm as jest.Mock).mockResolvedValue(undefined);
(fs.mkdir as jest.Mock).mockResolvedValue(undefined);
});
it('calls onQr with raw QR string when QR event fires', async () => {
const onQr = jest.fn();
await createWhatsAppSession('acc_1', '/sessions/1', jest.fn(), jest.fn(), jest.fn(), undefined, onQr);
await connectionUpdateHandler({ qr: 'test-qr-string' });
expect(onQr).toHaveBeenCalledWith('test-qr-string');
});
it('calls onStatus with "connected" when connection opens', async () => {
const onStatus = jest.fn();
await createWhatsAppSession('acc_1', '/sessions/1', jest.fn(), jest.fn(), jest.fn(), undefined, undefined, onStatus);
await connectionUpdateHandler({ connection: 'open' });
expect(onStatus).toHaveBeenCalledWith('connected');
});
it('calls onStatus with "disconnected" on non-logout close', async () => {
const onStatus = jest.fn();
await createWhatsAppSession('acc_1', '/sessions/1', jest.fn(), jest.fn(), jest.fn(), undefined, undefined, onStatus);
const error = new Boom('Connection lost', { statusCode: 408 });
await connectionUpdateHandler({ connection: 'close', lastDisconnect: { error } });
expect(onStatus).toHaveBeenCalledWith('disconnected');
});
it('calls onStatus with "logged_out" on loggedOut close', async () => {
const onStatus = jest.fn();
await createWhatsAppSession('acc_1', '/sessions/1', jest.fn(), jest.fn(), jest.fn(), undefined, undefined, onStatus);
const error = new Boom('Logged out', { statusCode: 401 });
await connectionUpdateHandler({ connection: 'close', lastDisconnect: { error } });
expect(onStatus).toHaveBeenCalledWith('logged_out');
});
it('deletes session directory on loggedOut', async () => {
await createWhatsAppSession('acc_1', '/sessions/1', jest.fn(), jest.fn(), jest.fn());
const error = new Boom('Logged out', { statusCode: 401 });
await connectionUpdateHandler({ connection: 'close', lastDisconnect: { error } });
expect(fs.rm).toHaveBeenCalledWith('/sessions/1', { recursive: true, force: true });
expect(fs.mkdir).toHaveBeenCalledWith('/sessions/1', { recursive: true });
});
it('does not delete session directory on non-logout close', async () => {
await createWhatsAppSession('acc_1', '/sessions/1', jest.fn(), jest.fn(), jest.fn());
const error = new Boom('Timeout', { statusCode: 408 });
await connectionUpdateHandler({ connection: 'close', lastDisconnect: { error } });
expect(fs.rm).not.toHaveBeenCalled();
});
});
+27 -9
View File
@@ -6,6 +6,7 @@ import makeWASocket, {
GroupMetadata,
} from '@whiskeysockets/baileys';
import { Boom } from '@hapi/boom';
import { rm, mkdir } from 'fs/promises';
import qrcode from 'qrcode-terminal';
import { NormalizedMessage, NormalizedReaction } from '@tower/types';
import { normalizeMessage, normalizeReaction } from './normalizer';
@@ -16,6 +17,8 @@ const logger = createLogger('whatsapp-session');
export type OnMessageCallback = (msg: NormalizedMessage) => Promise<void> | void;
export type OnReactionCallback = (reaction: NormalizedReaction) => Promise<void> | void;
export type OnGroupsCallback = (groups: Record<string, GroupMetadata>) => Promise<void> | void;
export type OnQrCallback = (qr: string) => Promise<void> | void;
export type OnStatusCallback = (status: 'connected' | 'disconnected' | 'logged_out') => Promise<void> | void;
export async function createWhatsAppSession(
accountId: string,
@@ -24,6 +27,8 @@ export async function createWhatsAppSession(
onReaction: OnReactionCallback,
onGroups: OnGroupsCallback,
onReconnect?: (newSocket: WASocket) => void,
onQr?: OnQrCallback,
onStatus?: OnStatusCallback,
): Promise<WASocket> {
const { state, saveCreds } = await useMultiFileAuthState(sessionPath);
const { version } = await fetchLatestBaileysVersion();
@@ -40,26 +45,39 @@ export async function createWhatsAppSession(
sock.ev.on('connection.update', async ({ connection, qr, lastDisconnect }) => {
if (qr) {
qrcode.generate(qr, { small: true });
await Promise.resolve(onQr?.(qr)).catch((err) =>
logger.error({ err }, 'Error storing QR'),
);
}
if (connection === 'close') {
const reason = (lastDisconnect?.error as Boom)?.output?.statusCode;
const shouldReconnect = reason !== DisconnectReason.loggedOut;
logger.info({ reason, shouldReconnect }, 'Connection closed');
if (shouldReconnect) {
const isLoggedOut = reason === DisconnectReason.loggedOut;
logger.info({ reason, isLoggedOut }, 'Connection closed');
if (isLoggedOut) {
logger.info({ accountId }, 'Logged out — clearing session and restarting for QR');
await onStatus?.('logged_out');
await rm(sessionPath, { recursive: true, force: true });
await mkdir(sessionPath, { recursive: true });
setTimeout(async () => {
const newSocket = await createWhatsAppSession(
accountId, sessionPath, onMessage, onReaction, onGroups, onReconnect, onQr, onStatus,
);
onReconnect?.(newSocket);
}, 1000);
} else {
await onStatus?.('disconnected');
logger.info('Reconnecting in 5s...');
setTimeout(async () => {
const newSocket = await createWhatsAppSession(
accountId,
sessionPath,
onMessage,
onReaction,
onGroups,
onReconnect,
accountId, sessionPath, onMessage, onReaction, onGroups, onReconnect, onQr, onStatus,
);
onReconnect?.(newSocket);
}, 5000);
}
} else if (connection === 'open') {
await onStatus?.('connected');
try {
logger.info({ accountId }, 'WhatsApp connected');
const groups = await sock.groupFetchAllParticipating();