From 7748297f2c86bc93fc39e08371ff3d1637ecf69b Mon Sep 17 00:00:00 2001
From: Steven Yuan <yuasteve@amazon.com>
Date: Wed, 13 Dec 2023 10:50:10 -0500
Subject: [PATCH] Add `eventStream` smithy context (#1111)

---
 .../integration/AddEventStreamDependency.java | 29 +++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/integration/AddEventStreamDependency.java b/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/integration/AddEventStreamDependency.java
index 3a5c78b74a9..03144e11a93 100644
--- a/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/integration/AddEventStreamDependency.java
+++ b/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/integration/AddEventStreamDependency.java
@@ -27,9 +27,13 @@
 import software.amazon.smithy.model.shapes.OperationShape;
 import software.amazon.smithy.model.shapes.ServiceShape;
 import software.amazon.smithy.typescript.codegen.LanguageTarget;
+import software.amazon.smithy.typescript.codegen.TypeScriptCodegenContext;
 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.sections.SmithyContextCodeSection;
+import software.amazon.smithy.utils.CodeInterceptor;
+import software.amazon.smithy.utils.CodeSection;
 import software.amazon.smithy.utils.ListUtils;
 import software.amazon.smithy.utils.MapUtils;
 import software.amazon.smithy.utils.SmithyInternalApi;
@@ -99,6 +103,31 @@ public Map<String, Consumer<TypeScriptWriter>> getRuntimeConfigWriters(
         }
     }
 
+    @Override
+    public List<? extends CodeInterceptor<? extends CodeSection, TypeScriptWriter>> interceptors(
+        TypeScriptCodegenContext codegenContext
+    ) {
+        return List.of(CodeInterceptor.appender(SmithyContextCodeSection.class, (w, s) -> {
+            EventStreamIndex eventStreamIndex = EventStreamIndex.of(s.getModel());
+            boolean input = eventStreamIndex.getInputInfo(s.getOperation()).isPresent();
+            boolean output = eventStreamIndex.getOutputInfo(s.getOperation()).isPresent();
+            // If not event streaming for I/O, don't write anything
+            if (!input && !output) {
+                return;
+            }
+            // Otherwise, write present input and output streaming
+            w.writeDocs("@internal");
+            w.openBlock("eventStream: {", "},", () -> {
+                if (input) {
+                    w.write("input: true,");
+                }
+                if (output) {
+                    w.write("output: true,");
+                }
+            });
+        }));
+    }
+
     private static boolean hasEventStream(
             Model model,
             ServiceShape service