From a12b2eb9acf1c14696ae02489545718bcc634964 Mon Sep 17 00:00:00 2001
From: Steven Yuan <yuasteve@amazon.com>
Date: Wed, 13 Dec 2023 13:39:04 -0500
Subject: [PATCH] feat(experimentalIdentityAndAuth): add event streaming input
 integrations relative to auth (#5578)

* feat(experimentalIdentityAndAuth): remove `AddEventStreamHandlingDependency` experimental branch

* feat(experimentalIdentityAndAuth): add order for `AddWebsocketPlugin`
---
 .../AddEventStreamHandlingDependency.java      | 18 +++++++-----------
 .../typescript/codegen/AddWebsocketPlugin.java |  7 +++++++
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AddEventStreamHandlingDependency.java b/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AddEventStreamHandlingDependency.java
index fd9f1849d4d72..de1936320c9b9 100644
--- a/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AddEventStreamHandlingDependency.java
+++ b/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AddEventStreamHandlingDependency.java
@@ -33,6 +33,7 @@
 import software.amazon.smithy.typescript.codegen.TypeScriptDependency;
 import software.amazon.smithy.typescript.codegen.TypeScriptSettings;
 import software.amazon.smithy.typescript.codegen.TypeScriptWriter;
+import software.amazon.smithy.typescript.codegen.auth.http.integration.AddHttpAuthSchemePlugin;
 import software.amazon.smithy.typescript.codegen.integration.RuntimeClientPlugin;
 import software.amazon.smithy.typescript.codegen.integration.TypeScriptIntegration;
 import software.amazon.smithy.utils.ListUtils;
@@ -45,6 +46,12 @@
  */
 @SmithyInternalApi
 public class AddEventStreamHandlingDependency implements TypeScriptIntegration {
+
+    @Override
+    public List<String> runAfter() {
+        return List.of(AddHttpAuthSchemePlugin.class.getCanonicalName());
+    }
+
     @Override
     public List<RuntimeClientPlugin> getClientPlugins() {
         return ListUtils.of(
@@ -52,13 +59,11 @@ public List<RuntimeClientPlugin> getClientPlugins() {
                         .withConventions(AwsDependency.MIDDLEWARE_EVENTSTREAM.dependency,
                                 "EventStream", HAS_CONFIG)
                         .servicePredicate(AddEventStreamHandlingDependency::hasEventStreamInput)
-                        .settingsPredicate((m, s, settings) -> !settings.getExperimentalIdentityAndAuth())
                         .build(),
                 RuntimeClientPlugin.builder()
                         .withConventions(AwsDependency.MIDDLEWARE_EVENTSTREAM.dependency,
                                 "EventStream", HAS_MIDDLEWARE)
                         .operationPredicate(AddEventStreamHandlingDependency::hasEventStreamInput)
-                        .settingsPredicate((m, s, settings) -> !settings.getExperimentalIdentityAndAuth())
                         .build()
         );
     }
@@ -70,10 +75,6 @@ public void addConfigInterfaceFields(
             SymbolProvider symbolProvider,
             TypeScriptWriter writer
     ) {
-        if (settings.getExperimentalIdentityAndAuth()) {
-            return;
-        }
-        // feat(experimentalIdentityAndAuth): control branch for event stream handler interface fields
         if (hasEventStreamInput(model, settings.getService(model))) {
             writer.addImport("EventStreamPayloadHandlerProvider", "__EventStreamPayloadHandlerProvider",
                     TypeScriptDependency.AWS_SDK_TYPES);
@@ -94,11 +95,6 @@ public Map<String, Consumer<TypeScriptWriter>> getRuntimeConfigWriters(
         if (!hasEventStreamInput(model, service)) {
             return Collections.emptyMap();
         }
-
-        if (settings.getExperimentalIdentityAndAuth()) {
-            return Collections.emptyMap();
-        }
-        // feat(experimentalIdentityAndAuth): control branch for event stream handler runtime config
         switch (target) {
             case NODE:
                 return MapUtils.of("eventStreamPayloadHandlerProvider", writer -> {
diff --git a/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AddWebsocketPlugin.java b/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AddWebsocketPlugin.java
index c68cbe151f899..618e121a3fe4d 100644
--- a/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AddWebsocketPlugin.java
+++ b/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AddWebsocketPlugin.java
@@ -31,6 +31,7 @@
 import software.amazon.smithy.typescript.codegen.TypeScriptDependency;
 import software.amazon.smithy.typescript.codegen.TypeScriptSettings;
 import software.amazon.smithy.typescript.codegen.TypeScriptWriter;
+import software.amazon.smithy.typescript.codegen.auth.http.integration.AddHttpAuthSchemePlugin;
 import software.amazon.smithy.typescript.codegen.integration.RuntimeClientPlugin;
 import software.amazon.smithy.typescript.codegen.integration.TypeScriptIntegration;
 import software.amazon.smithy.utils.ListUtils;
@@ -44,6 +45,12 @@
  **/
 @SmithyInternalApi
 public class AddWebsocketPlugin implements TypeScriptIntegration {
+
+    @Override
+    public List<String> runAfter() {
+        return List.of(AddHttpAuthSchemePlugin.class.getCanonicalName());
+    }
+
     @Override
     public List<RuntimeClientPlugin> getClientPlugins() {
         return ListUtils.of(