feat(client-bedrock-runtime): use WebSocketHandler in browser (#7175)

Author: trivikr

clients/client-bedrock-runtime/package.json

@@ -28,6 +28,7 @@
     "@aws-sdk/middleware-logger": "*",
     "@aws-sdk/middleware-recursion-detection": "*",
     "@aws-sdk/middleware-user-agent": "*",
+    "@aws-sdk/middleware-websocket": "*",
     "@aws-sdk/region-config-resolver": "*",
     "@aws-sdk/token-providers": "*",
     "@aws-sdk/types": "*",

clients/client-bedrock-runtime/src/BedrockRuntimeClient.ts

@@ -18,6 +18,7 @@ import {
   UserAgentInputConfig,
   UserAgentResolvedConfig,
 } from "@aws-sdk/middleware-user-agent";
+import { resolveWebSocketConfig, WebSocketInputConfig, WebSocketResolvedConfig } from "@aws-sdk/middleware-websocket";
 import { EventStreamPayloadHandlerProvider as __EventStreamPayloadHandlerProvider } from "@aws-sdk/types";
 import { RegionInputConfig, RegionResolvedConfig, resolveRegionConfig } from "@smithy/config-resolver";
 import {
@@ -294,6 +295,7 @@ export type BedrockRuntimeClientConfigType = Partial<__SmithyConfiguration<__Htt
   EventStreamSerdeInputConfig &
   HttpAuthSchemeInputConfig &
   EventStreamInputConfig &
+  WebSocketInputConfig &
   ClientInputEndpointParameters;
 /**
  * @public
@@ -316,6 +318,7 @@ export type BedrockRuntimeClientResolvedConfigType = __SmithyResolvedConfigurati
   EventStreamSerdeResolvedConfig &
   HttpAuthSchemeResolvedConfig &
   EventStreamResolvedConfig &
+  WebSocketResolvedConfig &
   ClientResolvedEndpointParameters;
 /**
  * @public
@@ -352,8 +355,9 @@ export class BedrockRuntimeClient extends __Client<
     const _config_7 = resolveEventStreamSerdeConfig(_config_6);
     const _config_8 = resolveHttpAuthSchemeConfig(_config_7);
     const _config_9 = resolveEventStreamConfig(_config_8);
-    const _config_10 = resolveRuntimeExtensions(_config_9, configuration?.extensions || []);
-    this.config = _config_10;
+    const _config_10 = resolveWebSocketConfig(_config_9);
+    const _config_11 = resolveRuntimeExtensions(_config_10, configuration?.extensions || []);
+    this.config = _config_11;
     this.middlewareStack.use(getUserAgentPlugin(this.config));
     this.middlewareStack.use(getRetryPlugin(this.config));
     this.middlewareStack.use(getContentLengthPlugin(this.config));

clients/client-bedrock-runtime/src/commands/InvokeModelWithBidirectionalStreamCommand.ts

@@ -1,5 +1,6 @@
 // smithy-typescript generated code
 import { getEventStreamPlugin } from "@aws-sdk/middleware-eventstream";
+import { getWebSocketPlugin } from "@aws-sdk/middleware-websocket";
 import { getEndpointPlugin } from "@smithy/middleware-endpoint";
 import { getSerdePlugin } from "@smithy/middleware-serde";
 import { Command as $Command } from "@smithy/smithy-client";
@@ -145,6 +146,9 @@ export class InvokeModelWithBidirectionalStreamCommand extends $Command
       getSerdePlugin(config, this.serialize, this.deserialize),
       getEndpointPlugin(config, Command.getEndpointParameterInstructions()),
       getEventStreamPlugin(config),
+      getWebSocketPlugin(config, {
+        headerPrefix: "x-amz-bedrock-",
+      }),
     ];
   })
   .s("AmazonBedrockFrontendService", "InvokeModelWithBidirectionalStream", {

clients/client-bedrock-runtime/src/models/models_0.ts

@@ -1,6 +1,5 @@
 // smithy-typescript generated code
 import { ExceptionOptionType as __ExceptionOptionType, SENSITIVE_STRING } from "@smithy/smithy-client";
-
 import { DocumentType as __DocumentType } from "@smithy/types";
 
 import { BedrockRuntimeServiceException as __BaseException } from "./BedrockRuntimeServiceException";

clients/client-bedrock-runtime/src/runtimeConfig.browser.ts

@@ -3,11 +3,16 @@
 import packageInfo from "../package.json"; // eslint-disable-line
 
 import { Sha256 } from "@aws-crypto/sha256-browser";
+import {
+  WebSocketFetchHandlerOptions,
+  WebSocketFetchHandler as WebSocketRequestHandler,
+  eventStreamPayloadHandlerProvider,
+} from "@aws-sdk/middleware-websocket";
 import { createDefaultUserAgentProvider } from "@aws-sdk/util-user-agent-browser";
 import { DEFAULT_USE_DUALSTACK_ENDPOINT, DEFAULT_USE_FIPS_ENDPOINT } from "@smithy/config-resolver";
 import { eventStreamSerdeProvider } from "@smithy/eventstream-serde-browser";
-import { FetchHttpHandler as RequestHandler, streamCollector } from "@smithy/fetch-http-handler";
-import { invalidFunction, invalidProvider } from "@smithy/invalid-dependency";
+import { FetchHttpHandler as HttpRequestHandler, streamCollector } from "@smithy/fetch-http-handler";
+import { invalidProvider } from "@smithy/invalid-dependency";
 import { calculateBodyLength } from "@smithy/util-body-length-browser";
 import { DEFAULT_MAX_ATTEMPTS, DEFAULT_RETRY_MODE } from "@smithy/util-retry";
 import { BedrockRuntimeClientConfig } from "./BedrockRuntimeClient";
@@ -33,13 +38,17 @@ export const getRuntimeConfig = (config: BedrockRuntimeClientConfig) => {
     defaultUserAgentProvider:
       config?.defaultUserAgentProvider ??
       createDefaultUserAgentProvider({ serviceId: clientSharedValues.serviceId, clientVersion: packageInfo.version }),
-    eventStreamPayloadHandlerProvider:
-      config?.eventStreamPayloadHandlerProvider ??
-      (() => ({ handle: invalidFunction("event stream request is not supported in browser.") })),
+    eventStreamPayloadHandlerProvider: config?.eventStreamPayloadHandlerProvider ?? eventStreamPayloadHandlerProvider,
     eventStreamSerdeProvider: config?.eventStreamSerdeProvider ?? eventStreamSerdeProvider,
     maxAttempts: config?.maxAttempts ?? DEFAULT_MAX_ATTEMPTS,
     region: config?.region ?? invalidProvider("Region is missing"),
-    requestHandler: RequestHandler.create(config?.requestHandler ?? defaultConfigProvider),
+    requestHandler: WebSocketRequestHandler.create(
+      (config?.requestHandler as
+        | WebSocketRequestHandler
+        | WebSocketFetchHandlerOptions
+        | (() => Promise<WebSocketFetchHandlerOptions>)) ?? defaultConfigProvider,
+      HttpRequestHandler.create(defaultConfigProvider)
+    ),
     retryMode: config?.retryMode ?? (async () => (await defaultConfigProvider()).retryMode || DEFAULT_RETRY_MODE),
     sha256: config?.sha256 ?? Sha256,
     streamCollector: config?.streamCollector ?? streamCollector,

codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AddWebsocketPlugin.java

@@ -126,7 +126,7 @@ public Map<String, Consumer<TypeScriptWriter>> getRuntimeConfigWriters(
     }
 
     private static boolean isWebsocketSupported(ServiceShape service) {
-        Set<String> websocketServices = SetUtils.of("Transcribe Streaming", "RekognitionStreaming");
+        Set<String> websocketServices = SetUtils.of("Transcribe Streaming", "RekognitionStreaming", "Bedrock Runtime");
         String serviceId = service.getTrait(ServiceTrait.class).map(ServiceTrait::getSdkId).orElse("");
         return websocketServices.contains(serviceId);
     }
@@ -141,6 +141,8 @@ private static Map<String, Object> getPluginFunctionParams(
             return MapUtils.of("headerPrefix", "x-amzn-transcribe-");
         } else if (serviceId.equals("RekognitionStreaming")) {
             return MapUtils.of("headerPrefix", "x-amz-rekognition-streaming-liveness-");
+        } else if (serviceId.equals("Bedrock Runtime")) {
+            return MapUtils.of("headerPrefix", "x-amz-bedrock-");
         } else {
             throw new CodegenException("Missing endpoint prefix for Websocket plugin of service " + serviceId);
         }

yarn.lock

@@ -3705,6 +3705,7 @@ __metadata:
     "@aws-sdk/middleware-logger": "npm:*"
     "@aws-sdk/middleware-recursion-detection": "npm:*"
     "@aws-sdk/middleware-user-agent": "npm:*"
+    "@aws-sdk/middleware-websocket": "npm:*"
     "@aws-sdk/region-config-resolver": "npm:*"
     "@aws-sdk/token-providers": "npm:*"
     "@aws-sdk/types": "npm:*"