Add WebSocket handler for Browser and Node

config/tsconfig.base.json

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

package.json

@@ -86,6 +86,7 @@
     "@types/sinon-chai": "3.2.12",
     "@types/tmp": "0.2.6",
     "@types/trusted-types": "2.0.7",
+    "@types/ws": "8.18.1",
     "@types/yargs": "17.0.33",
     "@typescript-eslint/eslint-plugin": "7.18.0",
     "@typescript-eslint/eslint-plugin-tslint": "7.0.2",

packages/ai/package.json

@@ -39,6 +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:integration": "karma start --integration",
     "api-report": "api-extractor run --local --verbose",
     "typings:public": "node ../../scripts/build/use_typings.js ./dist/ai-public.d.ts",
@@ -62,7 +63,8 @@
     "rollup": "2.79.2",
     "rollup-plugin-replace": "2.2.0",
     "rollup-plugin-typescript2": "0.36.0",
-    "typescript": "5.5.4"
+    "typescript": "5.5.4",
+    "ws": "8.18.3"
   },
   "repository": {
     "directory": "packages/ai",

packages/ai/rollup.config.js

@@ -22,6 +22,7 @@ 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';
 
 const deps = Object.keys(
@@ -57,12 +58,14 @@ const browserBuilds = [
     plugins: [
       ...buildPlugins,
       replace({
+        ...getEnvironmentReplacements('browser'),
         ...generateBuildTargetReplaceConfig('esm', 2020),
-        __PACKAGE_VERSION__: pkg.version
+        '__PACKAGE_VERSION__': pkg.version
       }),
       emitModulePackageFile()
     ],
-    external: id => deps.some(dep => id === dep || id.startsWith(`${dep}/`))
+    external: id =>
+      id === 'ws' || deps.some(dep => id === dep || id.startsWith(`${dep}/`))
   },
   {
     input: 'src/index.ts',
@@ -74,11 +77,13 @@ const browserBuilds = [
     plugins: [
       ...buildPlugins,
       replace({
+        ...getEnvironmentReplacements('browser'),
         ...generateBuildTargetReplaceConfig('cjs', 2020),
-        __PACKAGE_VERSION__: pkg.version
+        '__PACKAGE_VERSION__': pkg.version
       })
     ],
-    external: id => deps.some(dep => id === dep || id.startsWith(`${dep}/`))
+    external: id =>
+      id === 'ws' || deps.some(dep => id === dep || id.startsWith(`${dep}/`))
   }
 ];
 
@@ -93,10 +98,12 @@ const nodeBuilds = [
     plugins: [
       ...buildPlugins,
       replace({
+        ...getEnvironmentReplacements('node'),
         ...generateBuildTargetReplaceConfig('esm', 2020)
       })
     ],
-    external: id => deps.some(dep => id === dep || id.startsWith(`${dep}/`))
+    external: id =>
+      id === 'ws' || deps.some(dep => id === dep || id.startsWith(`${dep}/`))
   },
   {
     input: 'src/index.node.ts',
@@ -108,10 +115,12 @@ const nodeBuilds = [
     plugins: [
       ...buildPlugins,
       replace({
+        ...getEnvironmentReplacements('node'),
         ...generateBuildTargetReplaceConfig('cjs', 2020)
       })
     ],
-    external: id => deps.some(dep => id === dep || id.startsWith(`${dep}/`))
+    external: id =>
+      id === 'ws' || deps.some(dep => id === dep || id.startsWith(`${dep}/`))
   }
 ];
 

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

@@ -0,0 +1,273 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+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';
+
+use(sinonChai);
+use(chaiAsPromised);
+
+class MockBrowserWebSocket {
+  static CONNECTING = 0;
+  static OPEN = 1;
+  static CLOSING = 2;
+  static CLOSED = 3;
+
+  readyState: number = MockBrowserWebSocket.CONNECTING;
+  sentMessages: Array<string | ArrayBuffer> = [];
+  url: string;
+  private listeners: Map<string, Set<EventListener>> = new Map();
+
+  constructor(url: string) {
+    this.url = url;
+  }
+
+  send(data: string | ArrayBuffer): void {
+    if (this.readyState !== MockBrowserWebSocket.OPEN) {
+      throw new Error('WebSocket is not in OPEN state');
+    }
+    this.sentMessages.push(data);
+  }
+
+  close(): void {
+    if (
+      this.readyState === MockBrowserWebSocket.CLOSED ||
+      this.readyState === MockBrowserWebSocket.CLOSING
+    ) {
+      return;
+    }
+    this.readyState = MockBrowserWebSocket.CLOSING;
+    setTimeout(() => {
+      this.readyState = MockBrowserWebSocket.CLOSED;
+      this.dispatchEvent(new Event('close'));
+    }, 10);
+  }
+
+  addEventListener(type: string, listener: EventListener): void {
+    if (!this.listeners.has(type)) {
+      this.listeners.set(type, new Set());
+    }
+    this.listeners.get(type)!.add(listener);
+  }
+
+  removeEventListener(type: string, listener: EventListener): void {
+    this.listeners.get(type)?.delete(listener);
+  }
+
+  dispatchEvent(event: Event): void {
+    this.listeners.get(event.type)?.forEach(listener => listener(event));
+  }
+
+  triggerOpen(): void {
+    this.readyState = MockBrowserWebSocket.OPEN;
+    this.dispatchEvent(new Event('open'));
+  }
+
+  triggerMessage(data: any): void {
+    this.dispatchEvent(new MessageEvent('message', { data }));
+  }
+
+  triggerError(): void {
+    this.dispatchEvent(new Event('error'));
+  }
+}
+
+describe('BrowserWebSocketHandler', () => {
+  let handler: BrowserWebSocketHandler;
+  let mockWebSocket: MockBrowserWebSocket;
+  let clock: SinonFakeTimers;
+  let webSocketStub: SinonStub;
+
+  // Only run these tests in a browser environment
+  if (!isBrowser()) {
+    return;
+  }
+
+  beforeEach(() => {
+    webSocketStub = sinon.stub(window, 'WebSocket').callsFake((url: string) => {
+      mockWebSocket = new MockBrowserWebSocket(url);
+      return mockWebSocket as any;
+    });
+    clock = sinon.useFakeTimers();
+    handler = new BrowserWebSocketHandler();
+  });
+
+  afterEach(() => {
+    sinon.restore();
+    clock.restore();
+  });
+
+  describe('connect()', () => {
+    it('should resolve on open event', async () => {
+      const connectPromise = handler.connect('ws://test-url');
+      expect(webSocketStub).to.have.been.calledWith('ws://test-url');
+
+      await clock.tickAsync(1);
+      mockWebSocket.triggerOpen();
+
+      await expect(connectPromise).to.be.fulfilled;
+    });
+
+    it('should reject on error event', async () => {
+      const connectPromise = handler.connect('ws://test-url');
+      await clock.tickAsync(1);
+      mockWebSocket.triggerError();
+
+      await expect(connectPromise).to.be.rejectedWith(
+        AIError,
+        /Failed to establish WebSocket connection/
+      );
+    });
+  });
+
+  describe('listen()', () => {
+    beforeEach(async () => {
+      const connectPromise = handler.connect('ws://test');
+      mockWebSocket.triggerOpen();
+      await connectPromise;
+    });
+
+    it('should yield multiple messages as they arrive', async () => {
+      const generator = handler.listen();
+
+      const received: unknown[] = [];
+      const listenPromise = (async () => {
+        for await (const msg of generator) {
+          received.push(msg);
+        }
+      })();
+
+      // Use tickAsync to allow the consumer to start listening
+      await clock.tickAsync(1);
+      mockWebSocket.triggerMessage(new Blob([JSON.stringify({ foo: 1 })]));
+
+      await clock.tickAsync(10);
+      mockWebSocket.triggerMessage(new Blob([JSON.stringify({ foo: 2 })]));
+
+      await clock.tickAsync(5);
+      mockWebSocket.close();
+      await clock.runAllAsync(); // Let timers finish
+
+      await listenPromise; // Wait for the consumer to finish
+
+      expect(received).to.deep.equal([
+        {
+          foo: 1
+        },
+        {
+          foo: 2
+        }
+      ]);
+    });
+
+    it('should buffer messages that arrive before the consumer calls .next()', async () => {
+      const generator = handler.listen();
+
+      // Create a promise that will consume the generator in a separate async context
+      const received: unknown[] = [];
+      const consumptionPromise = (async () => {
+        for await (const message of generator) {
+          received.push(message);
+        }
+      })();
+
+      await clock.tickAsync(1);
+
+      mockWebSocket.triggerMessage(new Blob([JSON.stringify({ foo: 1 })]));
+      mockWebSocket.triggerMessage(new Blob([JSON.stringify({ foo: 2 })]));
+
+      await clock.tickAsync(1);
+      mockWebSocket.close();
+      await clock.runAllAsync();
+
+      await consumptionPromise;
+
+      expect(received).to.deep.equal([
+        {
+          foo: 1
+        },
+        {
+          foo: 2
+        }
+      ]);
+    });
+  });
+
+  describe('close()', () => {
+    it('should be idempotent and not throw if called multiple times', async () => {
+      const connectPromise = handler.connect('ws://test');
+      mockWebSocket.triggerOpen();
+      await connectPromise;
+
+      const closePromise1 = handler.close();
+      await clock.runAllAsync();
+      await closePromise1;
+
+      await expect(handler.close()).to.be.fulfilled;
+    });
+
+    it('should wait for the onclose event before resolving', async () => {
+      const connectPromise = handler.connect('ws://test');
+      mockWebSocket.triggerOpen();
+      await connectPromise;
+
+      let closed = false;
+      const closePromise = handler.close().then(() => {
+        closed = true;
+      });
+
+      // The promise should not have resolved yet
+      await clock.tickAsync(5);
+      expect(closed).to.be.false;
+
+      // Now, let the mock's setTimeout for closing run, which triggers onclose
+      await clock.tickAsync(10);
+
+      await expect(closePromise).to.be.fulfilled;
+      expect(closed).to.be.true;
+    });
+  });
+
+  describe('Interaction between listen() and close()', () => {
+    it('should allow close() to take precedence and resolve correctly, while also terminating the listener', async () => {
+      const connectPromise = handler.connect('ws://test');
+      mockWebSocket.triggerOpen();
+      await connectPromise;
+
+      const generator = handler.listen();
+      const listenPromise = (async () => {
+        // eslint-disable-next-line @typescript-eslint/no-unused-vars
+        for await (const _ of generator) {
+        }
+      })();
+
+      const closePromise = handler.close();
+
+      await clock.runAllAsync();
+
+      await expect(closePromise).to.be.fulfilled;
+      await expect(listenPromise).to.be.fulfilled;
+
+      expect(mockWebSocket.readyState).to.equal(MockBrowserWebSocket.CLOSED);
+    });
+  });
+});