diff --git a/package.json b/package.json index 6b39f9c4b90..39455ef1161 100644 --- a/package.json +++ b/package.json @@ -86,7 +86,6 @@ "@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", @@ -159,7 +158,6 @@ "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/rollup.config.js b/packages/ai/rollup.config.js index 016698824fb..7ebbff4f2f5 100644 --- a/packages/ai/rollup.config.js +++ b/packages/ai/rollup.config.js @@ -15,7 +15,6 @@ * 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'; @@ -24,7 +23,6 @@ import pkg from './package.json'; import tsconfig from './tsconfig.json'; import { generateBuildTargetReplaceConfig } from '../../scripts/build/rollup_replace_build_target'; 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) @@ -57,16 +55,14 @@ const browserBuilds = [ sourcemap: true }, plugins: [ - alias(generateAliasConfig('browser')), ...buildPlugins, replace({ ...generateBuildTargetReplaceConfig('esm', 2020), - '__PACKAGE_VERSION__': pkg.version + __PACKAGE_VERSION__: pkg.version }), emitModulePackageFile() ], - external: id => - id === 'ws' || deps.some(dep => id === dep || id.startsWith(`${dep}/`)) + external: id => deps.some(dep => id === dep || id.startsWith(`${dep}/`)) }, { input: 'src/index.ts', @@ -76,15 +72,13 @@ const browserBuilds = [ sourcemap: true }, plugins: [ - alias(generateAliasConfig('browser')), ...buildPlugins, replace({ ...generateBuildTargetReplaceConfig('cjs', 2020), - '__PACKAGE_VERSION__': pkg.version + __PACKAGE_VERSION__: pkg.version }) ], - external: id => - id === 'ws' || deps.some(dep => id === dep || id.startsWith(`${dep}/`)) + external: id => deps.some(dep => id === dep || id.startsWith(`${dep}/`)) } ]; @@ -97,14 +91,12 @@ const nodeBuilds = [ sourcemap: true }, plugins: [ - alias(generateAliasConfig('node')), ...buildPlugins, replace({ ...generateBuildTargetReplaceConfig('esm', 2020) }) ], - external: id => - id === 'ws' || deps.some(dep => id === dep || id.startsWith(`${dep}/`)) + external: id => deps.some(dep => id === dep || id.startsWith(`${dep}/`)) }, { input: 'src/index.node.ts', @@ -114,14 +106,12 @@ const nodeBuilds = [ sourcemap: true }, plugins: [ - alias(generateAliasConfig('node')), ...buildPlugins, replace({ ...generateBuildTargetReplaceConfig('cjs', 2020) }) ], - external: id => - id === 'ws' || deps.some(dep => id === dep || id.startsWith(`${dep}/`)) + external: id => deps.some(dep => id === dep || id.startsWith(`${dep}/`)) } ]; diff --git a/packages/ai/src/api.ts b/packages/ai/src/api.ts index af177f750d4..418c17bb49c 100644 --- a/packages/ai/src/api.ts +++ b/packages/ai/src/api.ts @@ -37,7 +37,7 @@ import { } from './models'; import { encodeInstanceIdentifier } from './helpers'; import { GoogleAIBackend } from './backend'; -import { createWebSocketHandler } from './platform/websocket'; +import { WebSocketHandlerImpl } from './websocket'; export { ChatSession } from './methods/chat-session'; export { LiveSession } from './methods/live-session'; @@ -164,6 +164,6 @@ export function getLiveGenerativeModel( `Must provide a model name for getLiveGenerativeModel. Example: getLiveGenerativeModel(ai, { model: 'my-model-name' })` ); } - const webSocketHandler = createWebSocketHandler(); + const webSocketHandler = new WebSocketHandlerImpl(); return new LiveGenerativeModel(ai, modelParams, webSocketHandler); } diff --git a/packages/ai/src/methods/live-session.test.ts b/packages/ai/src/methods/live-session.test.ts index 454c20402c2..7db9daaebe6 100644 --- a/packages/ai/src/methods/live-session.test.ts +++ b/packages/ai/src/methods/live-session.test.ts @@ -26,7 +26,7 @@ import { LiveServerToolCallCancellation } from '../types'; import { LiveSession } from './live-session'; -import { WebSocketHandler } from '../platform/websocket'; +import { WebSocketHandler } from '../websocket'; import { AIError } from '../errors'; import { logger } from '../logger'; diff --git a/packages/ai/src/methods/live-session.ts b/packages/ai/src/methods/live-session.ts index 40d35800737..b257d0a5787 100644 --- a/packages/ai/src/methods/live-session.ts +++ b/packages/ai/src/methods/live-session.ts @@ -26,7 +26,7 @@ import { } from '../public-types'; import { formatNewContent } from '../requests/request-helpers'; import { AIError } from '../errors'; -import { WebSocketHandler } from '../platform/websocket'; +import { WebSocketHandler } from '../websocket'; import { logger } from '../logger'; import { _LiveClientContent, diff --git a/packages/ai/src/models/live-generative-model.test.ts b/packages/ai/src/models/live-generative-model.test.ts index 6278d5e182f..495f340b846 100644 --- a/packages/ai/src/models/live-generative-model.test.ts +++ b/packages/ai/src/models/live-generative-model.test.ts @@ -20,7 +20,7 @@ import sinonChai from 'sinon-chai'; import chaiAsPromised from 'chai-as-promised'; import { AI } from '../public-types'; import { LiveSession } from '../methods/live-session'; -import { WebSocketHandler } from '../platform/websocket'; +import { WebSocketHandler } from '../websocket'; import { GoogleAIBackend } from '../backend'; import { LiveGenerativeModel } from './live-generative-model'; import { AIError } from '../errors'; diff --git a/packages/ai/src/models/live-generative-model.ts b/packages/ai/src/models/live-generative-model.ts index 67d74bba95e..251df095202 100644 --- a/packages/ai/src/models/live-generative-model.ts +++ b/packages/ai/src/models/live-generative-model.ts @@ -28,7 +28,7 @@ import { Tool, ToolConfig } from '../public-types'; -import { WebSocketHandler } from '../platform/websocket'; +import { WebSocketHandler } from '../websocket'; import { WebSocketUrl } from '../requests/request'; import { formatSystemInstruction } from '../requests/request-helpers'; import { _LiveClientSetup } from '../types/live-responses'; diff --git a/packages/ai/src/platform/node/websocket.test.ts b/packages/ai/src/platform/node/websocket.test.ts deleted file mode 100644 index 6fc6fc49c70..00000000000 --- a/packages/ai/src/platform/node/websocket.test.ts +++ /dev/null @@ -1,143 +0,0 @@ -/** - * @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 { TextEncoder } from 'util'; -import { MockWebSocketServer } from '../../../test-utils/mock-websocket-server'; -import { WebSocketHandler } from '../websocket'; -import { NodeWebSocketHandler } from './websocket'; - -use(sinonChai); -use(chaiAsPromised); - -const TEST_PORT = 9003; -const TEST_URL = `ws://localhost:${TEST_PORT}`; - -describe('NodeWebSocketHandler', () => { - 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', async () => { - const wrongPortUrl = `ws://wrongUrl:9000`; - await expect(handler.connect(wrongPortUrl)).to.be.rejected; - }); - }); - - 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/platform/node/websocket.ts b/packages/ai/src/platform/node/websocket.ts deleted file mode 100644 index 36a963953bf..00000000000 --- a/packages/ai/src/platform/node/websocket.ts +++ /dev/null @@ -1,211 +0,0 @@ -/** - * @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 { logger } from '../../logger'; -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.` - ); - } 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 ' + - 'modern browser windows, Web Workers with WebSocket support, and Node >= 22.' - ); - } -} - -/** - * A WebSocketHandler implementation for Node >= 22. - * - * Node 22 is the minimum version that offers the built-in global `WebSocket` API. - * - * @internal - */ -export class NodeWebSocketHandler implements WebSocketHandler { - private ws?: WebSocket; - - async connect(url: string): Promise { - return new Promise(async (resolve, reject) => { - this.ws = new WebSocket(url); - this.ws.binaryType = 'blob'; - this.ws!.addEventListener('open', () => resolve(), { once: true }); - this.ws.addEventListener( - 'error', - () => - reject( - new AIError( - AIErrorCode.FETCH_ERROR, - `Error event raised on WebSocket` - ) - ), - { 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 { - if (!this.ws) { - throw new AIError( - AIErrorCode.REQUEST_ERROR, - 'WebSocket is not connected.' - ); - } - - 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) { - data = await event.data.text(); - } else if (typeof event.data === 'string') { - data = event.data; - } else { - errorQueue.push( - new AIError( - AIErrorCode.PARSE_FAILED, - `Failed to parse WebSocket response. Expected data to be a Blob or string, but was ${typeof event.data}.` - ) - ); - if (resolvePromise) { - resolvePromise(); - resolvePromise = null; - } - 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 = (event: CloseEvent): void => { - if (event.reason) { - logger.warn( - `WebSocket connection closed by the server with reason: ${event.reason}` - ); - } - 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?.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 { - await new Promise(resolve => { - resolvePromise = resolve; - }); - } - } - - // 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 { - 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.readyState === WebSocket.CLOSED || - this.ws.readyState === WebSocket.CONNECTING - ) { - return resolve(); - } - - if (this.ws.readyState !== WebSocket.CLOSING) { - this.ws.close(code, reason); - } - }); - } -} diff --git a/packages/ai/src/platform/websocket.ts b/packages/ai/src/platform/websocket.ts deleted file mode 100644 index 65669c24809..00000000000 --- a/packages/ai/src/platform/websocket.ts +++ /dev/null @@ -1,86 +0,0 @@ -/** - * @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 { NodeWebSocketHandler } from './node/websocket'; - -/** - * 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; -} - -/** - * NOTE: Imports to this these APIs are renamed to either `platform/browser/websocket.ts` or - * `platform/node/websocket.ts` during build time. - * - * 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. - * - * 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 (typeof WebSocket === 'undefined') { - throw Error( - 'WebSocket API is not available. Make sure tests are being ran in Node >= 22.' - ); - } - - return new NodeWebSocketHandler(); -} diff --git a/packages/ai/src/platform/browser/websocket.test.ts b/packages/ai/src/websocket.test.ts similarity index 86% rename from packages/ai/src/platform/browser/websocket.test.ts rename to packages/ai/src/websocket.test.ts index 07e35cd8299..6d8f08282e7 100644 --- a/packages/ai/src/platform/browser/websocket.test.ts +++ b/packages/ai/src/websocket.test.ts @@ -19,20 +19,19 @@ import { expect, use } from 'chai'; import sinon, { SinonFakeTimers, SinonStub } from 'sinon'; import sinonChai from 'sinon-chai'; import chaiAsPromised from 'chai-as-promised'; -import { isBrowser } from '@firebase/util'; -import { BrowserWebSocketHandler } from './websocket'; -import { AIError } from '../../errors'; +import { WebSocketHandlerImpl } from './websocket'; +import { AIError } from './errors'; use(sinonChai); use(chaiAsPromised); -class MockBrowserWebSocket { +class MockWebSocket { static CONNECTING = 0; static OPEN = 1; static CLOSING = 2; static CLOSED = 3; - readyState: number = MockBrowserWebSocket.CONNECTING; + readyState: number = MockWebSocket.CONNECTING; sentMessages: Array = []; url: string; private listeners: Map> = new Map(); @@ -42,7 +41,7 @@ class MockBrowserWebSocket { } send(data: string | ArrayBuffer): void { - if (this.readyState !== MockBrowserWebSocket.OPEN) { + if (this.readyState !== MockWebSocket.OPEN) { throw new Error('WebSocket is not in OPEN state'); } this.sentMessages.push(data); @@ -50,14 +49,14 @@ class MockBrowserWebSocket { close(): void { if ( - this.readyState === MockBrowserWebSocket.CLOSED || - this.readyState === MockBrowserWebSocket.CLOSING + this.readyState === MockWebSocket.CLOSED || + this.readyState === MockWebSocket.CLOSING ) { return; } - this.readyState = MockBrowserWebSocket.CLOSING; + this.readyState = MockWebSocket.CLOSING; setTimeout(() => { - this.readyState = MockBrowserWebSocket.CLOSED; + this.readyState = MockWebSocket.CLOSED; this.dispatchEvent(new Event('close')); }, 10); } @@ -78,7 +77,7 @@ class MockBrowserWebSocket { } triggerOpen(): void { - this.readyState = MockBrowserWebSocket.OPEN; + this.readyState = MockWebSocket.OPEN; this.dispatchEvent(new Event('open')); } @@ -91,24 +90,21 @@ class MockBrowserWebSocket { } } -describe('BrowserWebSocketHandler', () => { - let handler: BrowserWebSocketHandler; - let mockWebSocket: MockBrowserWebSocket; +describe('WebSocketHandlerImpl', () => { + let handler: WebSocketHandlerImpl; + let mockWebSocket: MockWebSocket; 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; - }); + webSocketStub = sinon + .stub(globalThis, 'WebSocket') + .callsFake((url: string) => { + mockWebSocket = new MockWebSocket(url); + return mockWebSocket as any; + }); clock = sinon.useFakeTimers(); - handler = new BrowserWebSocketHandler(); + handler = new WebSocketHandlerImpl(); }); afterEach(() => { @@ -267,7 +263,7 @@ describe('BrowserWebSocketHandler', () => { await expect(closePromise).to.be.fulfilled; await expect(listenPromise).to.be.fulfilled; - expect(mockWebSocket.readyState).to.equal(MockBrowserWebSocket.CLOSED); + expect(mockWebSocket.readyState).to.equal(MockWebSocket.CLOSED); }); }); }); diff --git a/packages/ai/src/platform/browser/websocket.ts b/packages/ai/src/websocket.ts similarity index 68% rename from packages/ai/src/platform/browser/websocket.ts rename to packages/ai/src/websocket.ts index 884bef776e6..fa34f2d48c3 100644 --- a/packages/ai/src/platform/browser/websocket.ts +++ b/packages/ai/src/websocket.ts @@ -15,36 +15,76 @@ * limitations under the License. */ -import { AIError } from '../../errors'; -import { logger } from '../../logger'; -import { AIErrorCode } from '../../types'; -import { WebSocketHandler } from '../websocket'; - -export function createWebSocketHandler(): WebSocketHandler { - if (typeof WebSocket === 'undefined') { - throw new AIError( - AIErrorCode.UNSUPPORTED, - 'The WebSocket API is not available in this browser-like environment. ' + - 'The "Live" feature is not supported here. It is supported in ' + - 'modern browser windows, Web Workers with WebSocket support, and Node >= 22.' - ); - } +import { AIError } from './errors'; +import { logger } from './logger'; +import { AIErrorCode } from './types'; - return new BrowserWebSocketHandler(); +/** + * 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; } /** - * A WebSocketHandler implementation for the browser environment. - * It uses the native `WebSocket`. + * A wrapper for the native `WebSocket` available in both Browsers and Node >= 22. * * @internal */ -export class BrowserWebSocketHandler implements WebSocketHandler { +export class WebSocketHandlerImpl implements WebSocketHandler { private ws?: WebSocket; + constructor() { + if (typeof WebSocket === 'undefined') { + throw new AIError( + AIErrorCode.UNSUPPORTED, + 'The WebSocket API is not available in this environment. ' + + 'The "Live" feature is not supported here. It is supported in ' + + 'modern browser windows, Web Workers with WebSocket support, and Node >= 22.' + ); + } + } + connect(url: string): Promise { return new Promise((resolve, reject) => { this.ws = new WebSocket(url); + this.ws.binaryType = 'blob'; // Only important to set in Node this.ws.addEventListener('open', () => resolve(), { once: true }); this.ws.addEventListener( 'error', @@ -58,9 +98,11 @@ export class BrowserWebSocketHandler implements WebSocketHandler { { once: true } ); this.ws!.addEventListener('close', (closeEvent: CloseEvent) => { - logger.warn( - `WebSocket connection closed by server. Reason: '${closeEvent.reason}'` - ); + if (closeEvent.reason) { + logger.warn( + `WebSocket connection closed by server. Reason: '${closeEvent.reason}'` + ); + } }); }); } diff --git a/packages/ai/test-utils/mock-websocket-server.ts b/packages/ai/test-utils/mock-websocket-server.ts deleted file mode 100644 index e207c95b2be..00000000000 --- a/packages/ai/test-utils/mock-websocket-server.ts +++ /dev/null @@ -1,76 +0,0 @@ -/** - * @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, { binary: true }); - } - } - } - - 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/yarn.lock b/yarn.lock index cfc5a206b1f..530aafff84f 100644 --- a/yarn.lock +++ b/yarn.lock @@ -3321,13 +3321,6 @@ 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" @@ -17037,11 +17030,6 @@ 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"