3 files changed, 37 insertions, 39 deletions

diff --git a/frontend/src/plugin-loader.tsx b/frontend/src/plugin-loader.tsx
index e5f69f1f..86592016 100644
--- a/frontend/src/plugin-loader.tsx
+++ b/frontend/src/plugin-loader.tsx
@@ -34,6 +34,7 @@ import Toaster from './toaster';
 import { VerInfo, callUpdaterMethod } from './updater';
 import { getSetting, setSetting } from './utils/settings';
 import TranslationHelper, { TranslationClass } from './utils/TranslationHelper';
+import { WSRouter } from './wsrouter';
 
 const StorePage = lazy(() => import('./components/store/Store'));
 const SettingsPage = lazy(() => import('./components/settings'));
@@ -48,6 +49,8 @@ class PluginLoader extends Logger {
   public toaster: Toaster = new Toaster();
   private deckyState: DeckyState = new DeckyState();
 
+  public ws: WSRouter = new WSRouter();
+
   public hiddenPluginsService = new HiddenPluginsService(this.deckyState);
   public notificationService = new NotificationService(this.deckyState);
 
@@ -102,9 +105,11 @@ class PluginLoader extends Logger {
 
     initFilepickerPatches();
 
-    this.getUserInfo();
+    this.ws.connect().then(() => {
+      this.getUserInfo();
 
-    this.updateVersion();
+      this.updateVersion();
+    });
   }
 
   public async getUserInfo() {
diff --git a/frontend/src/utils/settings.ts b/frontend/src/utils/settings.ts
index cadfe935..d390d7ba 100644
--- a/frontend/src/utils/settings.ts
+++ b/frontend/src/utils/settings.ts
@@ -1,24 +1,8 @@
-interface GetSettingArgs<T> {
-  key: string;
-  default: T;
-}
-
-interface SetSettingArgs<T> {
-  key: string;
-  value: T;
-}
-
 export async function getSetting<T>(key: string, def: T): Promise<T> {
-  const res = (await window.DeckyPluginLoader.callServerMethod('get_setting', {
-    key,
-    default: def,
-  } as GetSettingArgs<T>)) as { result: T };
-  return res.result;
+  const res = await window.DeckyPluginLoader.ws.call<[string, T], T>('utilities/settings/get', key, def);
+  return res;
 }
 
 export async function setSetting<T>(key: string, value: T): Promise<void> {
-  await window.DeckyPluginLoader.callServerMethod('set_setting', {
-    key,
-    value,
-  } as SetSettingArgs<T>);
+  await window.DeckyPluginLoader.ws.call<[string, T], void>('utilities/settings/set', key, value);
 }
diff --git a/frontend/src/wsrouter.ts b/frontend/src/wsrouter.ts
index b6437568..3a36b5b0 100644
--- a/frontend/src/wsrouter.ts
+++ b/frontend/src/wsrouter.ts
@@ -41,10 +41,11 @@ interface PromiseResolver<T> {
   promise: Promise<T>;
 }
 
-class WSRouter extends Logger {
+export class WSRouter extends Logger {
   routes: Map<string, (...args: any) => any> = new Map();
   runningCalls: Map<number, PromiseResolver<any>> = new Map();
   ws?: WebSocket;
+  connectPromise?: Promise<void>;
   // Used to map results and errors to calls
   reqId: number = 0;
   constructor() {
@@ -52,30 +53,35 @@ class WSRouter extends Logger {
   }
 
   connect() {
-    this.ws = new WebSocket('ws://127.0.0.1:1337/ws');
-
-    this.ws.addEventListener('message', this.onMessage.bind(this));
-    this.ws.addEventListener('close', this.onError.bind(this));
-    this.ws.addEventListener('message', this.onError.bind(this));
+    return (this.connectPromise = new Promise<void>((resolve) => {
+      // Auth is a query param as JS WebSocket doesn't support headers
+      this.ws = new WebSocket(`ws://127.0.0.1:1337/ws?auth=${window.deckyAuthToken}`);
+
+      this.ws.addEventListener('open', () => {
+        this.debug('WS Connected');
+        resolve();
+        delete this.connectPromise;
+      });
+      this.ws.addEventListener('message', this.onMessage.bind(this));
+      this.ws.addEventListener('close', this.onError.bind(this));
+      // this.ws.addEventListener('error', this.onError.bind(this));
+    }));
   }
 
   createPromiseResolver<T>(): PromiseResolver<T> {
-    let resolver: PromiseResolver<T>;
+    let resolver: Partial<PromiseResolver<T>> = {};
     const promise = new Promise<T>((resolve, reject) => {
-      resolver = {
-        promise,
-        resolve,
-        reject,
-      };
-      this.debug('Created new PromiseResolver');
+      resolver.resolve = resolve;
+      resolver.reject = reject;
     });
-    this.debug('Returning new PromiseResolver');
+    resolver.promise = promise;
     // The promise will always run first
     // @ts-expect-error 2454
     return resolver;
   }
 
-  write(data: Message) {
+  async write(data: Message) {
+    if (this.connectPromise) await this.connectPromise;
     this.ws?.send(JSON.stringify(data));
   }
 
@@ -129,9 +135,9 @@ class WSRouter extends Logger {
     } catch (e) {
       this.error('Error parsing WebSocket message', e);
     }
-    this.call<[number, number], string>('methodName', 1, 2);
   }
 
+  // this.call<[number, number], string>('methodName', 1, 2);
   call<Args extends any[] = any[], Return = void>(route: string, ...args: Args): Promise<Return> {
     const resolver = this.createPromiseResolver<Return>();
 
@@ -139,12 +145,15 @@ class WSRouter extends Logger {
 
     this.runningCalls.set(id, resolver);
 
+    this.debug(`Calling PY method ${route} with args`, args);
+
     this.write({ type: MessageType.CALL, route, args, id });
 
     return resolver.promise;
   }
 
-  onError(error: any) {
-    this.error('WS ERROR', error);
+  async onError(error: any) {
+    this.error('WS DISCONNECTED', error);
+    await this.connect();
   }
 }