refactor(canvas): use native websocket clients

extensions/canvas/src/host/server.test.ts

@@ -27,6 +27,7 @@ type TrackingWebSocket = {
   sent: string[];
   on: (event: string, cb: () => void) => TrackingWebSocket;
   send: (message: string) => void;
+  terminate: () => void;
 };
 
 type CapturedResponse = {
@@ -482,10 +483,10 @@ describe("canvas host", () => {
     const constructorOptions: unknown[] = [];
     let connectionHandler: ((socket: TrackingWebSocket) => void) | undefined;
     class CapturingWebSocketServer {
-      on(event: string, cb: (socket: TrackingWebSocket) => void) {
-        if (event === "connection") {
-          connectionHandler = cb;
-        }
+      readonly clients = new Set<TrackingWebSocket>();
+
+      on(_event: string, cb: (socket: TrackingWebSocket) => void) {
+        connectionHandler = cb;
         return this;
       }
 
@@ -508,18 +509,9 @@ describe("canvas host", () => {
         noServer: true,
         maxPayload: canvasLiveReloadMaxInboundMessageBytes,
       });
-      const socketHandlers: string[] = [];
-      const socket: TrackingWebSocket = {
-        sent: [],
-        on: (event) => {
-          socketHandlers.push(event);
-          return socket;
-        },
-        send: vi.fn(),
-      };
-      expect(connectionHandler).toBeDefined();
-      connectionHandler?.(socket);
-      expect(socketHandlers).toEqual(expect.arrayContaining(["error", "close"]));
+      const socketOn = vi.fn();
+      connectionHandler?.({ on: socketOn } as unknown as TrackingWebSocket);
+      expect(socketOn).toHaveBeenCalledWith("error", expect.any(Function));
     } finally {
       await handler.close();
     }
@@ -593,22 +585,17 @@ describe("canvas host", () => {
 
     const watcherStart = watcherState.watchers.length;
     const TrackingWebSocketServerClass = class TrackingWebSocketServer {
-      static latestInstance: { connectionCount: number } | undefined;
       static latestSocket: TrackingWebSocket | undefined;
-      connectionCount = 0;
-      readonly handlers = new Map<string, Array<(...args: unknown[]) => void>>();
+      readonly clients = new Set<TrackingWebSocket>();
+      private connectionHandler?: (socket: TrackingWebSocket) => void;
 
-      on(event: string, cb: (...args: unknown[]) => void) {
-        const list = this.handlers.get(event) ?? [];
-        list.push(cb);
-        this.handlers.set(event, list);
+      on(_event: string, cb: (socket: TrackingWebSocket) => void) {
+        this.connectionHandler = cb;
         return this;
       }
 
-      emit(event: string, ...args: unknown[]) {
-        for (const cb of this.handlers.get(event) ?? []) {
-          cb(...args);
-        }
+      emit(_event: string, socket: TrackingWebSocket) {
+        this.connectionHandler?.(socket);
       }
 
       handleUpgrade(
@@ -620,15 +607,9 @@ describe("canvas host", () => {
         void req;
         void socket;
         void head;
-        const closeHandlers: Array<() => void> = [];
         const ws: TrackingWebSocket = {
           sent: [],
-          on: (event, handler) => {
-            if (event === "close") {
-              closeHandlers.push(handler);
-            }
-            return ws;
-          },
+          on: () => ws,
           send: (message: string) => {
             ws.sent.push(message);
             if (message === "reload") {
@@ -638,7 +619,9 @@ describe("canvas host", () => {
               resolveReload();
             }
           },
+          terminate: vi.fn(),
         };
+        this.clients.add(ws);
         TrackingWebSocketServerClass.latestSocket = ws;
         cb(ws);
       }
@@ -646,13 +629,6 @@ describe("canvas host", () => {
       close(cb?: (err?: Error) => void) {
         cb?.();
       }
-
-      constructor(..._args: unknown[]) {
-        TrackingWebSocketServerClass.latestInstance = this;
-        this.on("connection", () => {
-          this.connectionCount += 1;
-        });
-      }
     };
 
     const handler = await createTestCanvasHostHandler(dir, {
@@ -671,11 +647,6 @@ describe("canvas host", () => {
         Buffer.alloc(0),
       );
       expect(upgraded).toBe(true);
-      const latestServer = TrackingWebSocketServerClass.latestInstance;
-      if (!latestServer) {
-        throw new Error("expected Canvas host websocket server");
-      }
-      expect(latestServer.connectionCount).toBe(1);
       const ws = TrackingWebSocketServerClass.latestSocket;
       if (!ws) {
         throw new Error("expected Canvas host websocket");

extensions/canvas/src/host/server.ts

@@ -19,7 +19,7 @@ import {
   normalizeOptionalString,
 } from "openclaw/plugin-sdk/string-coerce-runtime";
 import { ensureDir, resolveUserPath } from "openclaw/plugin-sdk/text-utility-runtime";
-import { type WebSocket, WebSocketServer } from "ws";
+import { WebSocketServer } from "ws";
 import {
   CANVAS_HOST_PATH,
   CANVAS_WS_PATH,
@@ -289,24 +289,17 @@ export async function createCanvasHostHandler(
         maxPayload: CANVAS_LIVE_RELOAD_MAX_INBOUND_MESSAGE_BYTES,
       })
     : null;
-  const sockets = new Set<WebSocket>();
-  if (wss) {
-    wss.on("connection", (ws) => {
-      sockets.add(ws);
-      // ws emits error for maxPayload rejections; close handles final cleanup.
-      ws.on("error", () => {
-        sockets.delete(ws);
-      });
-      ws.on("close", () => sockets.delete(ws));
-    });
-  }
+  wss?.on("connection", (ws) => {
+    // Consume maxPayload errors; ws owns client tracking and close cleanup.
+    ws.on("error", () => {});
+  });
 
   let debounce: NodeJS.Timeout | null = null;
   const broadcastReload = () => {
-    if (!liveReload) {
+    if (!wss) {
       return;
     }
-    for (const ws of sockets) {
+    for (const ws of wss.clients) {
       try {
         ws.send("reload");
       } catch {
@@ -466,9 +459,9 @@ export async function createCanvasHostHandler(
       }
       watcherClosed = true;
       await watcher?.close().catch(() => {});
-      for (const ws of sockets) {
+      for (const ws of wss?.clients ?? []) {
         try {
-          ws.terminate?.();
+          ws.terminate();
         } catch {
           // ignore
         }