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

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 { ClientWebSocketManager, ClientWebSocketEvents, ClientWebSocketManagerOptions } from './ClientWebSocket'
+import { ClientOperation, ServerOperation } from '@revanced/bot-shared'
+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({
-            op: ClientOperation.ParseText,
-            d: {
-                text,
-                id: currentId,
-            },
-        })
-
-        type CorrectPacket = Packet<ServerOperation.ParsedText>
-
-        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
+                return Promise.race([
+                    awaitPkt(ServerOperation.ParsedText),
+                    awaitPkt(ServerOperation.ParseTextFailed, this.ws.timeout + 5000),
+                ])
+                    .then(pkt => {
+                        if (pkt.op === ServerOperation.ParsedText) return pkt.d
+                        throw new Error('Failed to parse text, the API encountered an error')
+                    })
+                    .catch(() => {
+                        throw new Error('Failed to parse text, the API did not respond in time')
+                    })
+            })
     }
 
     /**
@@ -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({
-            op: ClientOperation.ParseImage,
-            d: {
-                image_url: url,
-                id: currentId,
-            },
-        })
+                return Promise.race([
+                    awaitPkt(ServerOperation.ParsedImage),
+                    awaitPkt(ServerOperation.ParseImageFailed, this.ws.timeout + 5000),
+                ])
+                    .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) => {
-            const parsedImageListener = (packet: CorrectPacket) => {
-                if (packet.d.id !== currentId) return
-                this.ws.off('parsedImage', parsedImageListener)
-                rs(packet.d)
-            }
+        return await this.ws
+            .send({
+                op: ClientOperation.TrainMessage,
+                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>) => {
-                if (packet.d.id !== currentId) return
-                this.ws.off('parseImageFailed', parseImageFailedListener)
-                rj()
-            }
-
-            this.ws.on('parsedImage', parsedImageListener)
-            this.ws.on('parseImageFailed', parseImageFailedListener)
-        })
-
-        return await promise
+                return Promise.race([
+                    awaitPkt(ServerOperation.TrainedMessage),
+                    awaitPkt(ServerOperation.TrainMessageFailed, this.ws.timeout + 5000),
+                ])
+                    .then(pkt => {
+                        if (pkt.op === ServerOperation.TrainedMessage) return
+                        throw new Error('Failed to train message, the API encountered an error')
+                    })
+                    .catch(() => {
+                        throw new Error('Failed to train message, the API did not respond in time')
+                    })
+            })
     }
 
     /**
@@ -135,14 +163,18 @@ export default class Client {
      * @param handler The event handler
      * @returns The event handler function
      */
-    once<TOpName extends keyof ClientWebSocketEvents>(
-        name: TOpName,
-        handler: ClientWebSocketEvents[TOpName],
-    ) {
+    once<TOpName extends keyof ClientWebSocketEvents>(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
-    config: Readonly<Packet<ServerOperation.Hello>['d']> | null = null
+    disconnected: false | DisconnectReason = false
+    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(() => {
-                    if (!this.ready) throw new Error('WebSocket connection timed out')
-                    this.#socket.close()
+                const timeout = setTimeout(() => {
+                    if (!this.ready) {
+                        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}`)
-                    this.#handleDisconnect(DisconnectReason.Generic)
+                    clearTimeout(timeout)
+                    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>(
-        name: TOpName,
-        handler: ClientWebSocketEvents[typeof name],
-    ) {
+    on<TOpName extends keyof ClientWebSocketEvents>(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>(
-        name: TOpName,
-        handler: ClientWebSocketEvents[typeof name],
-    ) {
+    off<TOpName extends keyof ClientWebSocketEvents>(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>(
-        name: TOpName,
-        handler: ClientWebSocketEvents[typeof name],
-    ) {
+    once<TOpName extends keyof ClientWebSocketEvents>(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.config = data
-                    this.#emitter.emit('hello', data)
-                    this.#startHeartbeating()
+                    this.#emitter.emit('hello')
+                    this.ready = true
+                    this.#emitter.emit('ready')
                     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) {
-        clearTimeout(this.#hbTimeout)
-        this.disconnected = reason
-        this.#socket.close()
+    protected _handleDisconnect(reason: DisconnectReason | number, message?: string) {
+        this.disconnected = reason in DisconnectReason ? reason : DisconnectReason.Generic
+        this.#socket?.close(reason)
         this.#socket = null!
 
-        this.#emitter.emit('disconnect', reason)
-    }
-
-    #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,
-        })
+        this.#emitter.emit('disconnect', reason, message)
     }
 
     protected _toBuffer(data: RawData) {
@@ -217,16 +190,18 @@ export interface ClientWebSocketManagerOptions {
 
 export type ClientWebSocketEventName = keyof typeof ServerOperation
 
-export type ClientWebSocketEvents = {
-    [K in Uncapitalize<ClientWebSocketEventName>]: (
-        packet: Packet<(typeof ServerOperation)[Capitalize<K>]>,
-    ) => Promise<void> | void
-} & {
-    hello: (config: NonNullable<ClientWebSocketManager['config']>) => Promise<void> | void
+type ClientWebSocketPredefinedEvents = {
+    hello: () => 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",
-    "compilerOptions": {
-        "baseUrl": ".",
-        "rootDir": "./src",
-        "outDir": "dist",
-        "module": "ESNext",
-        "composite": true,
-        "noEmit": false
-    },
-    "exclude": ["node_modules", "dist"]
+  "extends": "../../tsconfig.packages.json",
+  "compilerOptions": {
+    "baseUrl": ".",
+    "rootDir": "./src",
+    "outDir": "dist",
+    "module": "ESNext",
+    "composite": true,
+    "noEmit": false
+  },
+  "exclude": ["node_modules", "dist"]
 }