feat(js,react): disconnect usage limited users (#356)

VojtechVidra · web-flow · commit 0ba72d8c8425 · 2025-04-07T15:51:48.000+02:00
* feat(js,react): disconnect usage limited users

* fix: websocket test in firefox
diff --git a/workspaces/e2e/tests/websocket.spec.ts b/workspaces/e2e/tests/websocket.spec.ts
@@ -13,49 +13,94 @@ const getBlock = (): Block => ({
   propertyMeta: [],
 });
 
-let ws: WebSocketRoute | null = null;
-test.beforeEach(async ({ page }) => {
-  await page.routeWebSocket(
-    (url) => url.pathname === "/ws/sdk/block-updates",
-    (_ws) => {
-      ws = _ws;
-    },
-  );
-});
-
 const run = (packageName: string) => {
-  test(`${packageName} - should display block that is received through websocket`, async ({
-    page,
-  }) => {
-    const block = getBlock();
-    await page.route("**/v2/sdk/blocks", (route) => {
-      route.fulfill({ json: { blocks: [] } });
+  test.describe(`mocked websocket`, () => {
+    let ws: WebSocketRoute | null = null;
+    test.beforeEach(async ({ page }) => {
+      await page.routeWebSocket(
+        (url) => url.pathname === "/ws/sdk/block-updates",
+        (_ws) => {
+          ws = _ws;
+        },
+      );
+    });
+
+    test(`${packageName} - should display block that is received through websocket`, async ({
+      page,
+    }) => {
+      const block = getBlock();
+      await page.route("**/v2/sdk/blocks", (route) => {
+        route.fulfill({ json: { blocks: [] } });
+      });
+      await page.goto(`/${packageName}.html`);
+      await expect(page.locator("h1")).toBeVisible();
+      await expect(page.getByText("Hello world", { exact: true })).toBeHidden();
+      const payload: BlockUpdatesPayload = {
+        exitedBlockIds: [],
+        updatedBlocks: [block],
+      };
+      await (ws as unknown as WebSocketRoute).send(JSON.stringify(payload));
+      await expect(page.getByText("Hello world", { exact: true })).toBeVisible();
+    });
+    test(`${packageName} - should exit block that is received through websocket`, async ({
+      page,
+    }) => {
+      const block = getBlock();
+      await page.route("**/v2/sdk/blocks", (route) => {
+        route.fulfill({ json: { blocks: [block] } });
+      });
+      await page.goto(`/${packageName}.html`);
+      await expect(page.getByText("Hello world", { exact: true })).toBeVisible();
+      const payload: BlockUpdatesPayload = {
+        exitedBlockIds: [block.id],
+        updatedBlocks: [],
+      };
+      await (ws as unknown as WebSocketRoute).send(JSON.stringify(payload));
+      await expect(page.getByText("Hello world", { exact: true })).toBeHidden();
     });
-    await page.goto(`/${packageName}.html`);
-    await expect(page.locator("h1")).toBeVisible();
-    await expect(page.getByText("Hello world", { exact: true })).toBeHidden();
-    const payload: BlockUpdatesPayload = {
-      exitedBlockIds: [],
-      updatedBlocks: [block],
-    };
-    await (ws as unknown as WebSocketRoute).send(JSON.stringify(payload));
-    await expect(page.getByText("Hello world", { exact: true })).toBeVisible();
   });
-  test(`${packageName} - should exit block that is received through websocket`, async ({
-    page,
-  }) => {
-    const block = getBlock();
-    await page.route("**/v2/sdk/blocks", (route) => {
-      route.fulfill({ json: { blocks: [block] } });
+
+  test.describe(`real websocket`, () => {
+    test(`${packageName} - should disconnect from websocket if the user is usage limited`, async ({
+      page,
+    }) => {
+      await page.route("**/v2/sdk/blocks", (route) => {
+        route.fulfill({ json: { blocks: [], meta: { usage_limited: true } } });
+      });
+      const wsPromise = page.waitForEvent("websocket");
+      await page.goto(`/${packageName}.html`);
+      const websocket = await wsPromise;
+
+      let wsWasClosed = false;
+      websocket.on("close", () => {
+        wsWasClosed = true;
+      });
+      await page.waitForRequest((req) => {
+        return req.url() === "https://api.flows-cloud.com/v2/sdk/blocks";
+      });
+
+      await expect(() => expect(wsWasClosed).toBe(true)).toPass();
+    });
+    test(`${packageName} - should keep websocket connection alive`, async ({ page }) => {
+      await page.route("**/v2/sdk/blocks", (route) => {
+        route.fulfill({ json: { blocks: [] } });
+      });
+      const wsPromise = page.waitForEvent("websocket");
+      await page.goto(`/${packageName}.html`);
+      const websocket = await wsPromise;
+
+      let wsWasClosed = false;
+      websocket.on("close", () => {
+        wsWasClosed = true;
+      });
+      await page.waitForRequest((req) => {
+        return req.url() === "https://api.flows-cloud.com/v2/sdk/blocks";
+      });
+
+      await new Promise((res) => setTimeout(res, 500));
+
+      await expect(() => expect(wsWasClosed).toBe(false)).toPass();
     });
-    await page.goto(`/${packageName}.html`);
-    await expect(page.getByText("Hello world", { exact: true })).toBeVisible();
-    const payload: BlockUpdatesPayload = {
-      exitedBlockIds: [block.id],
-      updatedBlocks: [],
-    };
-    await (ws as unknown as WebSocketRoute).send(JSON.stringify(payload));
-    await expect(page.getByText("Hello world", { exact: true })).toBeHidden();
   });
 };
 
diff --git a/workspaces/js/src/lib/blocks.ts b/workspaces/js/src/lib/blocks.ts
@@ -1,6 +1,6 @@
 import { type BlockUpdatesPayload, getApi, log, type UserProperties } from "@flows/shared";
 import { blocks } from "../store";
-import { websocket } from "./websocket";
+import { type Disconnect, websocket } from "./websocket";
 import { packageAndVersion } from "./constants";
 
 interface Props {
@@ -11,6 +11,8 @@ interface Props {
   userProperties?: UserProperties;
 }
 
+let disconnect: Disconnect | null = null;
+
 export const connectToWebsocketAndFetchBlocks = (props: Props): void => {
   const { environment, organizationId, userId, apiUrl } = props;
   const params = { environment, organizationId, userId };
@@ -24,6 +26,8 @@ export const connectToWebsocketAndFetchBlocks = (props: Props): void => {
       .getBlocks({ ...params, userProperties: props.userProperties })
       .then((res) => {
         blocks.value = res.blocks;
+        // Disconnect if the user is usage limited
+        if (res.meta?.usage_limited) disconnect?.();
       })
       .catch((err: unknown) => {
         log.error("Failed to load blocks", err);
@@ -41,5 +45,9 @@ export const connectToWebsocketAndFetchBlocks = (props: Props): void => {
     ];
   };
 
-  websocket({ url: wsUrl, onMessage, onOpen: fetchBlocks });
+  // Disconnect previous connection if it exists
+  disconnect?.();
+
+  const websocketResult = websocket({ url: wsUrl, onMessage, onOpen: fetchBlocks });
+  disconnect = websocketResult.disconnect;
 };
diff --git a/workspaces/js/src/lib/websocket.ts b/workspaces/js/src/lib/websocket.ts
@@ -4,19 +4,37 @@ interface Props {
   onOpen: () => void;
 }
 
+export type Disconnect = () => void;
+
 // TODO: implement better automatic reconnection
-export const websocket = (props: Props): void => {
-  const connect = (): void => {
+export const websocket = (props: Props): { disconnect: Disconnect } => {
+  const connect = (): { disconnect: Disconnect } => {
     const socket = new WebSocket(props.url);
 
-    socket.addEventListener("message", props.onMessage);
-    socket.addEventListener("open", props.onOpen);
-    socket.addEventListener("close", () => {
+    const handleClose = (): void => {
       setTimeout(() => {
         connect();
       }, 2000);
-    });
+    };
+
+    socket.addEventListener("message", props.onMessage);
+    socket.addEventListener("open", props.onOpen);
+    socket.addEventListener("close", handleClose);
+
+    const disconnect = (): void => {
+      socket.removeEventListener("message", props.onMessage);
+      socket.removeEventListener("open", props.onOpen);
+      socket.removeEventListener("close", handleClose);
+
+      if (socket.readyState === WebSocket.CONNECTING) {
+        socket.addEventListener("open", () => {
+          socket.close();
+        });
+      } else socket.close();
+    };
+
+    return { disconnect };
   };
 
-  connect();
+  return connect();
 };
diff --git a/workspaces/react/src/hooks/use-blocks.ts b/workspaces/react/src/hooks/use-blocks.ts
@@ -33,6 +33,7 @@ export const useBlocks = ({
   userProperties,
 }: Props): Return => {
   const [blocks, setBlocks] = useState<Block[]>([]);
+  const [usageLimited, setUsageLimited] = useState(false);
 
   const params = useMemo(
     () => ({ environment, organizationId, userId }),
@@ -49,16 +50,18 @@ export const useBlocks = ({
       .getBlocks({ ...params, userProperties: userPropertiesRef.current })
       .then((res) => {
         setBlocks(res.blocks);
+        if (res.meta?.usage_limited) setUsageLimited(true);
       })
       .catch((err: unknown) => {
         log.error("Failed to load blocks", err);
       });
   }, [apiUrl, params]);
 
-  const url = useMemo(() => {
+  const websocketUrl = useMemo(() => {
+    if (usageLimited) return;
     const baseUrl = apiUrl.replace("https://", "wss://").replace("http://", "ws://");
     return `${baseUrl}/ws/sdk/block-updates?${new URLSearchParams(params).toString()}`;
-  }, [apiUrl, params]);
+  }, [apiUrl, params, usageLimited]);
 
   const onMessage = useCallback((event: MessageEvent<unknown>) => {
     // TODO: add debug logging
@@ -73,7 +76,7 @@ export const useBlocks = ({
       ...data.updatedBlocks,
     ]);
   }, []);
-  useWebsocket({ url, onMessage, onOpen: fetchBlocks });
+  useWebsocket({ url: websocketUrl, onMessage, onOpen: fetchBlocks });
 
   // Log error about slottable blocks without slotId
   useEffect(() => {
diff --git a/workspaces/react/src/hooks/use-websocket.ts b/workspaces/react/src/hooks/use-websocket.ts
@@ -1,7 +1,7 @@
 import { useCallback, useEffect, useRef, useState } from "react";
 
 interface Props {
-  url: string;
+  url?: string;
   onMessage: (event: MessageEvent<unknown>) => void;
   onOpen?: () => void;
 }
@@ -24,12 +24,14 @@ export const useWebsocket = ({ url, onMessage, onOpen }: Props): void => {
     onMessageRef.current(event);
   }, []);
 
-  const connect = useCallback(() => {
+  const connect = useCallback((): (() => void) | undefined => {
     if (cleanupRef.current) {
       cleanupRef.current();
       cleanupRef.current = undefined;
     }
 
+    if (!url) return;
+
     const socket = new WebSocket(url);
     setWs(socket);
 
@@ -66,7 +68,7 @@ export const useWebsocket = ({ url, onMessage, onOpen }: Props): void => {
   useEffect(() => {
     const cleanup = connect();
     return () => {
-      cleanup();
+      cleanup?.();
     };
   }, [connect]);
 
diff --git a/workspaces/shared/src/api.ts b/workspaces/shared/src/api.ts
@@ -28,8 +28,13 @@ interface GetBlocksRequest {
   userProperties?: Record<string, unknown>;
 }
 
+interface BlockResponseMeta {
+  usage_limited?: boolean;
+}
+
 interface BlocksResponse {
   blocks: Block[];
+  meta?: BlockResponseMeta;
 }
 
 export interface EventRequest {

Original file line number	Diff line number	Diff line change
`@@ -28,8 +28,13 @@ interface GetBlocksRequest {`
`28`	`28`	`userProperties?: Record<string, unknown>;`
`29`	`29`	`}`
`30`	`30`
	`31`	`+interface BlockResponseMeta {`
	`32`	`+ usage_limited?: boolean;`
	`33`	`+}`
	`34`	`+`
`31`	`35`	`interface BlocksResponse {`
`32`	`36`	`blocks: Block[];`
	`37`	`+ meta?: BlockResponseMeta;`
`33`	`38`	`}`
`34`	`39`
`35`	`40`	`export interface EventRequest {`