Use platform/ convention

Author: dlarocque

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",

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 renamed to 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);
@@ -270,4 +270,4 @@ describe('BrowserWebSocketHandler', () => {
       expect(mockWebSocket.readyState).to.equal(MockBrowserWebSocket.CLOSED);
     });
   });
-});
+});

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

@@ -15,13 +15,27 @@
  * 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') {
+    return new BrowserWebSocketHandler();
+  } else {
+    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.'
+    );
+  }
+}
 
 /**
  * A WebSocketHandler implementation for the browser environment.
  * It uses the native `WebSocket`.
+ *
  * @internal
  */
 export class BrowserWebSocketHandler implements WebSocketHandler {
@@ -63,7 +77,6 @@ export class BrowserWebSocketHandler implements WebSocketHandler {
   }
 
   async *listen(): AsyncGenerator<unknown> {
-    console.log('listener started');
     if (!this.ws) {
       throw new AIError(
         AIErrorCode.REQUEST_ERROR,
@@ -77,16 +90,21 @@ export class BrowserWebSocketHandler implements WebSocketHandler {
 
     const messageListener = async (event: MessageEvent): Promise<void> => {
       if (event.data instanceof Blob) {
-        const obj = JSON.parse(await event.data.text()) as unknown;
-        messageQueue.push(obj);
-        if (resolvePromise) {
-          resolvePromise();
-          resolvePromise = null;
+        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, response was not a Blob'
+          `Failed to parse WebSocket response to JSON. ` +
+            `Expected data to be a Blob, but was ${typeof event.data}.`
         );
       }
     };

packages/ai/src/ws/node-websocket-handler.test.ts renamed to packages/ai/src/platform/node/websocket.test.ts

@@ -19,11 +19,11 @@ 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';
+import { AIError } from '../../errors';
 
 use(sinonChai);
 use(chaiAsPromised);
@@ -144,4 +144,4 @@ describe('NodeWebSocketHandler (Integration Tests)', () => {
       await expect(consumerPromise).to.be.fulfilled;
     });
   });
-});
+});

packages/ai/src/ws/node-websocket-handler.ts renamed to packages/ai/src/platform/node/websocket.ts

@@ -15,28 +15,53 @@
  * limitations under the License.
  */
 
-import { AIError } from '../errors';
-import { AIErrorCode } from '../public-types';
-import { WebSocketHandler } from './websocket-handler';
+// import { WebSocket, MessageEvent } from 'ws'; // External dependency on native Node module
+import { AIError } from '../../errors';
+import { AIErrorCode } from '../../types';
+import { WebSocketHandler } from '../websocket';
+
+export function createWebSocketHandler(): WebSocketHandler {
+  if (typeof process === 'object' && process.versions?.node) {
+    const [major] = process.versions.node.split('.').map(Number);
+    if (major < 22) {
+      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` +
+          `for native WebSocket support.`
+      );
+    }
+    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.'
+    );
+  }
+}
 
 /**
  * A WebSocketHandler implementation for Node >= 22.
- * It uses the native, built-in 'ws' module, which must be imported.
  *
  * @internal
  */
 export class NodeWebSocketHandler implements WebSocketHandler {
   private ws?: import('ws').WebSocket;
 
   async connect(url: string): Promise<void> {
-    // This dynamic import is why we need a separate class.
-    // It is only ever executed in a Node environment, preventing browser
-    // bundlers from attempting to resolve this Node-specific module.
-    // eslint-disable-next-line import/no-extraneous-dependencies
-    const { WebSocket } = await import('ws');
-
-    return new Promise((resolve, reject) => {
-      this.ws = new WebSocket(url);
+    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!.addEventListener('open', () => resolve(), { once: true });
       this.ws!.addEventListener(
         'error',
@@ -84,8 +109,11 @@ export class NodeWebSocketHandler implements WebSocketHandler {
         const decoder = new TextDecoder();
         textData = decoder.decode(event.data);
       } else {
-        console.warn('Received unexpected WebSocket message type:', event.data);
-        return;
+        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}.`
+        );
       }
 
       try {
@@ -106,6 +134,7 @@ export class NodeWebSocketHandler implements WebSocketHandler {
         resolvePromise();
         resolvePromise = null;
       }
+      // Clean up listeners to prevent memory leaks
       this.ws?.removeEventListener('message', messageListener);
       this.ws?.removeEventListener('close', closeListener);
     };
@@ -126,7 +155,11 @@ export class NodeWebSocketHandler implements WebSocketHandler {
 
   close(code?: number, reason?: string): Promise<void> {
     return new Promise(resolve => {
-      if (!this.ws || this.ws.readyState === this.ws.CLOSED) {
+      if (
+        !this.ws ||
+        this.ws.readyState === this.ws.CLOSED ||
+        this.ws.readyState === this.ws.CLOSING
+      ) {
         return resolve();
       }
       this.ws.addEventListener('close', () => resolve(), { once: true });

packages/ai/src/ws/websocket-handler.ts renamed to packages/ai/src/platform/websocket.ts

@@ -17,9 +17,9 @@
 
 import { isBrowser, isNode } from '@firebase/util';
 import { AIError } from '../errors';
-import { AIErrorCode } from '../public-types';
-import { NodeWebSocketHandler } from './node-websocket-handler';
-import { BrowserWebSocketHandler } from './browser-websocket-handler';
+import { AIErrorCode } from '../types';
+import { NodeWebSocketHandler } from './node/websocket';
+import { BrowserWebSocketHandler } from './browser/websocket';
 
 /**
  * A standardized interface for interacting with a WebSocket connection.
@@ -71,10 +71,7 @@ export interface WebSocketHandler {
  *
  * 1. Module Loading: The primary difference is how the `WebSocket` class is
  *    accessed. In browsers, it's a global (`window.WebSocket`). In Node, it
- *    must be imported from the built-in `'ws'` module. Isolating the Node
- *    `import('ws')` call in its own class prevents browser-bundling tools
- *    (like Webpack, Vite) from trying to resolve a Node-specific module, which
- *    would either fail the build or include unnecessary polyfills.
+ *    must be imported from the built-in `'ws'` module.
  *
  * 2. Type Safety: TypeScript's type definitions for the browser's WebSocket
  *    (from `lib.dom.d.ts`) and Node's WebSocket (from `@types/node`) are
@@ -83,11 +80,7 @@ export interface WebSocketHandler {
  * @internal
  */
 export function createWebSocketHandler(): WebSocketHandler {
-  // `isNode()` is replaced with a static boolean during build time so this block will be
-  // tree-shaken in browser builds.
   if (isNode()) {
-    // At this point we're certain we're in a Node bundle, but we still need to have checks
-    // to be certain we're in a Node environment, and not something like Deno, Bun, or Edge workers.
     if (typeof process === 'object' && process.versions?.node) {
       const [major] = process.versions.node.split('.').map(Number);
       if (major < 22) {
@@ -102,12 +95,7 @@ export function createWebSocketHandler(): WebSocketHandler {
     }
   }
 
-  // `isBrowser()` is replaced with a static boolean during build time so this block will be
-  // tree-shaken in Node builds.
   if (isBrowser()) {
-    // At this point we're certain we're in a browser build, but we still need to check for the
-    // existence of the `WebSocket` API. This check would fail in environments that use a browser
-    // bundle, but don't support WebSockets (Web workers and SSR).
     if (typeof WebSocket !== 'undefined') {
       return new BrowserWebSocketHandler();
     } else {

scripts/build/rollup_generate_alias_config.js

@@ -0,0 +1,27 @@
+/**
+ * @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.
+ */
+
+export function generateAliasConfig(platform) {
+  return {
+    entries: [
+      {
+        find: /^(.*)\/platform\/([^.\/]*)(\.ts)?$/,
+        replacement: `$1\/platform/${platform}/$2.ts`
+      }
+    ]
+  };
+}