Fixed cookie problems

src/client.ts

@@ -3,7 +3,7 @@ import { LiveRunMove } from "./types";
 import { randomUUID } from "crypto";
 
 let socket: WebSocket | null = null;
-let cookies: string | undefined;
+let cookies: Record<string, string> | undefined;
 
 type PendingResolver = {
   resolve: (value: any) => void;
@@ -12,11 +12,11 @@ type PendingResolver = {
 
 const pending = new Map<string, PendingResolver>();
 
-export function getCookies(): string | undefined {
+export function getCookies(): Record<string, string> | undefined {
   return cookies;
 }
 
-export function setCookies(c: string) {
+export function setCookies(c: Record<string, string>) {
   cookies = c;
 }
 
@@ -26,14 +26,37 @@ function getWsOrigin(useLocalhost: boolean) {
     : "wss://dev-run.selimaj.dev/api/run";
 }
 
-function ensureSocket(useLocalhost: boolean): WebSocket {
+async function ensureSocket(useLocalhost: boolean): Promise<WebSocket> {
-  if (socket && socket.readyState === WebSocket.OPEN) {
+  if (
+    socket &&
+    socket.url === getWsOrigin(useLocalhost) &&
+    socket.readyState === WebSocket.OPEN
+  ) {
     return socket;
   }
 
-  socket = new WebSocket(getWsOrigin(useLocalhost));
+  if (socket && socket.url !== getWsOrigin(useLocalhost)) {
+    socket.close();
+  }
 
-  socket.onmessage = (e) => {
+  return new Promise((resolve, reject) => {
+    const ws = new WebSocket(getWsOrigin(useLocalhost));
+
+    ws.onopen = () => {
+      socket = ws;
+
+      console.log("Authenticating with cookies:", getCookies());
+
+      sendUnsafe(useLocalhost, { cookies: getCookies() });
+
+      resolve(ws);
+    };
+
+    ws.onerror = (err) => {
+      reject(err);
+    };
+
+    ws.onmessage = (e) => {
       const msg = JSON.parse(e.data.toString());
       const { requestId, ok, data, error } = msg;
 
@@ -47,34 +70,36 @@ function ensureSocket(useLocalhost: boolean): WebSocket {
       ok ? pendingReq.resolve(data) : pendingReq.reject(error);
     };
 
-  socket.onclose = () => {
+    ws.onclose = () => {
       socket = null;
     };
-
+  });
-  return socket;
 }
 
-function send<T>(
+async function sendUnsafe<T>(
   useLocalhost: boolean,
   payload: Record<string, any>,
 ): Promise<T> {
-  return new Promise((resolve, reject) => {
+  if (!socket) {
-    const ws = ensureSocket(useLocalhost);
+    throw new Error("WebSocket is not connected");
+  }
 
   const requestId = randomUUID();
 
+  return new Promise<T>((resolve, reject) => {
     pending.set(requestId, { resolve, reject });
-
+    socket?.send(JSON.stringify({ ...payload, requestId }));
-    ws.send(
-      JSON.stringify({
-        ...payload,
-        requestId,
-        cookies, // optional: only needed if you authenticate via cookies manually
-      }),
-    );
   });
 }
 
+async function send<T>(
+  useLocalhost: boolean,
+  payload: Record<string, any>,
+): Promise<T> {
+  await ensureSocket(useLocalhost);
+  return await sendUnsafe(useLocalhost, payload);
+}
+
 export async function addRun(
   useLocalhost: boolean,
   problem: string,