feat!: big feature changes

BREAKING CHANGES:
- Heartbeating removed
- `config.consoleLogLevel` -> `config.logLevel`

NEW FEATURES:
- Training messages
- Sequence number system
- WebSocket close codes used instead of disconnect packets

FIXES:
- Improved error handling
- Some performance improvements
- Made code more clean
- Updated dependencies

apis/websocket/config.json

@@ -4,6 +4,5 @@
   "address": "127.0.0.1",
   "port": 3000,
   "ocrConcurrentQueues": 1,
-  "clientHeartbeatInterval": 5000,
+  "logLevel": "debug"
-  "consoleLogLevel": "log"
 }

apis/websocket/config.revanced.json

@@ -4,6 +4,5 @@
   "address": "127.0.0.1",
   "port": 3000,
   "ocrConcurrentQueues": 3,
-  "clientHeartbeatInterval": 5000,
+  "logLevel": "log"
-  "consoleLogLevel": "log"
 }

apis/websocket/config.schema.json

@@ -17,12 +17,7 @@
       "type": "integer",
       "default": 1
     },
-    "clientHeartbeatInterval": {
+    "logLevel": {
-      "description": "Time in milliseconds to wait for a client to send a heartbeat packet, if no packet is received, the server will wait for `clientHeartbeatExtraTime` milliseconds before disconnecting the client",
-      "type": "integer",
-      "default": 60000
-    },
-    "consoleLogLevel": {
       "description": "The log level to print to console",
       "type": "string",
       "enum": ["error", "warn", "info", "log", "debug", "trace", "none"],

apis/websocket/docs/1_configuration.md

@@ -7,7 +7,6 @@ This is the default configuration (provided in [config.json](../config.json)):
     "address": "127.0.0.1",
     "port": 3000,
     "ocrConcurrentQueues": 1,
-    "clientHeartbeatInterval": 60000,
     "consoleLogLevel": "log"
 }
 ```
@@ -25,22 +24,18 @@ Amount of concurrent queues that can be run at a time.
 > [!WARNING]
 > Setting this too high may cause performance issues.
 
-### `config.clientHeartbeatInterval`
+### `config.logLevel`
-
-Heartbeat interval for clients. See [**💓 Heartbeating**](./3_packets.md#💓-heartbeating).
-
-### `config.consoleLogLevel`
 
 The level of logs to print to console. If the level is more important or equally important to set level, it will be forwarded to the console.
 
 The possible levels (sorted by their importance descendingly) are:
 
+-   `none` (no messages)
 -   `fatal`
 -   `error`
 -   `warn`
 -   `info`
 -   `log`
--   `trace`
 -   `debug`
 
 ## ⏭️ What's next

apis/websocket/docs/2_running.md

@@ -36,7 +36,7 @@ bun bundle
 ```
 
 The files will be placed in the `dist` directory. **Configurations and `.env` files will NOT be copied automatically.**  
-You can run these files after using a runtime, eg. `bun run .` or `node .`.
+You can run these files using the command `bun run index.js`.
 
 ## ⏭️ What's next
 

apis/websocket/docs/3_packets.md

@@ -19,15 +19,13 @@ Operation codes are numbers that communicate an action.
 
 Data fields include additional information for the server to process. They are **either an object with specific fields or just `null`**.
 
+### `packet.s` (server packets)
+
+A sequence number, exclusively for server packets. The WebSocket server contacts other APIs and they may not be reliable at all times, this makes race conditions. A sequence number cleanly solves this issue by letting the client know what the next packet sequence number would be by giving the current number.
+
 #### 📦 Schemas and constants
 
 Schemas for packets and their respective data[^1], and the list of possible operation codes[^2] can be found in the `@revanced/bot-shared` package, with typings as well.
 
 [^1]: [`@revanced/bot-shared/src/schemas/Packet.ts`](../../../packages/shared/src/schemas/Packet.ts)
 [^2]: [`@revanced/bot-shared/src/constants/Operation`](../../../packages/shared/src/constants/Operation.ts)
-
-## 💓 Heartbeating
-
-Heartbeating is a process where the client regularly send each other signals to confirm that they are still connected and functioning. If the server doesn't receive a heartbeat from the client within a specified timeframe, it assume the client has disconnected and closes the socket.
-
-You can configure the interval in the configuration file. See [**📝 Configuration > `config.clientHeartbeatInterval`**](./1_configuration.md#configclientheartbeatinterval).

apis/websocket/package.json

@@ -1,40 +1,40 @@
 {
-    "name": "@revanced/bot-websocket-api",
+  "name": "@revanced/bot-websocket-api",
-    "type": "module",
+  "type": "module",
-    "private": true,
+  "private": true,
-    "version": "0.1.0",
+  "version": "0.1.0",
-    "description": "🧦 WebSocket API server for bots assisting ReVanced",
+  "description": "🧦 WebSocket API server for bots assisting ReVanced",
-    "main": "dist/index.js",
+  "main": "dist/index.js",
-    "scripts": {
+  "scripts": {
-        "bundle": "bun build src/index.ts --outdir=dist --target=bun --minify --sourcemap=external",
+    "bundle": "bun build src/index.ts --outdir=dist --target=bun",
-        "dev": "bun run src/index.ts --watch",
+    "dev": "bun run src/index.ts --watch",
-        "build": "bun bundle",
+    "build": "bun bundle",
-        "watch": "bun dev"
+    "watch": "bun dev"
-    },
+  },
-    "repository": {
+  "repository": {
-        "type": "git",
+    "type": "git",
-        "url": "git+https://github.com/revanced/revanced-helper.git",
+    "url": "git+https://github.com/revanced/revanced-helper.git",
-        "directory": "apis/websocket"
+    "directory": "apis/websocket"
-    },
+  },
-    "author": "Palm <palmpasuthorn@gmail.com> (https://github.com/PalmDevs)",
+  "author": "Palm <palmpasuthorn@gmail.com> (https://github.com/PalmDevs)",
-    "contributors": [
+  "contributors": [
-        "Palm <palmpasuthorn@gmail.com> (https://github.com/PalmDevs)",
+    "Palm <palmpasuthorn@gmail.com> (https://github.com/PalmDevs)",
-        "ReVanced <nosupport@revanced.app> (https://github.com/revanced)"
+    "ReVanced <nosupport@revanced.app> (https://github.com/revanced)"
-    ],
+  ],
-    "license": "GPL-3.0-or-later",
+  "license": "GPL-3.0-or-later",
-    "bugs": {
+  "bugs": {
-        "url": "https://github.com/revanced/revanced-helper/issues"
+    "url": "https://github.com/revanced/revanced-helper/issues"
-    },
+  },
-    "homepage": "https://github.com/revanced/revanced-helper#readme",
+  "homepage": "https://github.com/revanced/revanced-helper#readme",
-    "dependencies": {
+  "dependencies": {
-        "@revanced/bot-shared": "workspace:*",
+    "@revanced/bot-shared": "workspace:*",
-        "@sapphire/async-queue": "^1.5.1",
+    "@sapphire/async-queue": "^1.5.2",
-        "chalk": "^5.3.0",
+    "chalk": "^5.3.0",
-        "node-wit": "^6.6.0",
+    "node-wit": "^6.6.0",
-        "tesseract.js": "^5.0.4"
+    "tesseract.js": "^5.0.5"
-    },
+  },
-    "devDependencies": {
+  "devDependencies": {
-        "@types/node-wit": "^6.0.3",
+    "@types/node-wit": "^6.1.0",
-        "typed-emitter": "^2.1.0"
+    "typed-emitter": "^2.1.0"
-    }
+  }
 }

apis/websocket/src/classes/Client.ts

@@ -2,7 +2,7 @@ import { EventEmitter } from 'events'
 import {
     ClientOperation,
     DisconnectReason,
-    Packet,
+    type Packet,
     ServerOperation,
     deserializePacket,
     isClientPacket,
@@ -17,38 +17,30 @@ export default class Client {
     id: string
     disconnected: DisconnectReason | false = false
     ready = false
+    currentSequence = 0
 
-    lastHeartbeat: number = null!
-    heartbeatInterval: number
-
-    #hbTimeout: NodeJS.Timeout = null!
     #emitter = new EventEmitter() as TypedEmitter<ClientEventHandlers>
     #socket: WebSocket
 
     constructor(options: ClientOptions) {
         this.#socket = options.socket
-        this.heartbeatInterval = options.heartbeatInterval ?? 60000
         this.id = options.id
 
-        this.#socket.on('error', () => this.forceDisconnect())
+        this.#socket.on('error', () => this.disconnect(DisconnectReason.ServerError))
-        this.#socket.on('close', () => this.forceDisconnect())
+        this.#socket.on('close', code => this._handleDisconnect(code))
-        this.#socket.on('unexpected-response', () => this.forceDisconnect())
+        this.#socket.on('unexpected-response', () => this.disconnect(DisconnectReason.InvalidPacket))
 
         this.send({
             op: ServerOperation.Hello,
-            d: {
+            d: null,
-                heartbeatInterval: this.heartbeatInterval,
-            },
         })
             .then(() => {
-                this.#listen()
+                this._listen()
-                this.#listenHeartbeat()
                 this.ready = true
                 this.#emitter.emit('ready')
             })
             .catch(() => {
-                if (this.disconnected === false) this.disconnect(DisconnectReason.ServerError)
+                this.disconnect(DisconnectReason.ServerError)
-                else this.forceDisconnect(DisconnectReason.ServerError)
             })
     }
 
@@ -64,54 +56,40 @@ export default class Client {
         this.#emitter.off(name, handler)
     }
 
-    send<TOp extends ServerOperation>(packet: Packet<TOp>) {
+    send<TOp extends ServerOperation>(packet: Omit<Packet<TOp>, 's'>, sequence?: number) {
         return new Promise<void>((resolve, reject) => {
-            try {
+            this.#throwIfDisconnected('Cannot send packet to client that has already disconnected')
-                this.#throwIfDisconnected('Cannot send packet to client that has already disconnected')
+            this.#socket.send(
-                this.#socket.send(serializePacket(packet))
+                serializePacket({ ...packet, s: sequence ?? this.currentSequence++ } as Packet<TOp>),
-                resolve()
+                err => (err ? reject(err) : resolve()),
-            } catch (e) {
+            )
-                reject(e)
-            }
         })
     }
 
-    async disconnect(reason: DisconnectReason = DisconnectReason.Generic) {
+    async disconnect(reason: DisconnectReason | number = DisconnectReason.Generic) {
         this.#throwIfDisconnected('Cannot disconnect client that has already disconnected')
 
-        try {
+        this.#socket.close(reason)
-            await this.send({ op: ServerOperation.Disconnect, d: { reason } })
+        this._handleDisconnect(reason)
-        } catch (err) {
-            throw new Error(`Cannot send disconnect reason to client ${this.id}: ${err}`)
-        } finally {
-            this.forceDisconnect(reason)
-        }
-    }
-
-    forceDisconnect(reason: DisconnectReason = DisconnectReason.Generic) {
-        if (this.disconnected !== false) return
-
-        // It's so weird because if I moved this down a few lines
-        // it would just fire the disconnect event twice because of a race condition
-        this.disconnected = reason
-        this.ready = false
-
-        if (this.#hbTimeout) clearTimeout(this.#hbTimeout)
-        this.#socket.close()
-
-        this.#emitter.emit('disconnect', reason)
     }
 
     #throwIfDisconnected(errorMessage: string) {
         if (this.disconnected !== false) throw new Error(errorMessage)
 
         if (this.#socket.readyState !== this.#socket.OPEN) {
-            this.forceDisconnect(DisconnectReason.Generic)
+            this.#socket.close(DisconnectReason.NoOpenSocket)
             throw new Error(errorMessage)
         }
     }
 
-    #listen() {
+    protected _handleDisconnect(code: number) {
+        this.disconnected = code
+        this.ready = false
+
+        this.#emitter.emit('disconnect', code)
+    }
+
+    protected _listen() {
         this.#socket.on('message', data => {
             this.#emitter.emit('message', data)
             try {
@@ -136,38 +114,6 @@ export default class Client {
         })
     }
 
-    #listenHeartbeat() {
-        this.lastHeartbeat = Date.now()
-        this.#startHeartbeatTimeout()
-
-        this.on('heartbeat', () => {
-            this.lastHeartbeat = Date.now()
-            this.#hbTimeout.refresh()
-
-            this.send({
-                op: ServerOperation.HeartbeatAck,
-                d: {
-                    nextHeartbeat: this.lastHeartbeat + this.heartbeatInterval,
-                },
-            }).catch(() => {})
-        })
-    }
-
-    #startHeartbeatTimeout() {
-        this.#hbTimeout = setTimeout(() => {
-            if (Date.now() - this.lastHeartbeat > 0) {
-                // TODO: put into config
-                // 5000 is extra time to account for latency
-                const interval = setTimeout(() => this.disconnect(DisconnectReason.TimedOut), 5000)
-
-                this.once('heartbeat', () => clearTimeout(interval))
-                // This should never happen but it did in my testing so I'm adding this just in case
-                this.once('disconnect', () => clearTimeout(interval))
-                // Technically we don't have to do this, but JUST IN CASE!
-            } else this.#hbTimeout.refresh()
-        }, this.heartbeatInterval)
-    }
-
     protected _toBuffer(data: RawData) {
         if (data instanceof Buffer) return data
         if (data instanceof ArrayBuffer) return Buffer.from(data)
@@ -178,7 +124,6 @@ export default class Client {
 export interface ClientOptions {
     id: string
     socket: WebSocket
-    heartbeatInterval?: number
 }
 
 export type ClientPacketObject<TOp extends ClientOperation> = Packet<TOp> & {

apis/websocket/src/events/index.ts

@@ -1,20 +1,33 @@
 import type { ClientOperation } from '@revanced/bot-shared'
 import type { Logger } from '@revanced/bot-shared'
-import type { Wit } from 'node-wit'
 import type { Worker as TesseractWorker } from 'tesseract.js'
-import { ClientPacketObject } from '../classes/Client'
+import type { ClientPacketObject } from '../classes/Client'
-import type { Config } from '../utils/getConfig'
+import type { Config } from '../utils/config'
 
 export { default as parseTextEventHandler } from './parseText'
 export { default as parseImageEventHandler } from './parseImage'
+export { default as trainMessageEventHandler } from './trainMessage'
 
 export type EventHandler<POp extends ClientOperation> = (
     packet: ClientPacketObject<POp>,
     context: EventContext,
 ) => void | Promise<void>
+
 export type EventContext = {
-    witClient: Wit
+    wit: {
-    tesseractWorker: TesseractWorker
+        train(text: string, label: string): Promise<void>
+        message(text: string): Promise<WitMessageResponse>
+    }
+    tesseract: TesseractWorker
     logger: Logger
     config: Config
 }
+
+export interface WitMessageResponse {
+    text: string
+    intents: Array<{
+        id: string
+        name: string
+        confidence: number
+    }>
+}

apis/websocket/src/events/parseImage.ts

@@ -1,19 +1,21 @@
-import { ClientOperation, ServerOperation } from '@revanced/bot-shared'
+import { type ClientOperation, ServerOperation } from '@revanced/bot-shared'
 import { AsyncQueue } from '@sapphire/async-queue'
 
-import type { EventHandler } from './index'
+import type { EventHandler } from '.'
 
 const queue = new AsyncQueue()
 
 const parseImageEventHandler: EventHandler<ClientOperation.ParseImage> = async (
     packet,
-    { tesseractWorker, logger, config },
+    { tesseract, logger, config },
 ) => {
     const {
         client,
-        d: { image_url: imageUrl, id },
+        d: { image_url: imageUrl },
     } = packet
 
+    const nextSeq = client.currentSequence++
+
     logger.debug(`Client ${client.id} requested to parse image from URL:`, imageUrl)
     logger.debug(`Queue currently has ${queue.remaining}/${config.ocrConcurrentQueues} items in it`)
 
@@ -23,24 +25,27 @@ const parseImageEventHandler: EventHandler<ClientOperation.ParseImage> = async (
     try {
         logger.debug(`Recognizing image from URL for client ${client.id}`)
 
-        const { data, jobId } = await tesseractWorker.recognize(imageUrl)
+        const { data, jobId } = await tesseract.recognize(imageUrl)
 
         logger.debug(`Recognized image from URL for client ${client.id} (job ${jobId}):`, data.text)
-        await client.send({
+        await client.send(
-            op: ServerOperation.ParsedImage,
+            {
-            d: {
+                op: ServerOperation.ParsedImage,
-                id,
+                d: {
-                text: data.text,
+                    text: data.text,
+                },
             },
-        })
+            nextSeq,
+        )
     } catch {
         logger.error(`Failed to parse image from URL for client ${client.id}:`, imageUrl)
-        await client.send({
+        await client.send(
-            op: ServerOperation.ParseImageFailed,
+            {
-            d: {
+                op: ServerOperation.ParseImageFailed,
-                id,
+                d: null,
             },
-        })
+            nextSeq,
+        )
     } finally {
         queue.shift()
         logger.debug(

apis/websocket/src/events/parseText.ts

@@ -1,35 +1,41 @@
-import { ClientOperation, ServerOperation } from '@revanced/bot-shared'
+import { type ClientOperation, ServerOperation } from '@revanced/bot-shared'
 
 import { inspect as inspectObject } from 'util'
 
-import type { EventHandler } from './index'
+import type { EventHandler } from '.'
 
-const parseTextEventHandler: EventHandler<ClientOperation.ParseText> = async (packet, { witClient, logger }) => {
+const parseTextEventHandler: EventHandler<ClientOperation.ParseText> = async (packet, { wit, logger }) => {
     const {
         client,
-        d: { text, id },
+        d: { text },
     } = packet
 
-    logger.debug(`Client ${client.id} requested to parse text:`, text)
+    const nextSeq = client.currentSequence++
+    const actualText = text.slice(0, 279)
+
+    logger.debug(`Client ${client.id} requested to parse text:`, actualText)
 
     try {
-        const { intents } = await witClient.message(text, {})
+        const { intents } = await wit.message(actualText)
         const intentsWithoutIds = intents.map(({ id, ...rest }) => rest)
 
-        await client.send({
+        await client.send(
-            op: ServerOperation.ParsedText,
+            {
-            d: {
+                op: ServerOperation.ParsedText,
-                id,
+                d: {
-                labels: intentsWithoutIds,
+                    labels: intentsWithoutIds,
+                },
             },
-        })
+            nextSeq,
+        )
     } catch (e) {
-        await client.send({
+        await client.send(
-            op: ServerOperation.ParseTextFailed,
+            {
-            d: {
+                op: ServerOperation.ParseTextFailed,
-                id,
+                d: null,
             },
-        })
+            nextSeq,
+        )
 
         if (e instanceof Error) logger.error(e.stack ?? e.message)
         else logger.error(inspectObject(e))

apis/websocket/src/events/trainMessage.ts

@@ -0,0 +1,43 @@
+import { type ClientOperation, ServerOperation } from '@revanced/bot-shared'
+
+import { inspect as inspectObject } from 'util'
+
+import type { EventHandler } from '.'
+
+const trainMessageEventHandler: EventHandler<ClientOperation.TrainMessage> = async (packet, { wit, logger }) => {
+    const {
+        client,
+        d: { text, label },
+    } = packet
+
+    const nextSeq = client.currentSequence++
+    const actualText = text.slice(0, 279)
+
+    logger.debug(`Client ${client.id} requested to train label ${label} with:`, actualText)
+
+    try {
+        await wit.train(actualText, label)
+        await client.send(
+            {
+                op: ServerOperation.TrainedMessage,
+                d: null,
+            },
+            nextSeq,
+        )
+
+        logger.debug(`Trained label ${label} with:`, actualText)
+    } catch (e) {
+        await client.send(
+            {
+                op: ServerOperation.TrainMessageFailed,
+                d: null,
+            },
+            nextSeq,
+        )
+
+        if (e instanceof Error) logger.error(e.stack ?? e.message)
+        else logger.error(inspectObject(e))
+    }
+}
+
+export default trainMessageEventHandler

apis/websocket/src/index.ts

@@ -1,43 +1,93 @@
-import witPkg from 'node-wit'
 import { createWorker as createTesseractWorker } from 'tesseract.js'
-const { Wit } = witPkg
 
 import { inspect as inspectObject } from 'util'
 
 import Client from './classes/Client'
 
-import { EventContext, parseImageEventHandler, parseTextEventHandler } from './events/index'
+import {
+    type EventContext,
+    type WitMessageResponse,
+    parseImageEventHandler,
+    parseTextEventHandler,
+    trainMessageEventHandler,
+} from './events'
 
 import { DisconnectReason, HumanizedDisconnectReason, createLogger } from '@revanced/bot-shared'
-import { checkEnvironment, getConfig } from './utils/index'
+import { getConfig } from './utils/config'
 
 import { createServer } from 'http'
-import { WebSocket, WebSocketServer } from 'ws'
+import { type WebSocket, WebSocketServer } from 'ws'
 
 // Load config, init logger, check environment
 
 const config = getConfig()
 const logger = createLogger({
-    level: config['consoleLogLevel'] === 'none' ? Infinity : config['consoleLogLevel'],
+    level: config.logLevel === 'none' ? Number.MAX_SAFE_INTEGER : config.logLevel,
 })
 
-checkEnvironment(logger)
+if (!process.env['NODE_ENV']) logger.warn('NODE_ENV not set, defaulting to `development`')
+const environment = (process.env['NODE_ENV'] ?? 'development') as NodeEnvironment
+
+if (!['development', 'production'].includes(environment)) {
+    logger.error('NODE_ENV is neither `development` nor `production`, unable to determine environment')
+    logger.info('Set NODE_ENV to blank to use `development` mode')
+    process.exit(1)
+}
+
+logger.info(`Running in ${environment} mode...`)
+
+if (environment === 'production' && process.env['IS_USING_DOT_ENV']) {
+    logger.warn('You seem to be using .env files, this is generally not a good idea in production...')
+}
+
+if (!process.env['WIT_AI_TOKEN']) {
+    logger.error('WIT_AI_TOKEN is not defined in the environment variables')
+    process.exit(1)
+}
 
 // Workers and API clients
 
-const tesseractWorker = await createTesseractWorker('eng')
+const tesseract = await createTesseractWorker('eng')
-const witClient = new Wit({
+const wit = {
-    accessToken: process.env['WIT_AI_TOKEN']!,
+    token: process.env['WIT_AI_TOKEN']!,
-})
+    async fetch(route: string, options?: RequestInit) {
+        const res = await fetch(`https://api.wit.ai${route}`, {
+            headers: {
+                Authorization: `Bearer ${this.token}`,
+                'Content-Type': 'application/json',
+            },
+            ...options,
+        })
+
+        if (!res.ok) throw new Error(`Failed to fetch from Wit.ai: ${res.statusText} (${res.status})`)
+
+        return await res.json()
+    },
+    message(text: string) {
+        return this.fetch(`/message?q=${encodeURIComponent(text)}&n=8`) as Promise<WitMessageResponse>
+    },
+    async train(text: string, label: string) {
+        await this.fetch('/utterances', {
+            body: JSON.stringify([
+                {
+                    text,
+                    intent: label,
+                    entities: [],
+                    traits: [],
+                },
+            ]),
+            method: 'POST',
+        })
+    },
+} as const
 
 // Server logic
 
-const clients = new Set<Client>()
+const clientMap = new WeakMap<WebSocket, Client>()
-const clientSocketMap = new WeakMap<WebSocket, Client>()
 const eventContext: EventContext = {
-    tesseractWorker,
+    tesseract,
     logger,
-    witClient,
+    wit,
     config,
 }
 
@@ -61,25 +111,24 @@ wss.on('connection', async (socket, request) => {
         const client = new Client({
             socket,
             id: `${request.socket.remoteAddress}:${request.socket.remotePort}`,
-            heartbeatInterval: config['clientHeartbeatInterval'],
         })
 
-        clientSocketMap.set(socket, client)
+        clientMap.set(socket, client)
-        clients.add(client)
 
         logger.debug(`Client ${client.id}'s instance has been added`)
-        logger.info(`New client connected (now ${clients.size} clients) with ID:`, client.id)
+        logger.info(`New client connected with ID: ${client.id}`)
 
         client.on('disconnect', reason => {
-            clients.delete(client)
+            logger.info(
-            logger.info(`Client ${client.id} disconnected because client ${HumanizedDisconnectReason[reason]}`)
+                `Client ${client.id} disconnected because client ${HumanizedDisconnectReason[reason]} (${reason})`,
+            )
         })
 
-        client.on('parseText', async packet => parseTextEventHandler(packet, eventContext))
+        client.on('parseText', packet => parseTextEventHandler(packet, eventContext))
+        client.on('parseImage', packet => parseImageEventHandler(packet, eventContext))
+        client.on('trainMessage', packet => trainMessageEventHandler(packet, eventContext))
 
-        client.on('parseImage', async packet => parseImageEventHandler(packet, eventContext))
+        if (['debug', 'trace'].includes(config.logLevel)) {
-
-        if (['debug', 'trace'].includes(config['consoleLogLevel'])) {
             logger.debug('Debug logs enabled, attaching debug events...')
 
             client.on('packet', ({ client, ...rawPacket }) =>
@@ -87,14 +136,12 @@ wss.on('connection', async (socket, request) => {
             )
 
             client.on('message', d => logger.debug(`Message from client ${client.id}:`, d))
-
-            client.on('heartbeat', () => logger.debug('Heartbeat received from client', client.id))
         }
     } catch (e) {
         if (e instanceof Error) logger.error(e.stack ?? e.message)
         else logger.error(inspectObject(e))
 
-        const client = clientSocketMap.get(socket)
+        const client = clientMap.get(socket)
 
         if (!client) {
             logger.error(
@@ -104,9 +151,6 @@ wss.on('connection', async (socket, request) => {
         }
 
         if (client.disconnected === false) client.disconnect(DisconnectReason.ServerError)
-        else client.forceDisconnect()
-
-        clients.delete(client)
 
         logger.debug(`Client ${client.id} disconnected because of an internal error`)
     }
@@ -114,7 +158,7 @@ wss.on('connection', async (socket, request) => {
 
 // Start the server
 
-server.listen(config['port'], config['address'])
+server.listen(config.port, config.address)
 
 logger.debug(`Starting with these configurations: ${inspectObject(config)}`)
 

apis/websocket/src/utils/checkEnvironment.ts

@@ -1,23 +0,0 @@
-import type { Logger } from '@revanced/bot-shared'
-
-export default function checkEnvironment(logger: Logger) {
-    if (!process.env['NODE_ENV']) logger.warn('NODE_ENV not set, defaulting to `development`')
-    const environment = (process.env['NODE_ENV'] ?? 'development') as NodeEnvironment
-
-    if (!['development', 'production'].includes(environment)) {
-        logger.error('NODE_ENV is neither `development` nor `production`, unable to determine environment')
-        logger.info('Set NODE_ENV to blank to use `development` mode')
-        process.exit(1)
-    }
-
-    logger.info(`Running in ${environment} mode...`)
-
-    if (environment === 'production' && process.env['IS_USING_DOT_ENV']) {
-        logger.warn('You seem to be using .env files, this is generally not a good idea in production...')
-    }
-
-    if (!process.env['WIT_AI_TOKEN']) {
-        logger.error('WIT_AI_TOKEN is not defined in the environment variables')
-        process.exit(1)
-    }
-}

apis/websocket/src/utils/config.ts

@@ -7,7 +7,7 @@ const configPath = resolvePath(process.cwd(), 'config.json')
 const userConfig: Partial<Config> = existsSync(configPath)
     ? (
           await import(pathToFileURL(configPath).href, {
-              assert: {
+              with: {
                   type: 'json',
               },
           })
@@ -28,10 +28,9 @@ export const defaultConfig: Config = {
     address: '127.0.0.1',
     port: 8080,
     ocrConcurrentQueues: 1,
-    clientHeartbeatInterval: 60000,
+    logLevel: 'info',
-    consoleLogLevel: 'info',
 }
 
-export default function getConfig() {
+export function getConfig() {
     return Object.assign(defaultConfig, userConfig) satisfies Config
 }

apis/websocket/src/utils/index.ts

@@ -1,2 +0,0 @@
-export { default as getConfig } from './getConfig'
-export { default as checkEnvironment } from './checkEnvironment'

apis/websocket/tsconfig.json

@@ -1,11 +1,14 @@
 {
-    "extends": "../../tsconfig.json",
+  "extends": "../../tsconfig.json",
-    "compilerOptions": {
+  "compilerOptions": {
-        "baseUrl": ".",
+    "baseUrl": ".",
-        "outDir": "dist",
+    "outDir": "dist",
-        "module": "ESNext",
+    "module": "ESNext",
-        "composite": false
+    "target": "ESNext",
-    },
+    "lib": ["ESNext"],
-    "exclude": ["node_modules", "dist"],
+    "composite": false,
-    "include": ["./*.json", "src/**/*.ts"]
+    "skipLibCheck": true
+  },
+  "exclude": ["node_modules", "dist"],
+  "include": ["./*.json", "src/**/*.ts"]
 }

packages/api/package.json

@@ -8,7 +8,7 @@
   "scripts": {
     "build": "bun bundle && bun types",
     "watch": "bunx conc --raw \"bun bundle:watch\" \"bun types:watch\"",
-    "bundle": "bun build src/index.ts --outdir=dist --sourcemap=external --target=bun --minify",
+    "bundle": "bun build src/index.ts --outdir=dist --target=bun",
     "bundle:watch": "bun run bundle --watch",
     "types": "tsc --declaration --emitDeclarationOnly",
     "types:watch": "bun types --watch --preserveWatchOutput",

packages/api/src/classes/Client.ts

@@ -1,5 +1,10 @@
-import { ClientOperation, Packet, ServerOperation } from '@revanced/bot-shared'
+import { ClientOperation, ServerOperation } from '@revanced/bot-shared'
-import { ClientWebSocketManager, ClientWebSocketEvents, ClientWebSocketManagerOptions } from './ClientWebSocket'
+import { awaitPacket } from 'src/utils/packets'
+import {
+    type ClientWebSocketEvents,
+    ClientWebSocketManager,
+    type ClientWebSocketManagerOptions,
+} from './ClientWebSocket'
 
 /**
  * The client that connects to the API.
@@ -7,7 +12,6 @@ import { ClientWebSocketManager, ClientWebSocketEvents, ClientWebSocketManagerOp
 export default class Client {
     ready = false
     ws: ClientWebSocketManager
-    #parseId = 0
 
     constructor(options: ClientOptions) {
         this.ws = new ClientWebSocketManager(options.api.websocket)
@@ -15,7 +19,7 @@ export default class Client {
             this.ready = true
         })
         this.ws.on('disconnect', () => {
-
+            this.ready = false
         })
     }
 
@@ -35,36 +39,34 @@ export default class Client {
     async parseText(text: string) {
         this.#throwIfNotReady()
 
-        const currentId = (this.#parseId++).toString()
+        return await this.ws
+            .send({
+                op: ClientOperation.ParseText,
+                d: {
+                    text,
+                },
+            })
+            .then(() => {
+                // Since we don't have heartbeats anymore, this is fine.
+                // But if we add anything similar, this will cause another race condition
+                // To fix this, we can try adding a instanced function that would return the currentSequence
+                // and it would be updated every time a "heartbeat ack" packet is received
+                const expectedNextSeq = this.ws.currentSequence + 1
+                const awaitPkt = (op: ServerOperation, timeout = this.ws.timeout) =>
+                    awaitPacket(this.ws, op, expectedNextSeq, timeout)
 
-        this.ws.send({
+                return Promise.race([
-            op: ClientOperation.ParseText,
+                    awaitPkt(ServerOperation.ParsedText),
-            d: {
+                    awaitPkt(ServerOperation.ParseTextFailed, this.ws.timeout + 5000),
-                text,
+                ])
-                id: currentId,
+                    .then(pkt => {
-            },
+                        if (pkt.op === ServerOperation.ParsedText) return pkt.d
-        })
+                        throw new Error('Failed to parse text, the API encountered an error')
-
+                    })
-        type CorrectPacket = Packet<ServerOperation.ParsedText>
+                    .catch(() => {
-
+                        throw new Error('Failed to parse text, the API did not respond in time')
-        const promise = new Promise<CorrectPacket['d']>((rs, rj) => {
+                    })
-            const parsedTextListener = (packet: CorrectPacket) => {
+            })
-                if (packet.d.id !== currentId) return
-                this.ws.off('parsedText', parsedTextListener)
-                rs(packet.d)
-            }
-
-            const parseTextFailedListener = (packet: Packet<ServerOperation.ParseTextFailed>) => {
-                if (packet.d.id !== currentId) return
-                this.ws.off('parseTextFailed', parseTextFailedListener)
-                rj()
-            }
-
-            this.ws.on('parsedText', parsedTextListener)
-            this.ws.on('parseTextFailed', parseTextFailedListener)
-        })
-
-        return await promise
     }
 
     /**
@@ -75,36 +77,62 @@ export default class Client {
     async parseImage(url: string) {
         this.#throwIfNotReady()
 
-        const currentId = (this.#parseId++).toString()
+        return await this.ws
+            .send({
+                op: ClientOperation.ParseImage,
+                d: {
+                    image_url: url,
+                },
+            })
+            .then(() => {
+                // See line 48
+                const expectedNextSeq = this.ws.currentSequence + 1
+                const awaitPkt = (op: ServerOperation, timeout = this.ws.timeout) =>
+                    awaitPacket(this.ws, op, expectedNextSeq, timeout)
 
-        this.ws.send({
+                return Promise.race([
-            op: ClientOperation.ParseImage,
+                    awaitPkt(ServerOperation.ParsedImage),
-            d: {
+                    awaitPkt(ServerOperation.ParseImageFailed, this.ws.timeout + 5000),
-                image_url: url,
+                ])
-                id: currentId,
+                    .then(pkt => {
-            },
+                        if (pkt.op === ServerOperation.ParsedImage) return pkt.d
-        })
+                        throw new Error('Failed to parse image, the API encountered an error')
+                    })
+                    .catch(() => {
+                        throw new Error('Failed to parse image, the API did not respond in time')
+                    })
+            })
+    }
 
-        type CorrectPacket = Packet<ServerOperation.ParsedImage>
+    async trainMessage(text: string, label: string) {
+        this.#throwIfNotReady()
 
-        const promise = new Promise<CorrectPacket['d']>((rs, rj) => {
+        return await this.ws
-            const parsedImageListener = (packet: CorrectPacket) => {
+            .send({
-                if (packet.d.id !== currentId) return
+                op: ClientOperation.TrainMessage,
-                this.ws.off('parsedImage', parsedImageListener)
+                d: {
-                rs(packet.d)
+                    label,
-            }
+                    text,
+                },
+            })
+            .then(() => {
+                // See line 48
+                const expectedNextSeq = this.ws.currentSequence + 1
+                const awaitPkt = (op: ServerOperation, timeout = this.ws.timeout) =>
+                    awaitPacket(this.ws, op, expectedNextSeq, timeout)
 
-            const parseImageFailedListener = (packet: Packet<ServerOperation.ParseImageFailed>) => {
+                return Promise.race([
-                if (packet.d.id !== currentId) return
+                    awaitPkt(ServerOperation.TrainedMessage),
-                this.ws.off('parseImageFailed', parseImageFailedListener)
+                    awaitPkt(ServerOperation.TrainMessageFailed, this.ws.timeout + 5000),
-                rj()
+                ])
-            }
+                    .then(pkt => {
-
+                        if (pkt.op === ServerOperation.TrainedMessage) return
-            this.ws.on('parsedImage', parsedImageListener)
+                        throw new Error('Failed to train message, the API encountered an error')
-            this.ws.on('parseImageFailed', parseImageFailedListener)
+                    })
-        })
+                    .catch(() => {
-
+                        throw new Error('Failed to train message, the API did not respond in time')
-        return await promise
+                    })
+            })
     }
 
     /**
@@ -135,14 +163,18 @@ export default class Client {
      * @param handler The event handler
      * @returns The event handler function
      */
-    once<TOpName extends keyof ClientWebSocketEvents>(
+    once<TOpName extends keyof ClientWebSocketEvents>(name: TOpName, handler: ClientWebSocketEvents[TOpName]) {
-        name: TOpName,
-        handler: ClientWebSocketEvents[TOpName],
-    ) {
         this.ws.once(name, handler)
         return handler
     }
 
+    /**
+     * Disconnects the client from the API
+     */
+    disconnect() {
+        this.ws.disconnect()
+    }
+
     #throwIfNotReady() {
         if (!this.isReady()) throw new Error('Client is not ready')
     }

packages/api/src/classes/ClientWebSocket.ts

@@ -1,8 +1,8 @@
 import { EventEmitter } from 'events'
 import {
-    ClientOperation,
+    type ClientOperation,
     DisconnectReason,
-    Packet,
+    type Packet,
     ServerOperation,
     deserializePacket,
     isServerPacket,
@@ -10,7 +10,7 @@ import {
     uncapitalize,
 } from '@revanced/bot-shared'
 import type TypedEmitter from 'typed-emitter'
-import { RawData, WebSocket } from 'ws'
+import { type RawData, WebSocket } from 'ws'
 
 /**
  * The class that handles the WebSocket connection to the server.
@@ -21,10 +21,9 @@ export class ClientWebSocketManager {
     timeout: number
 
     ready = false
-    disconnected: boolean | DisconnectReason = DisconnectReason.NeverConnected
+    disconnected: false | DisconnectReason = false
-    config: Readonly<Packet<ServerOperation.Hello>['d']> | null = null
+    currentSequence = 0
 
-    #hbTimeout: NodeJS.Timeout = null!
     #socket: WebSocket = null!
     #emitter = new EventEmitter() as TypedEmitter<ClientWebSocketEvents>
 
@@ -42,26 +41,27 @@ export class ClientWebSocketManager {
             try {
                 this.#socket = new WebSocket(this.url)
 
-                setTimeout(() => {
+                const timeout = setTimeout(() => {
-                    if (!this.ready) throw new Error('WebSocket connection timed out')
+                    if (!this.ready) {
-                    this.#socket.close()
+                        this.#socket?.close(DisconnectReason.TooSlow)
+                        throw new Error('WebSocket connection was not readied in time')
+                    }
                 }, this.timeout)
 
                 this.#socket.on('open', () => {
-                    this.disconnected = false
+                    clearTimeout(timeout)
                     this.#listen()
-                    this.ready = true
-                    this.#emitter.emit('ready')
                     rs()
                 })
 
-                this.#socket.on('error', (err) => {
+                this.#socket.on('error', err => {
+                    clearTimeout(timeout)
                     throw err
                 })
 
                 this.#socket.on('close', (code, reason) => {
-                    if (code === 1006) throw new Error(`Failed to connect to WebSocket server: ${reason}`)
+                    clearTimeout(timeout)
-                    this.#handleDisconnect(DisconnectReason.Generic)
+                    this._handleDisconnect(code, reason.toString())
                 })
             } catch (e) {
                 rj(e)
@@ -75,10 +75,7 @@ export class ClientWebSocketManager {
      * @param handler The event handler
      * @returns The event handler function
      */
-    on<TOpName extends keyof ClientWebSocketEvents>(
+    on<TOpName extends keyof ClientWebSocketEvents>(name: TOpName, handler: ClientWebSocketEvents[typeof name]) {
-        name: TOpName,
-        handler: ClientWebSocketEvents[typeof name],
-    ) {
         this.#emitter.on(name, handler)
     }
 
@@ -88,10 +85,7 @@ export class ClientWebSocketManager {
      * @param handler The event handler to remove
      * @returns The removed event handler function
      */
-    off<TOpName extends keyof ClientWebSocketEvents>(
+    off<TOpName extends keyof ClientWebSocketEvents>(name: TOpName, handler: ClientWebSocketEvents[typeof name]) {
-        name: TOpName,
-        handler: ClientWebSocketEvents[typeof name],
-    ) {
         this.#emitter.off(name, handler)
     }
 
@@ -101,10 +95,7 @@ export class ClientWebSocketManager {
      * @param handler The event handler
      * @returns The event handler function
      */
-    once<TOpName extends keyof ClientWebSocketEvents>(
+    once<TOpName extends keyof ClientWebSocketEvents>(name: TOpName, handler: ClientWebSocketEvents[typeof name]) {
-        name: TOpName,
-        handler: ClientWebSocketEvents[typeof name],
-    ) {
         this.#emitter.once(name, handler)
     }
 
@@ -126,7 +117,7 @@ export class ClientWebSocketManager {
      */
     disconnect() {
         this.#throwIfDisconnected('Cannot disconnect when already disconnected from the server')
-        this.#handleDisconnect(DisconnectReason.PlannedDisconnect)
+        this._handleDisconnect(DisconnectReason.PlannedDisconnect)
     }
 
     /**
@@ -143,22 +134,22 @@ export class ClientWebSocketManager {
 
             if (!isServerPacket(packet)) return this.#emitter.emit('invalidPacket', packet)
 
+            this.currentSequence = packet.s
             this.#emitter.emit('packet', packet)
 
             switch (packet.op) {
                 case ServerOperation.Hello: {
-                    const data = Object.freeze((packet as Packet<ServerOperation.Hello>).d)
+                    this.#emitter.emit('hello')
-                    this.config = data
+                    this.ready = true
-                    this.#emitter.emit('hello', data)
+                    this.#emitter.emit('ready')
-                    this.#startHeartbeating()
                     break
                 }
                 case ServerOperation.Disconnect:
-                    return this.#handleDisconnect((packet as Packet<ServerOperation.Disconnect>).d.reason)
+                    return this._handleDisconnect((packet as Packet<ServerOperation.Disconnect>).d.reason)
                 default:
                     return this.#emitter.emit(
                         uncapitalize(ServerOperation[packet.op] as ClientWebSocketEventName),
-                        // @ts-expect-error TypeScript doesn't know that the lines above negate the type enough
+                        // @ts-expect-error: TS at it again
                         packet,
                     )
             }
@@ -170,30 +161,12 @@ export class ClientWebSocketManager {
         if (this.#socket.readyState !== this.#socket.OPEN) throw new Error(errorMessage)
     }
 
-    #handleDisconnect(reason: DisconnectReason) {
+    protected _handleDisconnect(reason: DisconnectReason | number, message?: string) {
-        clearTimeout(this.#hbTimeout)
+        this.disconnected = reason in DisconnectReason ? reason : DisconnectReason.Generic
-        this.disconnected = reason
+        this.#socket?.close(reason)
-        this.#socket.close()
         this.#socket = null!
 
-        this.#emitter.emit('disconnect', reason)
+        this.#emitter.emit('disconnect', reason, message)
-    }
-
-    #startHeartbeating() {
-        this.on('heartbeatAck', packet => {
-            this.#hbTimeout = setTimeout(() => {
-                this.send({
-                    op: ClientOperation.Heartbeat,
-                    d: null,
-                })
-            }, packet.d.nextHeartbeat - Date.now())
-        })
-
-        // Immediately send a heartbeat so we can get when to send the next one
-        this.send({
-            op: ClientOperation.Heartbeat,
-            d: null,
-        })
     }
 
     protected _toBuffer(data: RawData) {
@@ -217,16 +190,18 @@ export interface ClientWebSocketManagerOptions {
 
 export type ClientWebSocketEventName = keyof typeof ServerOperation
 
-export type ClientWebSocketEvents = {
+type ClientWebSocketPredefinedEvents = {
-    [K in Uncapitalize<ClientWebSocketEventName>]: (
+    hello: () => Promise<void> | void
-        packet: Packet<(typeof ServerOperation)[Capitalize<K>]>,
-    ) => Promise<void> | void
-} & {
-    hello: (config: NonNullable<ClientWebSocketManager['config']>) => Promise<void> | void
     ready: () => Promise<void> | void
     packet: (packet: Packet<ServerOperation>) => Promise<void> | void
     invalidPacket: (packet: Packet) => Promise<void> | void
-    disconnect: (reason: DisconnectReason) => Promise<void> | void
+    disconnect: (reason: DisconnectReason | number, message?: string) => Promise<void> | void
 }
 
+export type ClientWebSocketEvents = {
+    [K in Exclude<Uncapitalize<ClientWebSocketEventName>, keyof ClientWebSocketPredefinedEvents>]: (
+        packet: Packet<(typeof ServerOperation)[Capitalize<K>]>,
+    ) => Promise<void> | void
+} & ClientWebSocketPredefinedEvents
+
 export type ReadiedClientWebSocketManager = RequiredProperty<InstanceType<typeof ClientWebSocketManager>>

packages/api/src/index.ts

@@ -1 +1 @@
-export * from './classes/index'
+export * from './classes'

packages/api/src/utils/packets.ts

@@ -0,0 +1,26 @@
+import type { Packet, ServerOperation } from '@revanced/bot-shared'
+import type { ClientWebSocketManager } from 'src/classes'
+
+export function awaitPacket<TOp extends ServerOperation>(
+    ws: ClientWebSocketManager,
+    op: TOp,
+    expectedSeq: number,
+    timeout = 10000,
+): Promise<Packet<TOp>> {
+    return new Promise((resolve, reject) => {
+        const timer = setTimeout(() => {
+            ws.off('packet', handler)
+            reject('Awaiting packet timed out')
+        }, timeout)
+
+        function handler(packet: Packet) {
+            if (packet.op === op && packet.s === expectedSeq) {
+                clearTimeout(timer)
+                ws.off('packet', handler)
+                resolve(packet as Packet<TOp>)
+            }
+        }
+
+        ws.on('packet', handler)
+    })
+}

packages/api/tsconfig.json

@@ -1,12 +1,12 @@
 {
-    "extends": "../../tsconfig.packages.json",
+  "extends": "../../tsconfig.packages.json",
-    "compilerOptions": {
+  "compilerOptions": {
-        "baseUrl": ".",
+    "baseUrl": ".",
-        "rootDir": "./src",
+    "rootDir": "./src",
-        "outDir": "dist",
+    "outDir": "dist",
-        "module": "ESNext",
+    "module": "ESNext",
-        "composite": true,
+    "composite": true,
-        "noEmit": false
+    "noEmit": false
-    },
+  },
-    "exclude": ["node_modules", "dist"]
+  "exclude": ["node_modules", "dist"]
 }

packages/shared/package.json

@@ -1,40 +1,40 @@
 {
-    "name": "@revanced/bot-shared",
+  "name": "@revanced/bot-shared",
-    "type": "module",
+  "type": "module",
-    "version": "0.1.0",
+  "version": "0.1.0",
-    "description": "🙌🏻 Shared components for bots assisting ReVanced",
+  "description": "🙌🏻 Shared components for bots assisting ReVanced",
-    "main": "dist/index.js",
+  "main": "dist/index.js",
-    "types": "dist/index.d.ts",
+  "types": "dist/index.d.ts",
-    "scripts": {
+  "scripts": {
-        "build": "bun bundle && bun types",
+    "build": "bun bundle && bun types",
-        "watch": "conc --raw \"bun bundle:watch\" \"bun types:watch\"",
+    "watch": "conc --raw \"bun bundle:watch\" \"bun types:watch\"",
-        "bundle": "bun build src/index.ts --outdir=dist --sourcemap=external --target=bun --minify",
+    "bundle": "bun build src/index.ts --outdir=dist --sourcemap=external --target=bun --minify",
-        "bundle:watch": "bun run bundle --watch",
+    "bundle:watch": "bun run bundle --watch",
-        "types": "tsc --declaration --emitDeclarationOnly",
+    "types": "tsc --declaration --emitDeclarationOnly",
-        "types:watch": "bun types --watch --preserveWatchOutput",
+    "types:watch": "bun types --watch --preserveWatchOutput",
-        "types:clean": "bun types --build --clean"
+    "types:clean": "bun types --build --clean"
-    },
+  },
-    "repository": {
+  "repository": {
-        "type": "git",
+    "type": "git",
-        "url": "git+https://github.com/revanced/revanced-helper.git",
+    "url": "git+https://github.com/revanced/revanced-helper.git",
-        "directory": "packages/shared"
+    "directory": "packages/shared"
-    },
+  },
-    "author": "Palm <palmpasuthorn@gmail.com> (https://github.com/PalmDevs)",
+  "author": "Palm <palmpasuthorn@gmail.com> (https://github.com/PalmDevs)",
-    "contributors": [
+  "contributors": [
-        "Palm <palmpasuthorn@gmail.com> (https://github.com/PalmDevs)",
+    "Palm <palmpasuthorn@gmail.com> (https://github.com/PalmDevs)",
-        "ReVanced <nosupport@revanced.app> (https://github.com/revanced)"
+    "ReVanced <nosupport@revanced.app> (https://github.com/revanced)"
-    ],
+  ],
-    "license": "GPL-3.0-or-later",
+  "license": "GPL-3.0-or-later",
-    "bugs": {
+  "bugs": {
-        "url": "https://github.com/revanced/revanced-helper/issues"
+    "url": "https://github.com/revanced/revanced-helper/issues"
-    },
+  },
-    "homepage": "https://github.com/revanced/revanced-helper#readme",
+  "homepage": "https://github.com/revanced/revanced-helper#readme",
-    "dependencies": {
+  "dependencies": {
-        "bson": "^6.2.0",
+    "bson": "^6.5.0",
-        "chalk": "^5.3.0",
+    "chalk": "^5.3.0",
-        "supports-color": "^9.4.0",
+    "supports-color": "^9.4.0",
-        "tracer": "^1.3.0",
+    "tracer": "^1.3.0",
-        "valibot": "^0.21.0",
+    "valibot": "^0.30.0",
-        "zod": "^3.22.4"
+    "zod": "^3.22.4"
-    }
+  }
 }

packages/shared/src/constants/DisconnectReason.ts

@@ -3,29 +3,33 @@
  */
 enum DisconnectReason {
     /**
-     * Unknown reason
+     * The client disconnected on its own (**CLIENT-ONLY**)
      */
-    Generic = 1,
+    PlannedDisconnect = 1000,
-    /**
-     * The client did not respond in time
-     */
-    TimedOut = 2,
     /**
      * The client sent an invalid packet (unserializable or invalid JSON)
      */
-    InvalidPacket = 3,
+    InvalidPacket = 1007,
     /**
      * The server has encountered an internal error
      */
-    ServerError = 4,
+    ServerError = 1011,
     /**
-     * The client had never connected to the server (**CLIENT-ONLY**)
+     * Unknown reason
      */
-    NeverConnected = 5,
+    Generic = 4000,
     /**
-     * The client disconnected on its own (**CLIENT-ONLY**)
+     * The client did not respond with a heartbeat in time
      */
-    PlannedDisconnect = 6,
+    TimedOut = 4001,
+    /**
+     * The receiving end didn't have an open socket
+     */
+    NoOpenSocket = 4003,
+    /**
+     * The client was not ready in time (**CLIENT-ONLY**)
+     */
+    TooSlow = 4002,
 }
 
 export default DisconnectReason

packages/shared/src/constants/HumanizedDisconnectReason.ts

@@ -4,12 +4,14 @@ import DisconnectReason from './DisconnectReason'
  * Humanized disconnect reasons for logs
  */
 const HumanizedDisconnectReason = {
-    [DisconnectReason.InvalidPacket]: 'has sent invalid packet',
+    [1006]: 'the receiving end had unexpectedly closed the connection',
-    [DisconnectReason.Generic]: 'has been disconnected for unknown reasons',
+    [DisconnectReason.InvalidPacket]: 'the client has sent invalid packet',
-    [DisconnectReason.TimedOut]: 'has timed out',
+    [DisconnectReason.Generic]: '(unknown reason)',
-    [DisconnectReason.ServerError]: 'has been disconnected due to an internal server error',
+    [DisconnectReason.TimedOut]: 'the client did not respond with a heartbeat in time',
-    [DisconnectReason.NeverConnected]: 'had never connected to the server',
+    [DisconnectReason.ServerError]: 'the server had an internal server error',
-    [DisconnectReason.PlannedDisconnect]: 'has disconnected on its own',
+    [DisconnectReason.TooSlow]: 'the client was not ready in time',
-} as const satisfies Record<DisconnectReason, string>
+    [DisconnectReason.PlannedDisconnect]: 'the client has disconnected on its own',
+    [DisconnectReason.NoOpenSocket]: 'the receiving end did not have an open socket',
+} as const satisfies Record<DisconnectReason | number, string>
 
 export default HumanizedDisconnectReason

packages/shared/src/constants/Operation.ts

@@ -2,33 +2,28 @@
  * Client operation codes for the gateway
  */
 export enum ClientOperation {
-    /**
-     * Client's heartbeat (to check if the connection is dead or not)
-     */
-    Heartbeat = 100,
-
     /**
      * Client's request to parse text
      */
-    ParseText = 110,
+    ParseText = 100,
     /**
      * Client's request to parse image
      */
-    ParseImage = 111,
+    ParseImage = 101,
+    /**
+     * Client's request to train a message
+     */
+    TrainMessage = 102,
 }
 
 /**
  * Server operation codes for the gateway
  */
 export enum ServerOperation {
-    /**
-     * Server's acknowledgement of a client's heartbeat
-     */
-    HeartbeatAck = 1,
     /**
      * Server's initial response to a client's connection
      */
-    Hello = 2,
+    Hello = 1,
 
     /**
      * Server's response to client's request to parse text
@@ -46,6 +41,14 @@ export enum ServerOperation {
      * Server's failure response to client's request to parse image
      */
     ParseImageFailed = 13,
+    /**
+     * Server's response to client's request to train a message
+     */
+    TrainedMessage = 14,
+    /**
+     * Server's failure response to client's request to train a message
+     */
+    TrainMessageFailed = 15,
 
     /**
      * Server's disconnect message

packages/shared/src/index.ts

@@ -1,3 +1,3 @@
-export * from './constants/index'
+export * from './constants'
-export * from './schemas/index'
+export * from './schemas'
-export * from './utils/index'
+export * from './utils'

packages/shared/src/schemas/Packet.ts

@@ -1,13 +1,12 @@
 import {
     url,
-    AnySchema,
+    type AnySchema,
-    NullSchema,
+    type NullSchema,
-    ObjectSchema,
+    type ObjectSchema,
-    Output,
+    type Output,
     array,
     enum_,
     null_,
-    number,
     object,
     parse,
     special,
@@ -30,6 +29,8 @@ export const PacketSchema = special<Packet>(input => {
         'd' in input &&
         typeof input.d === 'object'
     ) {
+        if (input.op in ServerOperation && !('s' in input && typeof input.s === 'number')) return false
+
         try {
             parse(PacketDataSchemas[input.op as Operation], input.d)
             return true
@@ -44,14 +45,8 @@ export const PacketSchema = special<Packet>(input => {
  * Schema to validate packet data for each possible operations
  */
 export const PacketDataSchemas = {
-    [ServerOperation.Hello]: object({
+    [ServerOperation.Hello]: null_(),
-        heartbeatInterval: number(),
-    }),
-    [ServerOperation.HeartbeatAck]: object({
-        nextHeartbeat: number(),
-    }),
     [ServerOperation.ParsedText]: object({
-        id: string(),
         labels: array(
             object({
                 name: string(),
@@ -60,35 +55,38 @@ export const PacketDataSchemas = {
         ),
     }),
     [ServerOperation.ParsedImage]: object({
-        id: string(),
         text: string(),
     }),
-    [ServerOperation.ParseTextFailed]: object({
+    [ServerOperation.ParseTextFailed]: null_(),
-        id: string(),
+    [ServerOperation.ParseImageFailed]: null_(),
-    }),
-    [ServerOperation.ParseImageFailed]: object({
-        id: string(),
-    }),
     [ServerOperation.Disconnect]: object({
         reason: enum_(DisconnectReason),
     }),
+    [ServerOperation.TrainedMessage]: null_(),
+    [ServerOperation.TrainMessageFailed]: null_(),
 
-    [ClientOperation.Heartbeat]: null_(),
     [ClientOperation.ParseText]: object({
-        id: string(),
         text: string(),
     }),
     [ClientOperation.ParseImage]: object({
-        id: string(),
         image_url: string([url()]),
     }),
+    [ClientOperation.TrainMessage]: object({
+        text: string(),
+        label: string(),
+    }),
 } as const satisfies Record<
     Operation,
     // biome-ignore lint/suspicious/noExplicitAny: This is a schema, it's not possible to type it
     ObjectSchema<any> | AnySchema | NullSchema
 >
 
-export type Packet<TOp extends Operation = Operation> = {
+export type Packet<TOp extends Operation = Operation> = TOp extends ServerOperation
+    ? PacketWithSequenceNumber<TOp>
+    : Omit<PacketWithSequenceNumber<TOp>, 's'>
+
+type PacketWithSequenceNumber<TOp extends Operation> = {
     op: TOp
     d: Output<(typeof PacketDataSchemas)[TOp]>
+    s: number
 }

packages/shared/src/utils/guard.ts

@@ -1,5 +1,5 @@
-import { ClientOperation, Operation, ServerOperation } from '../constants/Operation'
+import { ClientOperation, type Operation, ServerOperation } from '../constants/Operation'
-import { Packet } from '../schemas/Packet'
+import type { Packet } from '../schemas/Packet'
 
 /**
  * Checks whether a packet is trying to do the given operation
@@ -21,7 +21,7 @@ export function isClientPacket(packet: Packet): packet is Packet<ClientOperation
 }
 
 /**
- * Checks whether this packet is a server packet **(this does NOT validate the data)**
+ * Checks whether this packet is a server packet **(this does NOT validate the data or the sequence number)**
  * @param packet A packet
  * @returns Whether this packet is a server packet
  */

packages/shared/src/utils/logger.ts

@@ -1,5 +1,5 @@
 import { Chalk, supportsColor, supportsColorStderr } from 'chalk'
-import { console as uncoloredConsole, Tracer, colorConsole } from 'tracer'
+import { type Tracer, colorConsole, console as uncoloredConsole } from 'tracer'
 
 const chalk = new Chalk()
 const DefaultConfig = {
@@ -8,19 +8,16 @@ const DefaultConfig = {
         '{{message}}',
         {
             error: `${chalk.bgRedBright.whiteBright(' ERROR ')} {{message}}\n${chalk.gray('{{stack}}')}`,
-            debug: chalk.gray('DEBUG: {{message}}\n{{stack}}'),
+            debug: chalk.gray('DEBUG: {{message}}'),
-            warn: `${chalk.bgYellowBright.whiteBright(' WARN ')} ${chalk.yellowBright('{{message}}')}\n${chalk.gray(
+            warn: `${chalk.bgYellowBright.whiteBright(' WARN ')} ${chalk.yellowBright('{{message}}')}`,
-                '{{stack}}',
-            )}`,
             info: `${chalk.bgBlueBright.whiteBright(' INFO ')} ${chalk.cyanBright('{{message}}')}`,
             fatal: `${chalk.bgRedBright.whiteBright(' FATAL ')} ${chalk.redBright('{{message}}')}\n${chalk.white(
                 '{{stack}}',
             )}`,
             log: '{{message}}',
-            trace: chalk.gray('[{{timestamp}}] TRACE: {{message}}\n{{stack}}'),
         },
     ],
-    methods: ['debug', 'trace', 'log', 'info', 'warn', 'error', 'fatal'],
+    methods: ['debug', 'log', 'info', 'warn', 'error', 'fatal'],
     filters: [],
 } satisfies Tracer.LoggerConfig
 

packages/shared/src/utils/serialization.ts

@@ -1,7 +1,7 @@
 import * as BSON from 'bson'
 import { parse } from 'valibot'
-import { Operation } from '../constants/index'
+import type { Operation } from '../constants'
-import { Packet, PacketSchema } from '../schemas/index'
+import { type Packet, PacketSchema } from '../schemas'
 
 /**
  * Compresses a packet into a buffer

packages/shared/tsconfig.json

@@ -1,12 +1,12 @@
 {
-    "extends": "../../tsconfig.packages.json",
+  "extends": "../../tsconfig.packages.json",
-    "compilerOptions": {
+  "compilerOptions": {
-        "baseUrl": ".",
+    "baseUrl": ".",
-        "rootDir": "./src",
+    "rootDir": "./src",
-        "outDir": "dist",
+    "outDir": "dist",
-        "module": "ESNext",
+    "module": "ESNext",
-        "composite": true,
+    "composite": true,
-        "noEmit": false
+    "noEmit": false
-    },
+  },
-    "exclude": ["node_modules", "dist"]
+  "exclude": ["node_modules", "dist"]
 }