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:
@@ -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();
|
||||
});
|
||||
});
|
||||
@@ -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();
|
||||
|
||||
Reference in New Issue
Block a user