From 9cc2187ed9fee9b4e8ce8ed067b46b461bfb79ae Mon Sep 17 00:00:00 2001 From: Daniel La Rocque Date: Thu, 31 Jul 2025 11:30:59 -0400 Subject: [PATCH 01/10] Add WebSocket handler for Browser and Node --- config/tsconfig.base.json | 2 +- package.json | 1 + packages/ai/package.json | 4 +- packages/ai/rollup.config.js | 21 +- .../src/ws/browser-websocket-handler.test.ts | 273 ++++++++++++++++++ .../ai/src/ws/browser-websocket-handler.ts | 132 +++++++++ .../ai/src/ws/node-websocket-handler.test.ts | 147 ++++++++++ packages/ai/src/ws/node-websocket-handler.ts | 136 +++++++++ packages/ai/src/ws/websocket-handler.ts | 107 +++++++ packages/ai/test-utils/convert-mocks.ts | 4 +- packages/ai/test-utils/mock-response.ts | 1 + .../ai/test-utils/mock-websocket-server.ts | 76 +++++ packages/ai/test-utils/types.ts | 18 ++ packages/util/src/websocket-handler.ts | 0 .../rollup_get_environment_replacements.js | 49 ++++ yarn.lock | 43 +-- 16 files changed, 976 insertions(+), 38 deletions(-) create mode 100644 packages/ai/src/ws/browser-websocket-handler.test.ts create mode 100644 packages/ai/src/ws/browser-websocket-handler.ts create mode 100644 packages/ai/src/ws/node-websocket-handler.test.ts create mode 100644 packages/ai/src/ws/node-websocket-handler.ts create mode 100644 packages/ai/src/ws/websocket-handler.ts create mode 100644 packages/ai/test-utils/mock-websocket-server.ts create mode 100644 packages/ai/test-utils/types.ts create mode 100644 packages/util/src/websocket-handler.ts create mode 100644 scripts/build/rollup_get_environment_replacements.js diff --git a/config/tsconfig.base.json b/config/tsconfig.base.json index ce58a6d700b..33a58c5b02d 100644 --- a/config/tsconfig.base.json +++ b/config/tsconfig.base.json @@ -9,7 +9,7 @@ "es2020", "esnext.WeakRef", ], - "module": "ES2015", + "module": "ES2020", "moduleResolution": "node", "resolveJsonModule": true, "esModuleInterop": true, diff --git a/package.json b/package.json index 39455ef1161..8f42b1dfce3 100644 --- a/package.json +++ b/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", diff --git a/packages/ai/package.json b/packages/ai/package.json index 97186afb1e1..785525cb07e 100644 --- a/packages/ai/package.json +++ b/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", diff --git a/packages/ai/rollup.config.js b/packages/ai/rollup.config.js index 7ebbff4f2f5..13389b427a1 100644 --- a/packages/ai/rollup.config.js +++ b/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}/`)) } ]; diff --git a/packages/ai/src/ws/browser-websocket-handler.test.ts b/packages/ai/src/ws/browser-websocket-handler.test.ts new file mode 100644 index 00000000000..4e1500f34d4 --- /dev/null +++ b/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 = []; + url: string; + private listeners: Map> = 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); + }); + }); +}); diff --git a/packages/ai/src/ws/browser-websocket-handler.ts b/packages/ai/src/ws/browser-websocket-handler.ts new file mode 100644 index 00000000000..1a55cf436df --- /dev/null +++ b/packages/ai/src/ws/browser-websocket-handler.ts @@ -0,0 +1,132 @@ +/** + * @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 { AIError } from '../errors'; +import { AIErrorCode } from '../types'; +import { WebSocketHandler } from './websocket-handler'; + +/** + * A WebSocketHandler implementation for the browser environment. + * It uses the native `WebSocket` class available on the `window` object. + * @internal + */ +export class BrowserWebSocketHandler implements WebSocketHandler { + private ws?: WebSocket; + + connect(url: string): Promise { + 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.addEventListener('open', () => resolve(), { once: true }); + this.ws.addEventListener( + 'error', + () => + reject( + new AIError( + AIErrorCode.FETCH_ERROR, + 'Failed to establish WebSocket connection' + ) + ), + { once: true } + ); + }); + } + + send(data: string | ArrayBuffer): void { + if (!this.ws || this.ws.readyState !== WebSocket.OPEN) { + throw new AIError(AIErrorCode.REQUEST_ERROR, 'WebSocket is not open.'); + } + this.ws.send(data); + } + + async *listen(): AsyncGenerator { + console.log('listener started'); + if (!this.ws) { + throw new AIError( + AIErrorCode.REQUEST_ERROR, + 'WebSocket is not connected.' + ); + } + + const messageQueue: unknown[] = []; + let resolvePromise: (() => void) | null = null; + let isClosed = false; + + const messageListener = async (event: MessageEvent): Promise => { + if (event.data instanceof Blob) { + const obj = JSON.parse(await event.data.text()) as unknown; + messageQueue.push(obj); + if (resolvePromise) { + resolvePromise(); + resolvePromise = null; + } + } else { + throw new AIError( + AIErrorCode.PARSE_FAILED, + 'Failed to parse WebSocket response to JSON, response was not a Blob' + ); + } + }; + + const closeListener = (): void => { + isClosed = true; + if (resolvePromise) { + resolvePromise(); + resolvePromise = null; + } + // Clean up listeners to prevent memory leaks + this.ws?.removeEventListener('message', messageListener); + this.ws?.removeEventListener('close', closeListener); + }; + + this.ws.addEventListener('message', messageListener); + this.ws.addEventListener('close', closeListener); + + while (!isClosed) { + if (messageQueue.length > 0) { + yield messageQueue.shift()!; + } else { + await new Promise(resolve => { + resolvePromise = resolve; + }); + } + } + } + + close(code?: number, reason?: string): Promise { + return new Promise(resolve => { + if ( + !this.ws || + this.ws.readyState === WebSocket.CLOSED || + this.ws.readyState === WebSocket.CLOSING + ) { + return resolve(); + } + this.ws.addEventListener('close', () => resolve(), { once: true }); + this.ws.close(code, reason); + }); + } +} diff --git a/packages/ai/src/ws/node-websocket-handler.test.ts b/packages/ai/src/ws/node-websocket-handler.test.ts new file mode 100644 index 00000000000..22f117332fc --- /dev/null +++ b/packages/ai/src/ws/node-websocket-handler.test.ts @@ -0,0 +1,147 @@ +/** + * @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 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'; + +use(sinonChai); +use(chaiAsPromised); + +const TEST_PORT = 9003; +const TEST_URL = `ws://localhost:${TEST_PORT}`; + +describe('NodeWebSocketHandler (Integration Tests)', () => { + let server: MockWebSocketServer; + let handler: WebSocketHandler; + + // Only run these tests in a Node environment + if (!isNode()) { + return; + } + + before(async () => { + server = new MockWebSocketServer(TEST_PORT); + }); + + after(async () => { + await server.close(); + }); + + beforeEach(() => { + handler = new NodeWebSocketHandler(); + server.reset(); + }); + + afterEach(async () => { + await handler.close().catch(() => {}); + }); + + describe('connect()', () => { + it('should successfully connect to a running server', async () => { + await handler.connect(TEST_URL); + // Allow a brief moment for the server to register the connection + await new Promise(r => setTimeout(r, 50)); + expect(server.connectionCount).to.equal(1); + 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/ + ); + }); + }); + + describe('listen()', () => { + beforeEach(async () => { + await handler.connect(TEST_URL); + // Wait for server to see the connection + await new Promise(r => setTimeout(r, 50)); + }); + + it('should yield parsed JSON objects from string data sent by the server', async () => { + const generator = handler.listen(); + const messageObj = { id: 1, text: 'test' }; + + const received: unknown[] = []; + const consumerPromise = (async () => { + for await (const msg of generator) { + received.push(msg); + } + })(); + + // Wait for the listener to be attached + await new Promise(r => setTimeout(r, 50)); + server.broadcast(JSON.stringify(messageObj)); + await new Promise(r => setTimeout(r, 50)); + await handler.close(); // Close client to terminate the loop + + await consumerPromise; + expect(received).to.deep.equal([messageObj]); + }); + + it('should correctly decode UTF-8 binary data sent by the server', async () => { + const generator = handler.listen(); + const messageObj = { text: '你好, 世界 🌍' }; + const encoder = new TextEncoder(); + const bufferData = encoder.encode(JSON.stringify(messageObj)); + + const received: unknown[] = []; + const consumerPromise = (async () => { + for await (const msg of generator) { + received.push(msg); + } + })(); + + await new Promise(r => setTimeout(r, 50)); + // The server's `send` method can handle Buffers/Uint8Arrays + server.clients.forEach(client => client.send(bufferData)); + await new Promise(r => setTimeout(r, 50)); + await handler.close(); + + await consumerPromise; + expect(received).to.deep.equal([messageObj]); + }); + + it('should terminate the generator when the server closes the connection', async () => { + const generator = handler.listen(); + const consumerPromise = (async () => { + // This loop should finish without error when the server closes + // eslint-disable-next-line @typescript-eslint/no-unused-vars + for await (const _ of generator) { + } + })(); + + await new Promise(r => setTimeout(r, 50)); + + await server.close(); + server = new MockWebSocketServer(TEST_PORT); + + // The consumer promise should resolve without timing out + await expect(consumerPromise).to.be.fulfilled; + }); + }); +}); diff --git a/packages/ai/src/ws/node-websocket-handler.ts b/packages/ai/src/ws/node-websocket-handler.ts new file mode 100644 index 00000000000..be21b73cd99 --- /dev/null +++ b/packages/ai/src/ws/node-websocket-handler.ts @@ -0,0 +1,136 @@ +/** + * @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 { AIError } from '../errors'; +import { AIErrorCode } from '../public-types'; +import { WebSocketHandler } from './websocket-handler'; + +/** + * 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 { + // 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); + this.ws!.addEventListener('open', () => resolve(), { once: true }); + this.ws!.addEventListener( + 'error', + () => + reject( + new AIError( + AIErrorCode.FETCH_ERROR, + 'Failed to establish WebSocket connection' + ) + ), + { once: true } + ); + }); + } + + send(data: string | ArrayBuffer): void { + if (!this.ws || this.ws.readyState !== this.ws.OPEN) { + throw new AIError(AIErrorCode.REQUEST_ERROR, 'WebSocket is not open.'); + } + this.ws.send(data); + } + + async *listen(): AsyncGenerator { + if (!this.ws) { + throw new AIError( + AIErrorCode.REQUEST_ERROR, + 'WebSocket is not connected.' + ); + } + + const messageQueue: unknown[] = []; + 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); + } else { + console.warn('Received unexpected WebSocket message type:', event.data); + return; + } + + 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 => { + isClosed = true; + if (resolvePromise) { + resolvePromise(); + resolvePromise = null; + } + this.ws?.removeEventListener('message', messageListener); + this.ws?.removeEventListener('close', closeListener); + }; + + this.ws.addEventListener('message', messageListener); + this.ws.addEventListener('close', closeListener); + + while (!isClosed) { + if (messageQueue.length > 0) { + yield messageQueue.shift()!; + } else { + await new Promise(resolve => { + resolvePromise = resolve; + }); + } + } + } + + close(code?: number, reason?: string): Promise { + return new Promise(resolve => { + if (!this.ws || this.ws.readyState === this.ws.CLOSED) { + return resolve(); + } + this.ws.addEventListener('close', () => resolve(), { once: true }); + this.ws.close(code, reason); + }); + } +} diff --git a/packages/ai/src/ws/websocket-handler.ts b/packages/ai/src/ws/websocket-handler.ts new file mode 100644 index 00000000000..754b030729e --- /dev/null +++ b/packages/ai/src/ws/websocket-handler.ts @@ -0,0 +1,107 @@ +/** + * @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 { 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'; + +/** + * A standardized interface for interacting with a WebSocket connection. + * This abstraction allows the SDK to use the appropriate WebSocket implementation + * for the current JS environment (Browser vs. Node) without + * changing the core logic of the `LiveSession`. + * @internal + */ +export interface WebSocketHandler { + /** + * Establishes a connection to the given URL. + * + * @param url The WebSocket URL (e.g., wss://...). + * @returns A promise that resolves on successful connection or rejects on failure. + */ + connect(url: string): Promise; + + /** + * Sends data over the WebSocket. + * + * @param data The string or binary data to send. + */ + send(data: string | ArrayBuffer): void; + + /** + * Returns an async generator that yields parsed JSON objects from the server. + * The yielded type is `unknown` because the handler cannot guarantee the shape of the data. + * The consumer is responsible for type validation. + * The generator terminates when the connection is closed. + * + * @returns A generator that allows consumers to pull messages using a `for await...of` loop. + */ + listen(): AsyncGenerator; + + /** + * Closes the WebSocket connection. + * + * @param code - A numeric status code explaining why the connection is closing. + * @param reason - A human-readable string explaining why the connection is closing. + */ + close(code?: number, reason?: string): Promise; +} + +/** + * Factory function to create the appropriate WebSocketHandler for the current environment. + * + * Even though the browser and Node >=22 WebSocket APIs are now very similar, + * we use two separate handler classes. There are two reasons for this: + * + * 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. + * + * 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 + * distinct. Using separate classes ensures type correctness for each environment. + * + * @internal + */ +export function createWebSocketHandler(): WebSocketHandler { + // `isNode()` is replaced with a static boolean during build time to enable tree shaking + if (isNode()) { + const [major] = process.versions.node.split('.').map(Number); + if (major < 22) { + throw new AIError( + AIErrorCode.UNSUPPORTED, + 'The Live feature requires Node version 22 or higher for native WebSocket support.' + ); + } + return new NodeWebSocketHandler(); + } + + // `isBrowser()` is replaced with a static boolean during build time to enable tree shaking + if (isBrowser() && window.WebSocket) { + return new BrowserWebSocketHandler(); + } + + throw new AIError( + AIErrorCode.UNSUPPORTED, + 'A WebSocket API is not available in this environment.' + ); +} diff --git a/packages/ai/test-utils/convert-mocks.ts b/packages/ai/test-utils/convert-mocks.ts index 4bac70d1d10..34233a73ace 100644 --- a/packages/ai/test-utils/convert-mocks.ts +++ b/packages/ai/test-utils/convert-mocks.ts @@ -19,6 +19,8 @@ const { readdirSync, readFileSync, writeFileSync } = require('node:fs'); const { join } = require('node:path'); +type BackendName = import('./types').BackendName; // Import type without triggering ES module detection + const MOCK_RESPONSES_DIR_PATH = join( __dirname, 'vertexai-sdk-test-data', @@ -26,8 +28,6 @@ const MOCK_RESPONSES_DIR_PATH = join( ); const MOCK_LOOKUP_OUTPUT_PATH = join(__dirname, 'mocks-lookup.ts'); -type BackendName = 'vertexAI' | 'googleAI'; - const mockDirs: Record = { vertexAI: join(MOCK_RESPONSES_DIR_PATH, 'vertexai'), googleAI: join(MOCK_RESPONSES_DIR_PATH, 'googleai') diff --git a/packages/ai/test-utils/mock-response.ts b/packages/ai/test-utils/mock-response.ts index 5128ddabe74..4963bcbb193 100644 --- a/packages/ai/test-utils/mock-response.ts +++ b/packages/ai/test-utils/mock-response.ts @@ -15,6 +15,7 @@ * limitations under the License. */ +import { BackendName } from './types'; import { vertexAIMocksLookup, googleAIMocksLookup } from './mocks-lookup'; const mockSetMaps: Record> = { diff --git a/packages/ai/test-utils/mock-websocket-server.ts b/packages/ai/test-utils/mock-websocket-server.ts new file mode 100644 index 00000000000..f2f83f569ff --- /dev/null +++ b/packages/ai/test-utils/mock-websocket-server.ts @@ -0,0 +1,76 @@ +/** + * @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 { WebSocketServer, WebSocket } from 'ws'; + +/** + * A mock WebSocket server for running integration tests against the + * `NodeWebSocketHandler`. It listens on a specified port, accepts connections, + * logs messages, and can broadcast messages to clients. + * + * This should only be used in a Node environment. + * + * @internal + */ +export class MockWebSocketServer { + private wss: WebSocketServer; + clients: Set = new Set(); + receivedMessages: string[] = []; + connectionCount = 0; + + constructor(public port: number) { + this.wss = new WebSocketServer({ port }); + + this.wss.on('connection', ws => { + this.connectionCount++; + this.clients.add(ws); + + ws.on('message', message => { + this.receivedMessages.push(message.toString()); + }); + + ws.on('close', () => { + this.clients.delete(ws); + }); + }); + } + + broadcast(message: string | Buffer): void { + for (const client of this.clients) { + if (client.readyState === WebSocket.OPEN) { + client.send(message); + } + } + } + + close(): Promise { + return new Promise(resolve => { + for (const client of this.clients) { + client.terminate(); + } + this.wss.close(() => { + this.reset(); + resolve(); + }); + }); + } + + reset(): void { + this.receivedMessages = []; + this.connectionCount = 0; + } +} diff --git a/packages/ai/test-utils/types.ts b/packages/ai/test-utils/types.ts new file mode 100644 index 00000000000..00b99eef55a --- /dev/null +++ b/packages/ai/test-utils/types.ts @@ -0,0 +1,18 @@ +/** + * @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 type BackendName = 'vertexAI' | 'googleAI'; diff --git a/packages/util/src/websocket-handler.ts b/packages/util/src/websocket-handler.ts new file mode 100644 index 00000000000..e69de29bb2d diff --git a/scripts/build/rollup_get_environment_replacements.js b/scripts/build/rollup_get_environment_replacements.js new file mode 100644 index 00000000000..5206a03eb03 --- /dev/null +++ b/scripts/build/rollup_get_environment_replacements.js @@ -0,0 +1,49 @@ +/** + * @license + * Copyright 2021 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. + */ + +/** + * Generates the `rollup-plugin-replace` configuration that enables tree-shaking + * of platform-specific code by replacing runtime checks like `isNode()` with boolean literals + * at build time. + * + * For a browser build, `if (isNode()) { ... }` becomes `if (false) { ... }`. + * This allows Rollup to identify and eliminate the Node-specific code branches + * and their imports as dead code. The reverse is true for Node builds. + */ +export function getEnvironmentReplacements(environment) { + const replacements = { + delimiters: ['', ''], // Set to empty strings to replace the entire function call (`isNode()`), not just a variable + preventAssignment: true + }; + + switch (environment) { + case 'browser': + return { + ...replacements, + 'isNode()': 'false', + 'isBrowser()': 'true' + }; + case 'node': + return { + ...replacements, + 'isNode()': 'true', + 'isBrowser()': 'false' + }; + default: + throw new Error(`Unknown build environment: ${environment}`); + } +} diff --git a/yarn.lock b/yarn.lock index fe69e44aead..cfc5a206b1f 100644 --- a/yarn.lock +++ b/yarn.lock @@ -3321,6 +3321,13 @@ tapable "^2.2.0" webpack "^5" +"@types/ws@8.18.1": + version "8.18.1" + resolved "https://registry.npmjs.org/@types/ws/-/ws-8.18.1.tgz#48464e4bf2ddfd17db13d845467f6070ffea4aa9" + integrity sha512-ThVF6DCVhA8kUGy+aazFQ4kXQ7E1Ty7A3ypFOe0IcJV8O/M511G99AW24irKrW56Wt44yG9+ij8FaqoBGkuBXg== + dependencies: + "@types/node" "*" + "@types/yargs-parser@*": version "21.0.3" resolved "https://registry.npmjs.org/@types/yargs-parser/-/yargs-parser-21.0.3.tgz#815e30b786d2e8f0dcd85fd5bcf5e1a04d008f15" @@ -15139,7 +15146,7 @@ string-argv@~0.3.1: resolved "https://registry.npmjs.org/string-argv/-/string-argv-0.3.2.tgz#2b6d0ef24b656274d957d54e0a4bbf6153dc02b6" integrity sha512-aqD2Q0144Z+/RqG52NeHEkZauTAUWJO8c6yTftGJKO3Tja5tUgIfmIl6kExvhtxSDP7fXB6DvzkfMpCd/F3G+Q== -"string-width-cjs@npm:string-width@^4.2.0": +"string-width-cjs@npm:string-width@^4.2.0", "string-width@^1.0.2 || 2 || 3 || 4", string-width@^4.0.0, string-width@^4.1.0, string-width@^4.2.0, string-width@^4.2.2, string-width@^4.2.3: version "4.2.3" resolved "https://registry.npmjs.org/string-width/-/string-width-4.2.3.tgz#269c7117d27b05ad2e536830a8ec895ef9c6d010" integrity sha512-wKyQRQpjJ0sIp62ErSZdGsjMJWsap5oRNihHhu6G7JVO/9jIB6UyevL+tXuOqrng8j/cxKTWyWUwvSTriiZz/g== @@ -15157,15 +15164,6 @@ string-width@^1.0.1, string-width@^1.0.2: is-fullwidth-code-point "^1.0.0" strip-ansi "^3.0.0" -"string-width@^1.0.2 || 2 || 3 || 4", string-width@^4.0.0, string-width@^4.1.0, string-width@^4.2.0, string-width@^4.2.2, string-width@^4.2.3: - version "4.2.3" - resolved "https://registry.npmjs.org/string-width/-/string-width-4.2.3.tgz#269c7117d27b05ad2e536830a8ec895ef9c6d010" - integrity sha512-wKyQRQpjJ0sIp62ErSZdGsjMJWsap5oRNihHhu6G7JVO/9jIB6UyevL+tXuOqrng8j/cxKTWyWUwvSTriiZz/g== - dependencies: - emoji-regex "^8.0.0" - is-fullwidth-code-point "^3.0.0" - strip-ansi "^6.0.1" - string-width@^2.1.1: version "2.1.1" resolved "https://registry.npmjs.org/string-width/-/string-width-2.1.1.tgz#ab93f27a8dc13d28cac815c462143a6d9012ae9e" @@ -15229,7 +15227,7 @@ string_decoder@~1.1.1: dependencies: safe-buffer "~5.1.0" -"strip-ansi-cjs@npm:strip-ansi@^6.0.1": +"strip-ansi-cjs@npm:strip-ansi@^6.0.1", strip-ansi@^6.0.0, strip-ansi@^6.0.1: version "6.0.1" resolved "https://registry.npmjs.org/strip-ansi/-/strip-ansi-6.0.1.tgz#9e26c63d30f53443e9489495b2105d37b67a85d9" integrity sha512-Y38VPSHcqkFrCpFnQ9vuSXmquuv5oXOKpGeT6aGrr3o3Gc9AlVa6JBfUSOCnbxGGZF+/0ooI7KrPuUSztUdU5A== @@ -15250,13 +15248,6 @@ strip-ansi@^4.0.0: dependencies: ansi-regex "^3.0.0" -strip-ansi@^6.0.0, strip-ansi@^6.0.1: - version "6.0.1" - resolved "https://registry.npmjs.org/strip-ansi/-/strip-ansi-6.0.1.tgz#9e26c63d30f53443e9489495b2105d37b67a85d9" - integrity sha512-Y38VPSHcqkFrCpFnQ9vuSXmquuv5oXOKpGeT6aGrr3o3Gc9AlVa6JBfUSOCnbxGGZF+/0ooI7KrPuUSztUdU5A== - dependencies: - ansi-regex "^5.0.1" - strip-ansi@^7.0.1: version "7.1.0" resolved "https://registry.npmjs.org/strip-ansi/-/strip-ansi-7.1.0.tgz#d5b6568ca689d8561370b0707685d22434faff45" @@ -16929,7 +16920,7 @@ workerpool@6.2.0: resolved "https://registry.npmjs.org/workerpool/-/workerpool-6.2.0.tgz#827d93c9ba23ee2019c3ffaff5c27fccea289e8b" integrity sha512-Rsk5qQHJ9eowMH28Jwhe8HEbmdYDX4lwoMWshiCXugjtHqMD9ZbiqSDLxcsfdqsETPzVUtX5s1Z5kStiIM6l4A== -"wrap-ansi-cjs@npm:wrap-ansi@^7.0.0": +"wrap-ansi-cjs@npm:wrap-ansi@^7.0.0", wrap-ansi@^7.0.0: version "7.0.0" resolved "https://registry.npmjs.org/wrap-ansi/-/wrap-ansi-7.0.0.tgz#67e145cff510a6a6984bdf1152911d69d2eb9e43" integrity sha512-YVGIj2kamLSTxw6NsZjoBxfSwsn0ycdesmc4p+Q21c5zPuZ1pl+NfxVdxPtdHvmNVOQ6XSYG4AUtyt/Fi7D16Q== @@ -16963,15 +16954,6 @@ wrap-ansi@^6.0.1, wrap-ansi@^6.2.0: string-width "^4.1.0" strip-ansi "^6.0.0" -wrap-ansi@^7.0.0: - version "7.0.0" - resolved "https://registry.npmjs.org/wrap-ansi/-/wrap-ansi-7.0.0.tgz#67e145cff510a6a6984bdf1152911d69d2eb9e43" - integrity sha512-YVGIj2kamLSTxw6NsZjoBxfSwsn0ycdesmc4p+Q21c5zPuZ1pl+NfxVdxPtdHvmNVOQ6XSYG4AUtyt/Fi7D16Q== - dependencies: - ansi-styles "^4.0.0" - string-width "^4.1.0" - strip-ansi "^6.0.0" - wrap-ansi@^8.1.0: version "8.1.0" resolved "https://registry.npmjs.org/wrap-ansi/-/wrap-ansi-8.1.0.tgz#56dc22368ee570face1b49819975d9b9a5ead214" @@ -17055,6 +17037,11 @@ write-pkg@^4.0.0: type-fest "^0.4.1" write-json-file "^3.2.0" +ws@8.18.3: + version "8.18.3" + resolved "https://registry.npmjs.org/ws/-/ws-8.18.3.tgz#b56b88abffde62791c639170400c93dcb0c95472" + integrity sha512-PEIGCY5tSlUt50cqyMXfCzX+oOPqN0vuGqWzbcJ2xvnkzkq46oOpz7dQaTDBdfICb4N14+GARUDw2XV2N4tvzg== + ws@^7.5.10: version "7.5.10" resolved "https://registry.npmjs.org/ws/-/ws-7.5.10.tgz#58b5c20dc281633f6c19113f39b349bd8bd558d9" From 3ab8bc04fee11735f6262ab773aa46e7ca61499b Mon Sep 17 00:00:00 2001 From: Daniel La Rocque Date: Fri, 1 Aug 2025 10:01:52 -0400 Subject: [PATCH 02/10] Add additional environment checks --- packages/ai/src/ws/websocket-handler.ts | 43 ++++++++++++++++++------- 1 file changed, 32 insertions(+), 11 deletions(-) diff --git a/packages/ai/src/ws/websocket-handler.ts b/packages/ai/src/ws/websocket-handler.ts index 754b030729e..3f42291639b 100644 --- a/packages/ai/src/ws/websocket-handler.ts +++ b/packages/ai/src/ws/websocket-handler.ts @@ -83,25 +83,46 @@ export interface WebSocketHandler { * @internal */ export function createWebSocketHandler(): WebSocketHandler { - // `isNode()` is replaced with a static boolean during build time to enable tree shaking + // `isNode()` is replaced with a static boolean during build time so this block will be + // tree-shaken in browser builds. if (isNode()) { - const [major] = process.versions.node.split('.').map(Number); - if (major < 22) { + // 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) { + 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.js ` + + `version 22 or higher for native WebSocket support.` + ); + } + return new NodeWebSocketHandler(); + } + } + + // `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 { throw new AIError( AIErrorCode.UNSUPPORTED, - 'The Live feature requires Node version 22 or higher for native WebSocket support.' + 'The WebSocket API is not available in this browser-like environment. ' + + 'The Firebase AI "Live" feature is not supported here. It is supported in ' + + 'standard browser windows, Web Workers with WebSocket support, and Node >= 22.' ); } - return new NodeWebSocketHandler(); - } - - // `isBrowser()` is replaced with a static boolean during build time to enable tree shaking - if (isBrowser() && window.WebSocket) { - return new BrowserWebSocketHandler(); } throw new AIError( AIErrorCode.UNSUPPORTED, - 'A WebSocket API is not available in this environment.' + 'This environment is not supported by the "Live" feature. ' + + 'Supported environments are modern web browsers and Node >= 22.' ); } From 5424470cd628337bec377c49cdbf50d76966b23f Mon Sep 17 00:00:00 2001 From: Daniel La Rocque Date: Fri, 1 Aug 2025 10:05:20 -0400 Subject: [PATCH 03/10] Clarify usage of `window` --- packages/ai/src/ws/browser-websocket-handler.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/ai/src/ws/browser-websocket-handler.ts b/packages/ai/src/ws/browser-websocket-handler.ts index 1a55cf436df..af617cfdac9 100644 --- a/packages/ai/src/ws/browser-websocket-handler.ts +++ b/packages/ai/src/ws/browser-websocket-handler.ts @@ -21,7 +21,7 @@ import { WebSocketHandler } from './websocket-handler'; /** * A WebSocketHandler implementation for the browser environment. - * It uses the native `WebSocket` class available on the `window` object. + * It uses the native `WebSocket`. * @internal */ export class BrowserWebSocketHandler implements WebSocketHandler { From 647eb21fd81761d685749077b50aeb03842bb4fe Mon Sep 17 00:00:00 2001 From: Daniel La Rocque Date: Fri, 1 Aug 2025 10:26:47 -0400 Subject: [PATCH 04/10] remove util websocket handler --- packages/util/src/websocket-handler.ts | 0 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 packages/util/src/websocket-handler.ts diff --git a/packages/util/src/websocket-handler.ts b/packages/util/src/websocket-handler.ts deleted file mode 100644 index e69de29bb2d..00000000000 From 7154af64c0f5f0345abc04b02cb9e60aa7d997c8 Mon Sep 17 00:00:00 2001 From: Daniel La Rocque Date: Tue, 5 Aug 2025 12:11:55 -0400 Subject: [PATCH 05/10] Use platform/ convention --- packages/ai/package.json | 2 +- packages/ai/rollup.config.js | 11 ++-- .../browser/websocket.test.ts} | 6 +- .../browser/websocket.ts} | 38 ++++++++--- .../node/websocket.test.ts} | 10 +-- .../node/websocket.ts} | 63 ++++++++++++++----- .../websocket.ts} | 20 ++---- scripts/build/rollup_generate_alias_config.js | 27 ++++++++ .../rollup_get_environment_replacements.js | 49 --------------- 9 files changed, 122 insertions(+), 104 deletions(-) rename packages/ai/src/{ws/browser-websocket-handler.test.ts => platform/browser/websocket.test.ts} (98%) rename packages/ai/src/{ws/browser-websocket-handler.ts => platform/browser/websocket.ts} (75%) rename packages/ai/src/{ws/node-websocket-handler.test.ts => platform/node/websocket.test.ts} (94%) rename packages/ai/src/{ws/node-websocket-handler.ts => platform/node/websocket.ts} (65%) rename packages/ai/src/{ws/websocket-handler.ts => platform/websocket.ts} (77%) create mode 100644 scripts/build/rollup_generate_alias_config.js delete mode 100644 scripts/build/rollup_get_environment_replacements.js diff --git a/packages/ai/package.json b/packages/ai/package.json index 785525cb07e..d335dfd3ab8 100644 --- a/packages/ai/package.json +++ b/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", diff --git a/packages/ai/rollup.config.js b/packages/ai/rollup.config.js index 13389b427a1..016698824fb 100644 --- a/packages/ai/rollup.config.js +++ b/packages/ai/rollup.config.js @@ -15,6 +15,7 @@ * 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'; @@ -22,8 +23,8 @@ 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) }) ], diff --git a/packages/ai/src/ws/browser-websocket-handler.test.ts b/packages/ai/src/platform/browser/websocket.test.ts similarity index 98% rename from packages/ai/src/ws/browser-websocket-handler.test.ts rename to packages/ai/src/platform/browser/websocket.test.ts index 4e1500f34d4..cc3c2cf0d46 100644 --- a/packages/ai/src/ws/browser-websocket-handler.test.ts +++ b/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); }); }); -}); +}); \ No newline at end of file diff --git a/packages/ai/src/ws/browser-websocket-handler.ts b/packages/ai/src/platform/browser/websocket.ts similarity index 75% rename from packages/ai/src/ws/browser-websocket-handler.ts rename to packages/ai/src/platform/browser/websocket.ts index af617cfdac9..32e47039001 100644 --- a/packages/ai/src/ws/browser-websocket-handler.ts +++ b/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 { - 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 => { 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}.` ); } }; diff --git a/packages/ai/src/ws/node-websocket-handler.test.ts b/packages/ai/src/platform/node/websocket.test.ts similarity index 94% rename from packages/ai/src/ws/node-websocket-handler.test.ts rename to packages/ai/src/platform/node/websocket.test.ts index 22f117332fc..020e2b3166d 100644 --- a/packages/ai/src/ws/node-websocket-handler.test.ts +++ b/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; }); }); -}); +}); \ No newline at end of file diff --git a/packages/ai/src/ws/node-websocket-handler.ts b/packages/ai/src/platform/node/websocket.ts similarity index 65% rename from packages/ai/src/ws/node-websocket-handler.ts rename to packages/ai/src/platform/node/websocket.ts index be21b73cd99..bfe48ac06ff 100644 --- a/packages/ai/src/ws/node-websocket-handler.ts +++ b/packages/ai/src/platform/node/websocket.ts @@ -15,13 +15,34 @@ * 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 */ @@ -29,14 +50,18 @@ export class NodeWebSocketHandler implements WebSocketHandler { private ws?: import('ws').WebSocket; async connect(url: string): Promise { - // 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 { 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 }); diff --git a/packages/ai/src/ws/websocket-handler.ts b/packages/ai/src/platform/websocket.ts similarity index 77% rename from packages/ai/src/ws/websocket-handler.ts rename to packages/ai/src/platform/websocket.ts index 3f42291639b..f4f563f9b20 100644 --- a/packages/ai/src/ws/websocket-handler.ts +++ b/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 { diff --git a/scripts/build/rollup_generate_alias_config.js b/scripts/build/rollup_generate_alias_config.js new file mode 100644 index 00000000000..95c435b9fa4 --- /dev/null +++ b/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` + } + ] + }; +} diff --git a/scripts/build/rollup_get_environment_replacements.js b/scripts/build/rollup_get_environment_replacements.js deleted file mode 100644 index 5206a03eb03..00000000000 --- a/scripts/build/rollup_get_environment_replacements.js +++ /dev/null @@ -1,49 +0,0 @@ -/** - * @license - * Copyright 2021 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. - */ - -/** - * Generates the `rollup-plugin-replace` configuration that enables tree-shaking - * of platform-specific code by replacing runtime checks like `isNode()` with boolean literals - * at build time. - * - * For a browser build, `if (isNode()) { ... }` becomes `if (false) { ... }`. - * This allows Rollup to identify and eliminate the Node-specific code branches - * and their imports as dead code. The reverse is true for Node builds. - */ -export function getEnvironmentReplacements(environment) { - const replacements = { - delimiters: ['', ''], // Set to empty strings to replace the entire function call (`isNode()`), not just a variable - preventAssignment: true - }; - - switch (environment) { - case 'browser': - return { - ...replacements, - 'isNode()': 'false', - 'isBrowser()': 'true' - }; - case 'node': - return { - ...replacements, - 'isNode()': 'true', - 'isBrowser()': 'false' - }; - default: - throw new Error(`Unknown build environment: ${environment}`); - } -} From d5b2ff33c09ce7db2f61f25bf76d86bdcebbffac Mon Sep 17 00:00:00 2001 From: Daniel La Rocque Date: Wed, 6 Aug 2025 10:20:27 -0400 Subject: [PATCH 06/10] use built-in global ws --- common/api-review/ai.api.md | 4 + package.json | 1 + packages/ai/package.json | 3 +- .../ai/src/platform/browser/websocket.test.ts | 2 +- packages/ai/src/platform/browser/websocket.ts | 35 ++++---- .../ai/src/platform/node/websocket.test.ts | 14 ++- packages/ai/src/platform/node/websocket.ts | 85 +++++++++---------- .../ai/test-utils/mock-websocket-server.ts | 2 +- 8 files changed, 69 insertions(+), 77 deletions(-) diff --git a/common/api-review/ai.api.md b/common/api-review/ai.api.md index 199b97b10a9..e55ccf4f1e5 100644 --- a/common/api-review/ai.api.md +++ b/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 diff --git a/package.json b/package.json index 8f42b1dfce3..6b39f9c4b90 100644 --- a/package.json +++ b/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" } } diff --git a/packages/ai/package.json b/packages/ai/package.json index d335dfd3ab8..d2ca2142b67 100644 --- a/packages/ai/package.json +++ b/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", diff --git a/packages/ai/src/platform/browser/websocket.test.ts b/packages/ai/src/platform/browser/websocket.test.ts index cc3c2cf0d46..5f211363576 100644 --- a/packages/ai/src/platform/browser/websocket.test.ts +++ b/packages/ai/src/platform/browser/websocket.test.ts @@ -270,4 +270,4 @@ describe('BrowserWebSocketHandler', () => { expect(mockWebSocket.readyState).to.equal(MockBrowserWebSocket.CLOSED); }); }); -}); \ No newline at end of file +}); diff --git a/packages/ai/src/platform/browser/websocket.ts b/packages/ai/src/platform/browser/websocket.ts index 32e47039001..f15afbcbe5d 100644 --- a/packages/ai/src/platform/browser/websocket.ts +++ b/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 { 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 { 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); + } }); } } diff --git a/packages/ai/src/platform/node/websocket.test.ts b/packages/ai/src/platform/node/websocket.test.ts index 020e2b3166d..6fc6fc49c70 100644 --- a/packages/ai/src/platform/node/websocket.test.ts +++ b/packages/ai/src/platform/node/websocket.test.ts @@ -23,7 +23,6 @@ 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); @@ -31,7 +30,7 @@ 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; }); }); -}); \ No newline at end of file +}); diff --git a/packages/ai/src/platform/node/websocket.ts b/packages/ai/src/platform/node/websocket.ts index bfe48ac06ff..dc4b95039e8 100644 --- a/packages/ai/src/platform/node/websocket.ts +++ b/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,16 +26,23 @@ 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.' ); } } @@ -44,24 +50,17 @@ export function createWebSocketHandler(): WebSocketHandler { /** * 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 { 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 => { + 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 { 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); + } }); } } diff --git a/packages/ai/test-utils/mock-websocket-server.ts b/packages/ai/test-utils/mock-websocket-server.ts index f2f83f569ff..e207c95b2be 100644 --- a/packages/ai/test-utils/mock-websocket-server.ts +++ b/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 }); } } } From 5bf69cb9d98864cab76107987f12efdcd5c24d4a Mon Sep 17 00:00:00 2001 From: Daniel La Rocque Date: Wed, 6 Aug 2025 10:34:38 -0400 Subject: [PATCH 07/10] revert tsconfig module bump --- config/tsconfig.base.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/config/tsconfig.base.json b/config/tsconfig.base.json index 33a58c5b02d..ce58a6d700b 100644 --- a/config/tsconfig.base.json +++ b/config/tsconfig.base.json @@ -9,7 +9,7 @@ "es2020", "esnext.WeakRef", ], - "module": "ES2020", + "module": "ES2015", "moduleResolution": "node", "resolveJsonModule": true, "esModuleInterop": true, From 4c81e4f2394f8eeb43109acdadd2fc86efceac57 Mon Sep 17 00:00:00 2001 From: Daniel La Rocque Date: Wed, 6 Aug 2025 10:36:34 -0400 Subject: [PATCH 08/10] cleanup --- common/api-review/ai.api.md | 4 ---- 1 file changed, 4 deletions(-) diff --git a/common/api-review/ai.api.md b/common/api-review/ai.api.md index e55ccf4f1e5..199b97b10a9 100644 --- a/common/api-review/ai.api.md +++ b/common/api-review/ai.api.md @@ -433,10 +433,6 @@ 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 From 685fc198d81b63f46ab314036d3e71964f9e6912 Mon Sep 17 00:00:00 2001 From: Daniel La Rocque Date: Wed, 6 Aug 2025 10:47:37 -0400 Subject: [PATCH 09/10] fix stub ws file --- packages/ai/src/platform/websocket.ts | 60 +++++++-------------------- 1 file changed, 15 insertions(+), 45 deletions(-) diff --git a/packages/ai/src/platform/websocket.ts b/packages/ai/src/platform/websocket.ts index f4f563f9b20..65669c24809 100644 --- a/packages/ai/src/platform/websocket.ts +++ b/packages/ai/src/platform/websocket.ts @@ -15,11 +15,7 @@ * limitations under the License. */ -import { isBrowser, isNode } from '@firebase/util'; -import { AIError } from '../errors'; -import { AIErrorCode } from '../types'; import { NodeWebSocketHandler } from './node/websocket'; -import { BrowserWebSocketHandler } from './browser/websocket'; /** * A standardized interface for interacting with a WebSocket connection. @@ -64,53 +60,27 @@ export interface WebSocketHandler { } /** - * Factory function to create the appropriate WebSocketHandler for the current environment. - * - * Even though the browser and Node >=22 WebSocket APIs are now very similar, - * we use two separate handler classes. There are two reasons for this: + * NOTE: Imports to this these APIs are renamed to either `platform/browser/websocket.ts` or + * `platform/node/websocket.ts` during build time. * - * 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. + * The types are still useful for type-checking during development. + * These are only used during the Node tests, which are ran against non-bundled code. + */ + +/** + * Factory function to create the appropriate WebSocketHandler for the current environment. * - * 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 - * distinct. Using separate classes ensures type correctness for each environment. + * This is only a stub for tests. See the real definitions in `./browser/websocket.ts` and + * `./node/websocket.ts`. * * @internal */ export function createWebSocketHandler(): WebSocketHandler { - if (isNode()) { - 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.js ` + - `version 22 or higher for native WebSocket support.` - ); - } - return new NodeWebSocketHandler(); - } - } - - if (isBrowser()) { - 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 Firebase AI "Live" feature is not supported here. It is supported in ' + - 'standard browser windows, Web Workers with WebSocket support, and Node >= 22.' - ); - } + if (typeof WebSocket === 'undefined') { + throw Error( + 'WebSocket API is not available. Make sure tests are being ran in Node >= 22.' + ); } - throw new AIError( - AIErrorCode.UNSUPPORTED, - 'This environment is not supported by the "Live" feature. ' + - 'Supported environments are modern web browsers and Node >= 22.' - ); + return new NodeWebSocketHandler(); } From 70985379ddaa7e2aaa6293d9d186152a32b291a3 Mon Sep 17 00:00:00 2001 From: Daniel La Rocque Date: Thu, 7 Aug 2025 11:56:21 -0400 Subject: [PATCH 10/10] parse strings, and create error queue --- packages/ai/src/platform/browser/websocket.ts | 69 ++++++++++++++---- packages/ai/src/platform/node/websocket.ts | 73 +++++++++++++++---- 2 files changed, 112 insertions(+), 30 deletions(-) diff --git a/packages/ai/src/platform/browser/websocket.ts b/packages/ai/src/platform/browser/websocket.ts index f15afbcbe5d..16e5bd843a9 100644 --- a/packages/ai/src/platform/browser/websocket.ts +++ b/packages/ai/src/platform/browser/websocket.ts @@ -75,28 +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 => { + let data: string; 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); - } + data = await event.data.text(); + } else if (typeof event.data === 'string') { + data = event.data; } else { - throw new AIError( - AIErrorCode.PARSE_FAILED, - `Failed to parse WebSocket response to JSON. ` + - `Expected data to be a Blob, but was ${typeof event.data}.` + 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; + } + 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 => { @@ -108,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 { @@ -122,6 +157,12 @@ 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 { diff --git a/packages/ai/src/platform/node/websocket.ts b/packages/ai/src/platform/node/websocket.ts index dc4b95039e8..156ee988785 100644 --- a/packages/ai/src/platform/node/websocket.ts +++ b/packages/ai/src/platform/node/websocket.ts @@ -92,28 +92,57 @@ export class NodeWebSocketHandler implements WebSocketHandler { } const messageQueue: unknown[] = []; + const errorQueue: Error[] = []; let resolvePromise: (() => void) | null = null; let isClosed = false; const messageListener = async (event: MessageEvent): Promise => { + let data: string; 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); - } + data = await event.data.text(); + } else if (typeof event.data === 'string') { + data = event.data; } else { - throw new AIError( - AIErrorCode.PARSE_FAILED, - `Failed to parse WebSocket response to JSON. ` + - `Expected data to be a Blob, but was ${typeof event.data}.` + 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; + } + 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 => { @@ -122,15 +151,21 @@ export class NodeWebSocketHandler implements WebSocketHandler { resolvePromise(); resolvePromise = null; } - // Clean up listeners to prevent memory leaks + // 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 { @@ -139,6 +174,12 @@ export class NodeWebSocketHandler 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 { @@ -148,7 +189,7 @@ export class NodeWebSocketHandler implements WebSocketHandler { } this.ws.addEventListener('close', () => resolve(), { once: true }); - // Calling 'close' during these states results in an error. + // Calling 'close' during these states results in an error if ( this.ws.readyState === WebSocket.CLOSED || this.ws.readyState === WebSocket.CONNECTING