Live Refactor WebSocket (#9212)

* refactor websocket

* remove test ws server

Author: dlarocque

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

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}/`))
   }
 ];
 

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);
 }

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';
 

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,

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';

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';