Add WebSocket handler for Browser and Node

config/tsconfig.base.json

@@ -9,7 +9,7 @@
       "es2020",
       "esnext.WeakRef",
     ],
-    "module": "ES2020",
+    "module": "ES2015",
     "moduleResolution": "node",
     "resolveJsonModule": true,
     "esModuleInterop": true,

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

@@ -39,7 +39,7 @@
     "test:ci": "yarn testsetup && node ../../scripts/run_tests_in_ci.js -s test",
     "test:skip-clone": "karma start",
     "test:browser": "yarn testsetup && karma start",
-    "test:node": "TS_NODE_COMPILER_OPTIONS='{\"module\":\"commonjs\"}' mocha --require ts-node/register --require src/index.node.ts src/**/*.test.ts --config ../../config/mocharc.node.js",
+    "test:node": "TS_NODE_COMPILER_OPTIONS='{\"module\":\"commonjs\"}' mocha --require ts-node/register --require src/index.node.ts 'src/**/*.test.ts' --config ../../config/mocharc.node.js",
     "test:integration": "karma start --integration",
     "api-report": "api-extractor run --local --verbose",
     "typings:public": "node ../../scripts/build/use_typings.js ./dist/ai-public.d.ts",
@@ -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/rollup.config.js

@@ -15,15 +15,16 @@
  * limitations under the License.
  */
 
+import alias from '@rollup/plugin-alias';
 import json from '@rollup/plugin-json';
 import typescriptPlugin from 'rollup-plugin-typescript2';
 import replace from 'rollup-plugin-replace';
 import typescript from 'typescript';
 import pkg from './package.json';
 import tsconfig from './tsconfig.json';
 import { generateBuildTargetReplaceConfig } from '../../scripts/build/rollup_replace_build_target';
-import { getEnvironmentReplacements } from '../../scripts/build/rollup_get_environment_replacements';
 import { emitModulePackageFile } from '../../scripts/build/rollup_emit_module_package_file';
+import { generateAliasConfig } from '../../scripts/build/rollup_generate_alias_config';
 
 const deps = Object.keys(
   Object.assign({}, pkg.peerDependencies, pkg.dependencies)
@@ -56,9 +57,9 @@ const browserBuilds = [
       sourcemap: true
     },
     plugins: [
+      alias(generateAliasConfig('browser')),
       ...buildPlugins,
       replace({
-        ...getEnvironmentReplacements('browser'),
         ...generateBuildTargetReplaceConfig('esm', 2020),
         '__PACKAGE_VERSION__': pkg.version
       }),
@@ -75,9 +76,9 @@ const browserBuilds = [
       sourcemap: true
     },
     plugins: [
+      alias(generateAliasConfig('browser')),
       ...buildPlugins,
       replace({
-        ...getEnvironmentReplacements('browser'),
         ...generateBuildTargetReplaceConfig('cjs', 2020),
         '__PACKAGE_VERSION__': pkg.version
       })
@@ -96,9 +97,9 @@ const nodeBuilds = [
       sourcemap: true
     },
     plugins: [
+      alias(generateAliasConfig('node')),
       ...buildPlugins,
       replace({
-        ...getEnvironmentReplacements('node'),
         ...generateBuildTargetReplaceConfig('esm', 2020)
       })
     ],
@@ -113,9 +114,9 @@ const nodeBuilds = [
       sourcemap: true
     },
     plugins: [
+      alias(generateAliasConfig('node')),
       ...buildPlugins,
       replace({
-        ...getEnvironmentReplacements('node'),
         ...generateBuildTargetReplaceConfig('cjs', 2020)
       })
     ],

packages/ai/src/ws/browser-websocket-handler.test.ts → packages/ai/src/platform/browser/websocket.test.ts

@@ -19,9 +19,9 @@ import { expect, use } from 'chai';
 import sinon, { SinonFakeTimers, SinonStub } from 'sinon';
 import sinonChai from 'sinon-chai';
 import chaiAsPromised from 'chai-as-promised';
-import { BrowserWebSocketHandler } from './browser-websocket-handler';
-import { AIError } from '../errors';
 import { isBrowser } from '@firebase/util';
+import { BrowserWebSocketHandler } from './websocket';
+import { AIError } from '../../errors';
 
 use(sinonChai);
 use(chaiAsPromised);

packages/ai/src/ws/browser-websocket-handler.ts → packages/ai/src/platform/browser/websocket.ts

@@ -15,31 +15,35 @@
  * limitations under the License.
  */
 
-import { AIError } from '../errors';
-import { AIErrorCode } from '../types';
-import { WebSocketHandler } from './websocket-handler';
+import { AIError } from '../../errors';
+import { AIErrorCode } from '../../types';
+import { WebSocketHandler } from '../websocket';
+
+export function createWebSocketHandler(): WebSocketHandler {
+  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 ' +
+        'modern browser windows, Web Workers with WebSocket support, and Node >= 22.'
+    );
+  }
+
+  return new BrowserWebSocketHandler();
+}
 
 /**
  * A WebSocketHandler implementation for the browser environment.
  * It uses the native `WebSocket`.
+ *
  * @internal
  */
 export class BrowserWebSocketHandler implements WebSocketHandler {
   private ws?: WebSocket;
 
   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',
@@ -63,7 +67,6 @@ export class BrowserWebSocketHandler implements WebSocketHandler {
   }
 
   async *listen(): AsyncGenerator<unknown> {
-    console.log('listener started');
     if (!this.ws) {
       throw new AIError(
         AIErrorCode.REQUEST_ERROR,
@@ -72,23 +75,57 @@ export class BrowserWebSocketHandler implements WebSocketHandler {
     }
 
     const messageQueue: unknown[] = [];
+    const errorQueue: Error[] = [];
     let resolvePromise: (() => void) | null = null;
     let isClosed = false;
 
     const messageListener = async (event: MessageEvent): Promise<void> => {
+      let data: string;
       if (event.data instanceof Blob) {
-        const obj = JSON.parse(await event.data.text()) as unknown;
-        messageQueue.push(obj);
+        data = await event.data.text();
+      } else if (typeof event.data === 'string') {
+        data = event.data;
+      } else {
+        errorQueue.push(
+          new AIError(
+            AIErrorCode.PARSE_FAILED,
+            `Failed to parse WebSocket response. Expected data to be a Blob or string, but was ${typeof event.data}.`
+          )
+        );
         if (resolvePromise) {
           resolvePromise();
           resolvePromise = null;
         }
-      } else {
-        throw new AIError(
-          AIErrorCode.PARSE_FAILED,
-          'Failed to parse WebSocket response to JSON, response was not a Blob'
+        return;
+      }
+
+      try {
+        const obj = JSON.parse(data) as unknown;
+        messageQueue.push(obj);
+      } catch (e) {
+        const err = e as Error;
+        errorQueue.push(
+          new AIError(
+            AIErrorCode.PARSE_FAILED,
+            `Error parsing WebSocket message to JSON: ${err.message}`
+          )
         );
       }
+
+      if (resolvePromise) {
+        resolvePromise();
+        resolvePromise = null;
+      }
+    };
+
+    const errorListener = (): void => {
+      errorQueue.push(
+        new AIError(AIErrorCode.FETCH_ERROR, 'WebSocket connection error.')
+      );
+      if (resolvePromise) {
+        resolvePromise();
+        resolvePromise = null;
+      }
     };
 
     const closeListener = (): void => {
@@ -100,12 +137,18 @@ export class BrowserWebSocketHandler implements WebSocketHandler {
       // Clean up listeners to prevent memory leaks
       this.ws?.removeEventListener('message', messageListener);
       this.ws?.removeEventListener('close', closeListener);
+      this.ws?.removeEventListener('error', errorListener);
     };
 
     this.ws.addEventListener('message', messageListener);
     this.ws.addEventListener('close', closeListener);
+    this.ws.addEventListener('error', errorListener);
 
     while (!isClosed) {
+      if (errorQueue.length > 0) {
+        const error = errorQueue.shift()!;
+        throw error;
+      }
       if (messageQueue.length > 0) {
         yield messageQueue.shift()!;
       } else {
@@ -114,19 +157,32 @@ export class BrowserWebSocketHandler implements WebSocketHandler {
         });
       }
     }
+
+    // If the loop terminated because isClosed is true, check for any final errors
+    if (errorQueue.length > 0) {
+      const error = errorQueue.shift()!;
+      throw error;
+    }
   }
 
   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/ws/node-websocket-handler.test.ts → packages/ai/src/platform/node/websocket.test.ts

@@ -19,19 +19,18 @@ import { expect, use } from 'chai';
 import sinonChai from 'sinon-chai';
 import chaiAsPromised from 'chai-as-promised';
 import { isNode } from '@firebase/util';
-import { NodeWebSocketHandler } from './node-websocket-handler';
-import { WebSocketHandler } from './websocket-handler';
-import { MockWebSocketServer } from '../../test-utils/mock-websocket-server';
-import { AIError } from '../errors';
 import { TextEncoder } from 'util';
+import { MockWebSocketServer } from '../../../test-utils/mock-websocket-server';
+import { WebSocketHandler } from '../websocket';
+import { NodeWebSocketHandler } from './websocket';
 
 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;
     });
   });