use built-in global ws

Author: dlarocque

common/api-review/ai.api.md

@@ -433,6 +433,10 @@ export class GenerativeModel extends AIModel {
     toolConfig?: ToolConfig;
     // (undocumented)
     tools?: Tool[];
+    // Warning: (ae-forgotten-export) The symbol "WebSocketHandler" needs to be exported by the entry point index.d.ts
+    //
+    // (undocumented)
+    ws: WebSocketHandler;
 }
 
 // @public

package.json

@@ -159,6 +159,7 @@
     "typescript": "5.5.4",
     "watch": "1.0.2",
     "webpack": "5.98.0",
+    "ws": "8.18.3",
     "yargs": "17.7.2"
   }
 }

packages/ai/package.json

@@ -63,8 +63,7 @@
     "rollup": "2.79.2",
     "rollup-plugin-replace": "2.2.0",
     "rollup-plugin-typescript2": "0.36.0",
-    "typescript": "5.5.4",
-    "ws": "8.18.3"
+    "typescript": "5.5.4"
   },
   "repository": {
     "directory": "packages/ai",

packages/ai/src/platform/browser/websocket.test.ts

@@ -270,4 +270,4 @@ describe('BrowserWebSocketHandler', () => {
       expect(mockWebSocket.readyState).to.equal(MockBrowserWebSocket.CLOSED);
     });
   });
-});
+});

packages/ai/src/platform/browser/websocket.ts

@@ -20,16 +20,16 @@ import { AIErrorCode } from '../../types';
 import { WebSocketHandler } from '../websocket';
 
 export function createWebSocketHandler(): WebSocketHandler {
-  if (typeof WebSocket !== 'undefined') {
-    return new BrowserWebSocketHandler();
-  } else {
+  if (typeof WebSocket === 'undefined') {
     throw new AIError(
       AIErrorCode.UNSUPPORTED,
       'The WebSocket API is not available in this browser-like environment. ' +
         'The "Live" feature is not supported here. It is supported in ' +
-        'standard browser windows, Web Workers with WebSocket support, and Node >= 22.'
+        'modern browser windows, Web Workers with WebSocket support, and Node >= 22.'
     );
   }
+
+  return new BrowserWebSocketHandler();
 }
 
 /**
@@ -43,17 +43,7 @@ export class BrowserWebSocketHandler implements WebSocketHandler {
 
   connect(url: string): Promise<void> {
     return new Promise((resolve, reject) => {
-      try {
-        this.ws = new WebSocket(url);
-      } catch (e) {
-        return reject(
-          new AIError(
-            AIErrorCode.ERROR,
-            `Internal Error: Invalid WebSocket URL: ${url}`
-          )
-        );
-      }
-
+      this.ws = new WebSocket(url);
       this.ws.addEventListener('open', () => resolve(), { once: true });
       this.ws.addEventListener(
         'error',
@@ -136,15 +126,22 @@ export class BrowserWebSocketHandler implements WebSocketHandler {
 
   close(code?: number, reason?: string): Promise<void> {
     return new Promise(resolve => {
+      if (!this.ws) {
+        return resolve();
+      }
+
+      this.ws.addEventListener('close', () => resolve(), { once: true });
+      // Calling 'close' during these states results in an error.
       if (
-        !this.ws ||
         this.ws.readyState === WebSocket.CLOSED ||
-        this.ws.readyState === WebSocket.CLOSING
+        this.ws.readyState === WebSocket.CONNECTING
       ) {
         return resolve();
       }
-      this.ws.addEventListener('close', () => resolve(), { once: true });
-      this.ws.close(code, reason);
+
+      if (this.ws.readyState !== WebSocket.CLOSING) {
+        this.ws.close(code, reason);
+      }
     });
   }
 }

packages/ai/src/platform/node/websocket.test.ts

@@ -23,15 +23,14 @@ import { TextEncoder } from 'util';
 import { MockWebSocketServer } from '../../../test-utils/mock-websocket-server';
 import { WebSocketHandler } from '../websocket';
 import { NodeWebSocketHandler } from './websocket';
-import { AIError } from '../../errors';
 
 use(sinonChai);
 use(chaiAsPromised);
 
 const TEST_PORT = 9003;
 const TEST_URL = `ws://localhost:${TEST_PORT}`;
 
-describe('NodeWebSocketHandler (Integration Tests)', () => {
+describe('NodeWebSocketHandler', () => {
   let server: MockWebSocketServer;
   let handler: WebSocketHandler;
 
@@ -66,12 +65,9 @@ describe('NodeWebSocketHandler (Integration Tests)', () => {
       expect(server.clients.size).to.equal(1);
     });
 
-    it('should reject if the connection fails (e.g., wrong port)', async () => {
-      const wrongPortUrl = `ws://localhost:${TEST_PORT + 1}`;
-      await expect(handler.connect(wrongPortUrl)).to.be.rejectedWith(
-        AIError,
-        /Failed to establish WebSocket connection/
-      );
+    it('should reject if the connection fails', async () => {
+      const wrongPortUrl = `ws://wrongUrl:9000`;
+      await expect(handler.connect(wrongPortUrl)).to.be.rejected;
     });
   });
 
@@ -144,4 +140,4 @@ describe('NodeWebSocketHandler (Integration Tests)', () => {
       await expect(consumerPromise).to.be.fulfilled;
     });
   });
-});
+});

packages/ai/src/platform/node/websocket.ts

@@ -15,7 +15,6 @@
  * limitations under the License.
  */
 
-// import { WebSocket, MessageEvent } from 'ws'; // External dependency on native Node module
 import { AIError } from '../../errors';
 import { AIErrorCode } from '../../types';
 import { WebSocketHandler } from '../websocket';
@@ -27,41 +26,41 @@ export function createWebSocketHandler(): WebSocketHandler {
       throw new AIError(
         AIErrorCode.UNSUPPORTED,
         `The "Live" feature is being used in a Node environment, but the ` +
-          `runtime version is ${process.versions.node}. This feature requires Node >= 22` +
+          `runtime version is ${process.versions.node}. This feature requires Node >= 22 ` +
           `for native WebSocket support.`
       );
+    } else if (typeof WebSocket === 'undefined') {
+      throw new AIError(
+        AIErrorCode.UNSUPPORTED,
+        `The "Live" feature is being used in a Node environment that does not offer the ` +
+          `'WebSocket' API in the global scope.`
+      );
     }
+
     return new NodeWebSocketHandler();
   } else {
     throw new AIError(
       AIErrorCode.UNSUPPORTED,
       'The "Live" feature is not supported in this Node-like environment. It is supported in ' +
-        'standard browser windows, Web Workers with WebSocket support, and Node >= 22.'
+        'modern browser windows, Web Workers with WebSocket support, and Node >= 22.'
     );
   }
 }
 
 /**
  * A WebSocketHandler implementation for Node >= 22.
  *
+ * Node 22 is the minimum version that offers the built-in global `WebSocket` API.
+ *
  * @internal
  */
 export class NodeWebSocketHandler implements WebSocketHandler {
-  private ws?: import('ws').WebSocket;
+  private ws?: WebSocket;
 
   async connect(url: string): Promise<void> {
     return new Promise(async (resolve, reject) => {
-      const { WebSocket } = await import('ws');
-      try {
-        this.ws = new WebSocket(url);
-      } catch (e) {
-        return reject(
-          new AIError(
-            AIErrorCode.ERROR,
-            `Internal Error: Invalid WebSocket URL: ${url}`
-          )
-        );
-      }
+      this.ws = new WebSocket(url);
+      this.ws.binaryType = 'blob';
       this.ws!.addEventListener('open', () => resolve(), { once: true });
       this.ws!.addEventListener(
         'error',
@@ -78,7 +77,7 @@ export class NodeWebSocketHandler implements WebSocketHandler {
   }
 
   send(data: string | ArrayBuffer): void {
-    if (!this.ws || this.ws.readyState !== this.ws.OPEN) {
+    if (!this.ws || this.ws.readyState !== WebSocket.OPEN) {
       throw new AIError(AIErrorCode.REQUEST_ERROR, 'WebSocket is not open.');
     }
     this.ws.send(data);
@@ -96,36 +95,25 @@ export class NodeWebSocketHandler implements WebSocketHandler {
     let resolvePromise: (() => void) | null = null;
     let isClosed = false;
 
-    const messageListener = (event: import('ws').MessageEvent): void => {
-      let textData: string;
-
-      if (typeof event.data === 'string') {
-        textData = event.data;
-      } else if (
-        event.data instanceof Buffer ||
-        event.data instanceof ArrayBuffer ||
-        event.data instanceof Uint8Array
-      ) {
-        const decoder = new TextDecoder();
-        textData = decoder.decode(event.data);
+    const messageListener = async (event: MessageEvent): Promise<void> => {
+      if (event.data instanceof Blob) {
+        try {
+          const obj = JSON.parse(await event.data.text()) as unknown;
+          messageQueue.push(obj);
+          if (resolvePromise) {
+            resolvePromise();
+            resolvePromise = null;
+          }
+        } catch (e) {
+          console.warn('Failed to parse WebSocket message to JSON:', e);
+        }
       } else {
         throw new AIError(
           AIErrorCode.PARSE_FAILED,
           `Failed to parse WebSocket response to JSON. ` +
-            `Expected data to be string, Buffer, ArrayBuffer, or Uint8Array, but was ${typeof event.data}.`
+            `Expected data to be a Blob, but was ${typeof event.data}.`
         );
       }
-
-      try {
-        const parsedObject = JSON.parse(textData);
-        messageQueue.push(parsedObject);
-        if (resolvePromise) {
-          resolvePromise();
-          resolvePromise = null;
-        }
-      } catch (e) {
-        console.warn('Failed to parse WebSocket message to JSON:', textData, e);
-      }
     };
 
     const closeListener = (): void => {
@@ -155,15 +143,22 @@ export class NodeWebSocketHandler implements WebSocketHandler {
 
   close(code?: number, reason?: string): Promise<void> {
     return new Promise(resolve => {
+      if (!this.ws) {
+        return resolve();
+      }
+
+      this.ws.addEventListener('close', () => resolve(), { once: true });
+      // Calling 'close' during these states results in an error.
       if (
-        !this.ws ||
-        this.ws.readyState === this.ws.CLOSED ||
-        this.ws.readyState === this.ws.CLOSING
+        this.ws.readyState === WebSocket.CLOSED ||
+        this.ws.readyState === WebSocket.CONNECTING
       ) {
         return resolve();
       }
-      this.ws.addEventListener('close', () => resolve(), { once: true });
-      this.ws.close(code, reason);
+
+      if (this.ws.readyState !== WebSocket.CLOSING) {
+        this.ws.close(code, reason);
+      }
     });
   }
 }

packages/ai/test-utils/mock-websocket-server.ts

@@ -52,7 +52,7 @@ export class MockWebSocketServer {
   broadcast(message: string | Buffer): void {
     for (const client of this.clients) {
       if (client.readyState === WebSocket.OPEN) {
-        client.send(message);
+        client.send(message, { binary: true });
       }
     }
   }