import { WebSocket, type Data } from 'ws'; import { Unicord } from '.'; import { UnicordCommandContext } from './context'; import type { DiscordMessage, Payload, UnicordEventContext } from './types'; import { fromDiscord, Message, UnicordEventType, User } from './types'; import { exists, isNil, isObject, log } from './utils'; export class UnicordGatewayManager { private socket: WebSocket | null = null; private heartbeat_interval: number | null = null; private heartbeat_timeout: ReturnType | null = null; private sequence: number | null = null; private awaiting_heartbeat_ack = false; private session_id: string | null = null; private resume_gateway_url: string | null = null; private should_resume = false; private reconnecting = false; private reconnect_attempts = 0; // Internal state user: User | null = null; constructor(private readonly self: Unicord) {} private onOpen = () => { log('Connected to gateway'); this.reconnecting = false; this.reconnect_attempts = 0; }; private onClose = (code: number, reason: Buffer) => { log('Gateway closed:', code, reason.toString()); this.clearHeartbeat(); if (!this.reconnecting) { this.reconnect(this.canResumeFromCloseCode(code)); } }; private onError = (error: Error) => { log('Gateway error:', error); if (!this.reconnecting) { this.reconnect(true); } }; private onMessage = async (packet: Data) => { // Allocate a variable to store the parsed payload let data: Payload | null = null; try { // Attempt to parse the packet data data = JSON.parse(packet.toString()); } catch (e) { console.error(e); } if (exists(data)) { if (exists(data.s)) { this.sequence = data.s; } // NOTE: Uncomment for debugging recieved event data // log('='.repeat(25)); // log('received: ', data.op); // log('received event: ', data.t ?? 'no event'); // log('received data: ', exists(data.d) ? Object.keys(data.d) : 'no data'); // log('='.repeat(25)); // If the data is successfully parsed check the opcode const code = data.op; if (code === 0) { // Event received await this.parseEvent(data); } else if (code === 1) { // Heartbeat requested this.sendHeartbeat(); } else if (code === 7) { // Reconnect requested this.reconnect(true); } else if (code === 9) { // Invalid session. Discord sends a boolean indicating whether resume is allowed. const canResume = data.d === true; if (!canResume) { this.session_id = null; this.sequence = null; this.should_resume = false; } this.reconnect(canResume); } else if (code === 10) { // Hello if (!isObject(data.d)) { throw new Error('Invalid HELLO payload received'); } // Store the heartbeat interval this.heartbeat_interval = data.d?.heartbeat_interval; // Kick off the heartbeat loop with initial jitter this.startHeartbeat(); // Resume or identify with the gateway if (this.should_resume && exists(this.session_id)) { this.resume(); } else { this.identify(); } } else if (code === 11) { // Heartbeat ACK this.awaiting_heartbeat_ack = false; } } }; private startHeartbeat() { if (!exists(this.heartbeat_interval)) { throw new Error('Heartbeat interval not set'); } this.clearHeartbeat(); const firstHeartbeatDelay = this.heartbeat_interval * Math.random(); this.heartbeat_timeout = setTimeout(() => { this.heartbeat(); }, firstHeartbeatDelay); } private heartbeat() { if (!exists(this.heartbeat_interval)) { throw new Error('Heartbeat interval not set'); } if (this.awaiting_heartbeat_ack) { log('Missed heartbeat ACK; reconnecting gateway'); this.reconnect(true); return; } this.sendHeartbeat(); this.heartbeat_timeout = setTimeout(() => { this.heartbeat(); }, this.heartbeat_interval); } private sendHeartbeat() { this.awaiting_heartbeat_ack = true; this.send({ op: 1, d: this.sequence }); } private clearHeartbeat() { if (exists(this.heartbeat_timeout)) { clearTimeout(this.heartbeat_timeout); this.heartbeat_timeout = null; } this.awaiting_heartbeat_ack = false; } protected send = (payload: Partial) => { if (isNil(this.socket) || this.socket.readyState !== WebSocket.OPEN) { throw new Error('Socket is not connected'); } // log('='.repeat(25)); // log('sending: ', JSON.stringify(payload, null, 4)); // log('='.repeat(25)); this.socket.send(JSON.stringify(payload)); }; private identify = () => { this.should_resume = false; this.send({ op: 2, d: { token: this.self.config.token, intents: this.self.config.intents.reduce((acc, intent) => acc | intent, 0), properties: { os: process.platform, browser: 'unicord', device: 'unicord', }, }, }); }; private resume = () => { this.send({ op: 6, d: { token: this.self.config.token, session_id: this.session_id, seq: this.sequence, }, }); }; private async parseEvent(payload: Payload) { if (isNil(payload.t)) { throw new Error('Event name is required'); } // log('Event:', payload.t); // Extract the event name and data const event_name = payload.t; const event_data = payload.d; // Simple validation on the event data if (!isObject(event_data)) { throw new Error('Invalid event data received'); } if (event_name === 'READY') { this.user = User[fromDiscord](event_data.user); this.session_id = event_data.session_id; this.resume_gateway_url = event_data.resume_gateway_url; this.should_resume = true; } else if (event_name === 'RESUMED') { this.should_resume = true; } else if (event_name === 'INTERACTION_CREATE') { // TODO: Application commands/interactions need a redesigned context/hydration model before implementation. } else if (event_name === 'MESSAGE_CREATE') { await this.handleMessageCreate(payload); } else if (event_name === 'MESSAGE_UPDATE') { // TODO: setting to call commands again here on message edit // await this.handleMessageUpdate(payload); } // TODO: Make a real UnicordEventContext class this.self.emitSystem(event_name, event_data as UnicordEventContext); } private handleMessageCreate = async (payload: Payload) => { if (isNil(this.user)) { throw new Error('Cannot handle MESSAGE_CREATE event without a user'); } if (!isObject(payload.d)) { throw new Error('Invalid MESSAGE_CREATE payload received'); } // Ignore bots if (payload.d.author.bot) { return; } // Ignore self if (payload.d.author.id === this.user.id) { return; } if (typeof payload.d.content !== 'string') { return; } // Extract the content const content = payload.d.content.trim(); const hasPrefixInConfig = exists(this.self.config.prefix) && this.self.config.prefix.length > 0; const userMentionPrefix = `<@${this.user.id}>`; // Is this a command? if ( (hasPrefixInConfig && content.startsWith(this.self.config.prefix!)) || content.startsWith(userMentionPrefix) ) { // Determine the prefix const prefix = hasPrefixInConfig && content.startsWith(this.self.config.prefix!) ? this.self.config.prefix! : userMentionPrefix; // Parse command const args = content .slice(prefix.length) .trim() .split(/\s+(?=(?:[^\'"]*[\'"][^\'"]*[\'"])*[^\'"]*$)/g); // Ignore empty commands if (args.length === 0) { return; } // Extract the command name // TODO: Case insensitive config? const command = args.shift()!; // Generate context const message = Message.fromDiscord(payload.d as DiscordMessage); const context: UnicordCommandContext = new UnicordCommandContext(this.self, { message }); // Dispatch event for the command this.self.emitCommand(UnicordEventType.ChatCommands, command, context, args); } }; private reconnect(resume: boolean) { if (this.reconnecting) { return; } this.reconnecting = true; this.should_resume = resume && exists(this.session_id); this.clearHeartbeat(); if (exists(this.socket)) { this.socket.removeAllListeners(); this.socket.terminate(); this.socket = null; } const delay = this.getReconnectDelay(); log(`Reconnecting gateway in ${delay}ms${this.should_resume ? ' with resume' : ''}`); setTimeout(() => { const gatewayUrl = this.should_resume && exists(this.resume_gateway_url) ? this.resume_gateway_url : undefined; this.connect(gatewayUrl).catch((error) => { log('Gateway reconnect failed:', error); this.reconnecting = false; this.reconnect(this.should_resume); }); }, delay); } private getReconnectDelay() { const baseDelay = 1_000; const maxDelay = 60_000; const delay = Math.min(maxDelay, baseDelay * 2 ** this.reconnect_attempts); this.reconnect_attempts += 1; return delay + Math.floor(Math.random() * 1_000); } private canResumeFromCloseCode(code: number) { // Discord documents these close codes as non-resumable. return ![4004, 4010, 4011, 4012, 4013, 4014].includes(code); } async connect(gatewayUrl?: string) { // Fetch the Discord gateway URL when no explicit URL was provided. const response = gatewayUrl ? { url: gatewayUrl } : await this.self.discordAPI.request<{ url: string }>('GET', '/gateway/bot'); // Append the query parameters to the URL const url = new URL(response.url); url.searchParams.set('v', '10'); url.searchParams.set('encoding', 'json'); // Open a new WebSocket connection to the Discord gateway this.socket = new WebSocket(url); // Register event listeners this.socket.on('open', this.onOpen); this.socket.on('message', this.onMessage); this.socket.on('close', this.onClose); this.socket.on('error', this.onError); } }