Niels you are so smart

Lots of refactoring based on Niels' feedback, including a subscription model that can be used by consumers

This will manage the websocket connection entirely

Author: anultravioletaurora

src/utils/url/constants.ts

@@ -10,16 +10,6 @@ export const HTTP_PROTOCOL = 'http:';
 /** The https protocol string. */
 export const HTTPS_PROTOCOL = 'https:';
 
-/**
- * The websocket protocol string
- */
-export const WS_PROTOCOL = 'ws:';
-
-/**
- * The secure websocket protocol string
- */
-export const WSS_PROTOCOL = 'wss:';
-
 /** The default port for the http protocol. */
 export const HTTP_PORT = 80;
 

src/utils/url/websocket.ts

@@ -1,6 +1,5 @@
-import { Api } from "../api";
-import { SessionsMessage } from "../generated-client";
-import getWebSocketUrl from "../utils/url/websocket";
+import { Api, AUTHORIZATION_PARAMETER } from "../api";
+import { OutboundWebSocketMessage, SessionsMessage } from "../generated-client";
 
 
 /**
@@ -9,49 +8,141 @@ import getWebSocketUrl from "../utils/url/websocket";
  */
 export class WebSocketService {
 
+    /**
+     * The {@link Api} instance used to construct the websocket URL
+     */
     private api;
+
+    /**
+     * The active websocket connection, if one exists
+     */
     private socket : WebSocket | undefined
 
+    /**
+     * A map of message type subscriptions to their respective handlers
+     */
+    private subscriptions: Map<string, SocketMessageHandler<any>[]> = new Map();
+
     constructor(api : Api) {
         this.api = api
     }
 
-    private shouldOpenWebSocket(accessToken : string) { 
-
-        const socketClosed = this.socket === undefined || this.socket.readyState === WebSocket.CLOSED;
+    /**
+     * Gets the websocket URL for the given API instance
+     * 
+     * @param api The authenticated {@link Api} instance
+     * @returns The websocket URL
+     */
+    private getWebSocketUrl(api: Api) : URL {
+        return new URL(
+            api.getUri("socket", {
+                [AUTHORIZATION_PARAMETER]: api.accessToken
+            }).replace(/^http/, "ws")
+        );
+    }
 
-        // TODO: Add validation around accessToken
-        return socketClosed && accessToken.length > 0;
+    /**
+     * Gets the current status of the websocket connection
+     * @returns The {@link WebSocket.readyState} status
+     */
+    getSocketStatus() : WebSocket['readyState'] | undefined {
+        return this.socket?.readyState;
     }
 
     /**
-     * Ensures a websocket connection to the server is active by establishing a new connection
-     * if a connection doesn't already exist or the previous connection is closed.
+     * Adds message listeners for the provided message types
+     * 
+     * Adding a listener will establish a websocket connection if one does not already exist
+     * 
+     * Listeners will be automatically re-added if the connection is lost and re-established
+     * 
+     * Listeners can be removed by invoking the returned unsubscribe function
      * 
-     * Depends on the {@link Api.accessToken} being valid and
-     * populated.
+     * @param messageTypes Any array of {@link OutboundWebSocketMessage} message types to listen for
+     * @param onMessage The callback to invoke when a message is received
      * 
-     * @param onOpen An optional callback to run when the socket is opened
-     * @param onClose An optional callback to run when the socket is closed
+     * @returns A function which can be invoked to remove the added listeners
      */
-    ensureWebSocket(
-        onOpen: (e?: Event) => Promise<void>, 
-        onMessage: (e: MessageEvent<SessionsMessage>) => Promise<void>,
-        onClose: (e?: CloseEvent) => Promise<void>
-    ) {
-
-        if (this.shouldOpenWebSocket(this.api.accessToken)) {
-
-            const webSocketUrl = getWebSocketUrl(this.api);
-            
-            this.socket = new WebSocket(webSocketUrl);
-            this.socket.onopen = onOpen;
-            this.socket.onclose = onClose;
-            this.socket.onmessage = onMessage;
+    subscribe<T extends OutboundWebSocketMessage['MessageType']>(messageTypes: T[], onMessage: SocketMessageHandler<T>) {
+        if (!this.socket || this.socket.readyState !== WebSocket.OPEN) {
+            const url = this.getWebSocketUrl(this.api);
+            this.socket = new WebSocket(url.toString());
+        }
+
+        // Send startMessages and stopMessages to the server for three specific message types
+        // Sessions
+        // ActivityLogEntry 
+        // ScheduledTaskSSSSSSSSSSSSSSSSSS🐍Info 
+
+        this.socket.addEventListener('message', (event) => {
+            const message = JSON.parse(event.data) as OutboundWebSocketMessage;
+            if (messageTypes.includes(message.MessageType as T)) {
+                onMessage(message as Extract<OutboundWebSocketMessage, { MessageType: T }>);
+            }
+        });
+
+        this.socket.addEventListener('close', () => {
+            /**
+             * If the socket is closed, reopen it if there are subscriptions 
+             * and re-add the listeners
+             */ 
+            if (this.subscriptions.size > 0) {
+                const url = this.getWebSocketUrl(this.api);
+                this.socket = new WebSocket(url.toString());
+
+                this.socket.addEventListener('open', () => {
+                    for (const [messageType, handlers] of this.subscriptions.entries()) {
+                        for (const handler of handlers) {
+                            this.socket?.addEventListener('message', (event) => {
+                                const message = JSON.parse(event.data) as OutboundWebSocketMessage;
+                                if (message.MessageType === messageType) {
+                                    handler(message);
+                                }
+                            });
+                        }
+                    }
+                });
+            } 
+            // Else close and dispose
+            else {
+                this.socket?.close();
+                this.socket = undefined;
+            }
+        });
+
+        // If the last subscription has been removed, close the socket
+        
+
+        // Catalog all existing subscriptions for the given message types
+        for (const messageType of messageTypes) {
+            if (!this.subscriptions.has(messageType)) {
+                this.subscriptions.set(messageType, []);
+            }
+            this.subscriptions.get(messageType)!.push(onMessage);
+        }
+
+        // Return an unsubscribe function
+        return () => {
+            for (const messageType of messageTypes) {
+                const handlers = this.subscriptions.get(messageType);
+                if (handlers) {
+                    const index = handlers.indexOf(onMessage);
+                    if (index !== -1) {
+                        handlers.splice(index, 1);
+                    }
+                    if (handlers.length === 0) {
+                        this.subscriptions.delete(messageType);
+                    }
+                }
+            }
+
+            // If there are no more subscriptions, close the socket
+            if (this.subscriptions.size === 0) {
+                this.socket?.close();
+                this.socket = undefined;
+            }
         }
     }
+}
 
-    sendMessage(message: SessionsMessage) {
-        this.socket?.send(JSON.stringify(message));
-    } 
-}
+type SocketMessageHandler<T extends OutboundWebSocketMessage['MessageType']> = (message: Extract<OutboundWebSocketMessage, { MessageType: T }>) => void