From 5122951a7216f04c20e1a6be87e90578ff3a082d Mon Sep 17 00:00:00 2001
From: George Fu <kuhe@users.noreply.github.com>
Date: Tue, 8 Oct 2024 20:03:38 +0000
Subject: [PATCH] chore: calculate user agent in websocket

---
 .../src/check-features.ts                     | 27 ++++++++++++++++++-
 .../src/user-agent-middleware.ts              | 18 ++++++-------
 2 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/packages/middleware-user-agent/src/check-features.ts b/packages/middleware-user-agent/src/check-features.ts
index e12e0829ab803..0e3b6ed61812b 100644
--- a/packages/middleware-user-agent/src/check-features.ts
+++ b/packages/middleware-user-agent/src/check-features.ts
@@ -6,7 +6,13 @@ import type {
   AwsSdkCredentialsFeatures,
 } from "@aws-sdk/types";
 import type { IHttpRequest } from "@smithy/protocol-http";
-import type { AwsCredentialIdentityProvider, BuildHandlerArguments, Provider } from "@smithy/types";
+import type {
+  AwsCredentialIdentityProvider,
+  BuildHandlerArguments,
+  Provider,
+  RetryStrategy,
+  RetryStrategyV2,
+} from "@smithy/types";
 
 /**
  * @internal
@@ -14,6 +20,7 @@ import type { AwsCredentialIdentityProvider, BuildHandlerArguments, Provider } f
 type PreviouslyResolved = Partial<{
   credentials?: AwsCredentialIdentityProvider;
   accountIdEndpointMode?: Provider<AccountIdEndpointMode>;
+  retryStrategy?: Provider<RetryStrategy | RetryStrategyV2>;
 }>;
 
 /**
@@ -33,6 +40,24 @@ export async function checkFeatures(
 ): Promise<void> {
   // eslint-disable-next-line
   const request = args.request as IHttpRequest;
+
+  if (request.headers?.["smithy-protocol"] === "rpc-v2-cbor") {
+    setFeature(context, "PROTOCOL_RPC_V2_CBOR", "M");
+  }
+
+  if (typeof config.retryStrategy === "function") {
+    const retryStrategy = await config.retryStrategy();
+    if (typeof (retryStrategy as RetryStrategyV2).acquireInitialRetryToken === "function") {
+      if (retryStrategy.constructor?.name?.includes("Adaptive")) {
+        setFeature(context, "RETRY_MODE_ADAPTIVE", "F");
+      } else {
+        setFeature(context, "RETRY_MODE_STANDARD", "E");
+      }
+    } else {
+      setFeature(context, "RETRY_MODE_LEGACY", "D");
+    }
+  }
+
   if (typeof config.accountIdEndpointMode === "function") {
     const endpointV2 = context.endpointV2;
     if (String(endpointV2?.url?.hostname).match(ACCOUNT_ID_ENDPOINT_REGEX)) {
diff --git a/packages/middleware-user-agent/src/user-agent-middleware.ts b/packages/middleware-user-agent/src/user-agent-middleware.ts
index 42c915d8526c7..4438ede668d07 100644
--- a/packages/middleware-user-agent/src/user-agent-middleware.ts
+++ b/packages/middleware-user-agent/src/user-agent-middleware.ts
@@ -3,10 +3,10 @@ import { getUserAgentPrefix } from "@aws-sdk/util-endpoints";
 import { HttpRequest } from "@smithy/protocol-http";
 import {
   AbsoluteLocation,
-  FinalizeHandler,
-  FinalizeHandlerArguments,
-  FinalizeHandlerOutput,
-  FinalizeRequestHandlerOptions,
+  BuildHandler,
+  BuildHandlerArguments,
+  BuildHandlerOptions,
+  BuildHandlerOutput,
   HandlerExecutionContext,
   MetadataBearer,
   Pluggable,
@@ -41,10 +41,10 @@ import { encodeFeatures } from "./encode-features";
 export const userAgentMiddleware =
   (options: UserAgentResolvedConfig) =>
   <Output extends MetadataBearer>(
-    next: FinalizeHandler<any, any>,
+    next: BuildHandler<any, any>,
     context: HandlerExecutionContext | AwsHandlerExecutionContext
-  ): FinalizeHandler<any, any> =>
-  async (args: FinalizeHandlerArguments<any>): Promise<FinalizeHandlerOutput<Output>> => {
+  ): BuildHandler<any, any> =>
+  async (args: BuildHandlerArguments<any>): Promise<BuildHandlerOutput<Output>> => {
     const { request } = args;
     if (!HttpRequest.isInstance(request)) {
       return next(args);
@@ -130,9 +130,9 @@ const escapeUserAgent = (userAgentPair: UserAgentPair): string => {
     }, "") as string;
 };
 
-export const getUserAgentMiddlewareOptions: FinalizeRequestHandlerOptions & AbsoluteLocation = {
+export const getUserAgentMiddlewareOptions: BuildHandlerOptions & AbsoluteLocation = {
   name: "getUserAgentMiddleware",
-  step: "finalizeRequest",
+  step: "build",
   priority: "low",
   tags: ["SET_USER_AGENT", "USER_AGENT"],
   override: true,