From 64e26f4cc3f8937cee01edacaefb520c91d8647a Mon Sep 17 00:00:00 2001
From: Jaykumar Gosar <gosar@amazon.com>
Date: Thu, 12 May 2022 12:47:28 -0700
Subject: [PATCH 1/2] Make SmithyIntegrations available to
 CreateContextDirective

For now, only made it available from CreateContextDirective, so subclasses
CodegenContext can access it. But it likely can be made part of parent
CodegenContext interface.
---
 .../core/directed/CodegenDirector.java        | 14 ++++++-------
 .../core/directed/CreateContextDirective.java | 20 +++++++++++++++++--
 .../core/directed/DirectedCodegen.java        |  5 +++--
 .../core/directed/CodegenDirectorTest.java    |  5 ++---
 .../core/directed/TestIntegration.java        |  2 ++
 5 files changed, 32 insertions(+), 14 deletions(-)
 create mode 100644 smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/TestIntegration.java

diff --git a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/CodegenDirector.java b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/CodegenDirector.java
index 1fdecdc1e3f..b880745dbbd 100644
--- a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/CodegenDirector.java
+++ b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/CodegenDirector.java
@@ -71,7 +71,7 @@ public final class CodegenDirector<
     private S settings;
     private FileManifest fileManifest;
     private Supplier<Iterable<I>> integrationFinder;
-    private DirectedCodegen<C, S> directedCodegen;
+    private DirectedCodegen<C, S, I> directedCodegen;
     private final List<BiFunction<Model, ModelTransformer, Model>> transforms = new ArrayList<>();
 
     /**
@@ -124,7 +124,7 @@ public void service(ShapeId service) {
      *
      * @param directedCodegen Directed code generator to run.
      */
-    public void directedCodegen(DirectedCodegen<C, S> directedCodegen) {
+    public void directedCodegen(DirectedCodegen<C, S, I> directedCodegen) {
         this.directedCodegen = directedCodegen;
     }
 
@@ -269,7 +269,7 @@ public void run() {
 
         SymbolProvider provider = createSymbolProvider(integrations, serviceShape);
 
-        C context = createContext(serviceShape, provider);
+        C context = createContext(serviceShape, provider, integrations);
 
         // After the context is created, it holds the model that should be used for rest of codegen. So `model` should
         // not really be used directly after this point in the flow. Setting the `model` to `context.model()` to avoid
@@ -356,10 +356,10 @@ private SymbolProvider createSymbolProvider(List<I> integrations, ServiceShape s
         return SymbolProvider.cache(provider);
     }
 
-    private C createContext(ServiceShape serviceShape, SymbolProvider provider) {
+    private C createContext(ServiceShape serviceShape, SymbolProvider provider, List<I> integrations) {
         LOGGER.fine(() -> "Creating a codegen context for " + directedCodegen.getClass().getName());
         return directedCodegen.createContext(new CreateContextDirective<>(
-                model, settings, serviceShape, provider, fileManifest));
+                model, settings, serviceShape, provider, fileManifest, integrations));
     }
 
     private void registerInterceptors(C context, List<I> integrations) {
@@ -404,9 +404,9 @@ private static class ShapeGenerator<
 
         private final C context;
         private final ServiceShape serviceShape;
-        private final DirectedCodegen<C, S> directedCodegen;
+        private final DirectedCodegen<C, S, ?> directedCodegen;
 
-        ShapeGenerator(C context, ServiceShape serviceShape, DirectedCodegen<C, S> directedCodegen) {
+        ShapeGenerator(C context, ServiceShape serviceShape, DirectedCodegen<C, S, ?> directedCodegen) {
             this.context = context;
             this.serviceShape = serviceShape;
             this.directedCodegen = directedCodegen;
diff --git a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/CreateContextDirective.java b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/CreateContextDirective.java
index 8982bf9c1c8..ec7a07fbb29 100644
--- a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/CreateContextDirective.java
+++ b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/CreateContextDirective.java
@@ -15,9 +15,12 @@
 
 package software.amazon.smithy.codegen.core.directed;
 
+import java.util.Collections;
+import java.util.List;
 import java.util.Map;
 import software.amazon.smithy.build.FileManifest;
 import software.amazon.smithy.codegen.core.CodegenContext;
+import software.amazon.smithy.codegen.core.SmithyIntegration;
 import software.amazon.smithy.codegen.core.SymbolProvider;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.knowledge.ServiceIndex;
@@ -29,23 +32,28 @@
  * Directive used to create a {@link CodegenContext}.
  *
  * @param <S> Codegen settings type.
+ * @param <I> {@link SmithyIntegration} type.
+
  * @see DirectedCodegen#createContext
  */
-public final class CreateContextDirective<S> extends Directive<S> {
+public final class CreateContextDirective<S, I extends SmithyIntegration<S, ?, ?>> extends Directive<S> {
 
     private final SymbolProvider symbolProvider;
     private final FileManifest fileManifest;
+    private final List<I> integrations;
 
     CreateContextDirective(
             Model model,
             S settings,
             ServiceShape service,
             SymbolProvider symbolProvider,
-            FileManifest fileManifest
+            FileManifest fileManifest,
+            List<I> integrations
     ) {
         super(model, settings, service);
         this.symbolProvider = symbolProvider;
         this.fileManifest = fileManifest;
+        this.integrations = Collections.unmodifiableList(integrations);
     }
 
     /**
@@ -62,6 +70,14 @@ public FileManifest fileManifest() {
         return fileManifest;
     }
 
+
+    /**
+     * @return Returns the list of Integrations used during codegen.
+     */
+    public List<I> integrations() {
+        return integrations;
+    }
+
     /**
      * Get a map of supported protocols on the service shape in the form of shape ID to
      * the definition of the trait.
diff --git a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/DirectedCodegen.java b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/DirectedCodegen.java
index 8325f69f74f..ee33901b931 100644
--- a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/DirectedCodegen.java
+++ b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/DirectedCodegen.java
@@ -30,8 +30,9 @@
  *
  * @param <C> Smithy {@link CodegenContext} to use in directed methods.
  * @param <S> Settings object passed to directed methods as part of the context.
+ * @param <I> {@link SmithyIntegration} type to use in directed methods.
  */
-public interface DirectedCodegen<C extends CodegenContext<S, ?>, S> {
+public interface DirectedCodegen<C extends CodegenContext<S, ?>, S, I extends SmithyIntegration<S, ?, ?>> {
 
     /**
      * Create the {@link SymbolProvider} used to map shapes to code symbols.
@@ -47,7 +48,7 @@ public interface DirectedCodegen<C extends CodegenContext<S, ?>, S> {
      * @param directive Directive context data.
      * @return Returns the created context object used by the rest of the directed generation.
      */
-    C createContext(CreateContextDirective<S> directive);
+    C createContext(CreateContextDirective<S, I> directive);
 
     /**
      * Generates the code needed for a service shape.
diff --git a/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/CodegenDirectorTest.java b/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/CodegenDirectorTest.java
index d6799da0d9a..0c42136bedd 100644
--- a/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/CodegenDirectorTest.java
+++ b/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/CodegenDirectorTest.java
@@ -31,14 +31,13 @@
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.node.ExpectationNotMetException;
 import software.amazon.smithy.model.node.Node;
-import software.amazon.smithy.model.shapes.Shape;
 import software.amazon.smithy.model.shapes.ShapeId;
 
 public class CodegenDirectorTest {
 
     interface TestIntegration extends SmithyIntegration<TestSettings, TestWriter, TestContext> {}
 
-    private static final class TestDirected implements DirectedCodegen<TestContext, TestSettings> {
+    private static final class TestDirected implements DirectedCodegen<TestContext, TestSettings, TestIntegration> {
         public final List<ShapeId> generatedShapes = new ArrayList<>();
 
         @Override
@@ -50,7 +49,7 @@ public SymbolProvider createSymbolProvider(CreateSymbolProviderDirective<TestSet
         }
 
         @Override
-        public TestContext createContext(CreateContextDirective<TestSettings> directive) {
+        public TestContext createContext(CreateContextDirective<TestSettings, TestIntegration> directive) {
             WriterDelegator<TestWriter> delegator = new WriterDelegator<>(
                     directive.fileManifest(),
                     directive.symbolProvider(),
diff --git a/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/TestIntegration.java b/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/TestIntegration.java
new file mode 100644
index 00000000000..ba9406fd185
--- /dev/null
+++ b/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/TestIntegration.java
@@ -0,0 +1,2 @@
+package software.amazon.smithy.codegen.core.directed;public class TestIntegration {
+}

From 67f7ae9a7effa40896bc26686991e19673e35641 Mon Sep 17 00:00:00 2001
From: Jaykumar Gosar <gosar@amazon.com>
Date: Thu, 2 Jun 2022 16:16:12 -0700
Subject: [PATCH 2/2] Make SmithyIntegrations available from CodegenContext

---
 .../amazon/smithy/codegen/core/CodegenContext.java       | 9 ++++++++-
 .../amazon/smithy/codegen/core/SmithyIntegration.java    | 2 +-
 .../smithy/codegen/core/directed/CodegenDirector.java    | 4 ++--
 .../codegen/core/directed/ContextualDirective.java       | 2 +-
 .../smithy/codegen/core/directed/CustomizeDirective.java | 2 +-
 .../smithy/codegen/core/directed/DirectedCodegen.java    | 2 +-
 .../codegen/core/directed/GenerateEnumDirective.java     | 2 +-
 .../codegen/core/directed/GenerateErrorDirective.java    | 2 +-
 .../codegen/core/directed/GenerateResourceDirective.java | 2 +-
 .../codegen/core/directed/GenerateServiceDirective.java  | 2 +-
 .../core/directed/GenerateStructureDirective.java        | 2 +-
 .../codegen/core/directed/GenerateUnionDirective.java    | 2 +-
 .../smithy/codegen/core/directed/ShapeDirective.java     | 2 +-
 .../codegen/core/IntegrationTopologicalSortTest.java     | 2 +-
 .../smithy/codegen/core/SmithyIntegrationTest.java       | 5 ++++-
 .../codegen/core/directed/CodegenDirectorTest.java       | 3 ---
 .../amazon/smithy/codegen/core/directed/TestContext.java | 9 ++++++++-
 .../smithy/codegen/core/directed/TestIntegration.java    | 8 +++++++-
 18 files changed, 41 insertions(+), 21 deletions(-)

diff --git a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/CodegenContext.java b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/CodegenContext.java
index 2b778f54f90..fe361d2003b 100644
--- a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/CodegenContext.java
+++ b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/CodegenContext.java
@@ -15,6 +15,7 @@
 
 package software.amazon.smithy.codegen.core;
 
+import java.util.List;
 import software.amazon.smithy.build.FileManifest;
 import software.amazon.smithy.model.Model;
 
@@ -24,8 +25,9 @@
  *
  * @param <S> The settings object used to configure the generator.
  * @param <W> The type of {@link SymbolWriter} used by the generator.
+ * @param <I> The type of {@link SmithyIntegration}s used by the generator.
  */
-public interface CodegenContext<S, W extends SymbolWriter<W, ?>> {
+public interface CodegenContext<S, W extends SymbolWriter<W, ?>, I extends SmithyIntegration<S, W, ?>> {
     /**
      * @return Gets the model being code generated.
      */
@@ -57,4 +59,9 @@ public interface CodegenContext<S, W extends SymbolWriter<W, ?>> {
      * @return Returns the writer delegator used by the generator.
      */
     WriterDelegator<W> writerDelegator();
+
+    /**
+     * @return Gets the SmithyIntegrations used for code generation.
+     */
+    List<I> integrations();
 }
diff --git a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/SmithyIntegration.java b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/SmithyIntegration.java
index 73fb03e30bd..5ab5ddad2b5 100644
--- a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/SmithyIntegration.java
+++ b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/SmithyIntegration.java
@@ -42,7 +42,7 @@
  * @param <C> The CodegenContext value used by the generator.
  */
 @SmithyUnstableApi
-public interface SmithyIntegration<S, W extends SymbolWriter<W, ?>, C extends CodegenContext<S, W>> {
+public interface SmithyIntegration<S, W extends SymbolWriter<W, ?>, C extends CodegenContext<S, W, ?>> {
     /**
      * Gets the name of the integration.
      *
diff --git a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/CodegenDirector.java b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/CodegenDirector.java
index b880745dbbd..2f9f4bff5fb 100644
--- a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/CodegenDirector.java
+++ b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/CodegenDirector.java
@@ -60,7 +60,7 @@
 public final class CodegenDirector<
         W extends SymbolWriter<W, ? extends ImportContainer>,
         I extends SmithyIntegration<S, W, C>,
-        C extends CodegenContext<S, W>,
+        C extends CodegenContext<S, W, I>,
         S> {
 
     private static final Logger LOGGER = Logger.getLogger(DirectedCodegen.class.getName());
@@ -399,7 +399,7 @@ private void applyIntegrationCustomizations(C context, List<I> integrations) {
 
     private static class ShapeGenerator<
             W extends SymbolWriter<W, ? extends ImportContainer>,
-            C extends CodegenContext<S, W>,
+            C extends CodegenContext<S, W, ?>,
             S> extends ShapeVisitor.Default<Void> {
 
         private final C context;
diff --git a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/ContextualDirective.java b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/ContextualDirective.java
index fff1c8dcb92..84909d86655 100644
--- a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/ContextualDirective.java
+++ b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/ContextualDirective.java
@@ -26,7 +26,7 @@
  * @param <C> CodegenContext type.
  * @param <S> Codegen settings type.
  */
-public abstract class ContextualDirective<C extends CodegenContext<S, ?>, S> extends Directive<S> {
+public abstract class ContextualDirective<C extends CodegenContext<S, ?, ?>, S> extends Directive<S> {
 
     private final C context;
 
diff --git a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/CustomizeDirective.java b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/CustomizeDirective.java
index d807b4c3cc6..77b9f2bc7fd 100644
--- a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/CustomizeDirective.java
+++ b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/CustomizeDirective.java
@@ -26,7 +26,7 @@
  * @see DirectedCodegen#customizeBeforeIntegrations
  * @see DirectedCodegen#customizeAfterIntegrations
  */
-public final class CustomizeDirective<C extends CodegenContext<S, ?>, S> extends ContextualDirective<C, S> {
+public final class CustomizeDirective<C extends CodegenContext<S, ?, ?>, S> extends ContextualDirective<C, S> {
     CustomizeDirective(C context, ServiceShape service) {
         super(context, service);
     }
diff --git a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/DirectedCodegen.java b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/DirectedCodegen.java
index ee33901b931..4d0860ebda2 100644
--- a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/DirectedCodegen.java
+++ b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/DirectedCodegen.java
@@ -32,7 +32,7 @@
  * @param <S> Settings object passed to directed methods as part of the context.
  * @param <I> {@link SmithyIntegration} type to use in directed methods.
  */
-public interface DirectedCodegen<C extends CodegenContext<S, ?>, S, I extends SmithyIntegration<S, ?, ?>> {
+public interface DirectedCodegen<C extends CodegenContext<S, ?, I>, S, I extends SmithyIntegration<S, ?, C>> {
 
     /**
      * Create the {@link SymbolProvider} used to map shapes to code symbols.
diff --git a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateEnumDirective.java b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateEnumDirective.java
index e0f49054050..79797251466 100644
--- a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateEnumDirective.java
+++ b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateEnumDirective.java
@@ -27,7 +27,7 @@
  * @param <S> Codegen settings type.
  * @see DirectedCodegen#generateEnumShape
  */
-public final class GenerateEnumDirective<C extends CodegenContext<S, ?>, S> extends ShapeDirective<Shape, C, S> {
+public final class GenerateEnumDirective<C extends CodegenContext<S, ?, ?>, S> extends ShapeDirective<Shape, C, S> {
 
     GenerateEnumDirective(C context, ServiceShape service, Shape shape) {
         super(context, service, validateShape(shape));
diff --git a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateErrorDirective.java b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateErrorDirective.java
index e0f1d53b5ca..ca85c5dc845 100644
--- a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateErrorDirective.java
+++ b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateErrorDirective.java
@@ -27,7 +27,7 @@
  * @param <S> Codegen settings type.
  * @see DirectedCodegen#generateError
  */
-public final class GenerateErrorDirective<C extends CodegenContext<S, ?>, S>
+public final class GenerateErrorDirective<C extends CodegenContext<S, ?, ?>, S>
         extends ShapeDirective<StructureShape, C, S> {
 
     GenerateErrorDirective(C context, ServiceShape service, StructureShape shape) {
diff --git a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateResourceDirective.java b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateResourceDirective.java
index 9b2904651e9..8b9301fc7a8 100644
--- a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateResourceDirective.java
+++ b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateResourceDirective.java
@@ -26,7 +26,7 @@
  * @param <S> Codegen settings type.
  * @see DirectedCodegen#generateResource
  */
-public final class GenerateResourceDirective<C extends CodegenContext<S, ?>, S>
+public final class GenerateResourceDirective<C extends CodegenContext<S, ?, ?>, S>
         extends ShapeDirective<ResourceShape, C, S> {
     GenerateResourceDirective(C context, ServiceShape service, ResourceShape shape) {
         super(context, service, shape);
diff --git a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateServiceDirective.java b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateServiceDirective.java
index e4a4bfd0638..ee5979252e0 100644
--- a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateServiceDirective.java
+++ b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateServiceDirective.java
@@ -35,7 +35,7 @@
  * @param <S> Codegen settings type.
  * @see DirectedCodegen#generateService
  */
-public final class GenerateServiceDirective<C extends CodegenContext<S, ?>, S>
+public final class GenerateServiceDirective<C extends CodegenContext<S, ?, ?>, S>
         extends ShapeDirective<ServiceShape, C, S> {
 
     GenerateServiceDirective(C context, ServiceShape service) {
diff --git a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateStructureDirective.java b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateStructureDirective.java
index 6d12817d571..b4ad1736401 100644
--- a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateStructureDirective.java
+++ b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateStructureDirective.java
@@ -29,7 +29,7 @@
  * @param <S> Codegen settings type.
  * @see DirectedCodegen#generateStructure
  */
-public final class GenerateStructureDirective<C extends CodegenContext<S, ?>, S>
+public final class GenerateStructureDirective<C extends CodegenContext<S, ?, ?>, S>
         extends ShapeDirective<StructureShape, C, S> {
 
     GenerateStructureDirective(C context, ServiceShape service, StructureShape shape) {
diff --git a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateUnionDirective.java b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateUnionDirective.java
index 08abee8327b..daee34a0865 100644
--- a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateUnionDirective.java
+++ b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateUnionDirective.java
@@ -27,7 +27,7 @@
  * @param <S> Codegen settings type.
  * @see DirectedCodegen#generateUnion
  */
-public final class GenerateUnionDirective<C extends CodegenContext<S, ?>, S>
+public final class GenerateUnionDirective<C extends CodegenContext<S, ?, ?>, S>
         extends ShapeDirective<UnionShape, C, S> {
 
     GenerateUnionDirective(C context, ServiceShape service, UnionShape shape) {
diff --git a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/ShapeDirective.java b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/ShapeDirective.java
index fe0709a6e3e..0c8c8d77ce9 100644
--- a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/ShapeDirective.java
+++ b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/ShapeDirective.java
@@ -27,7 +27,7 @@
  * @param <C> CodegenContext type.
  * @param <S> Codegen settings type.
  */
-public abstract class ShapeDirective<T extends Shape, C extends CodegenContext<S, ?>, S>
+public abstract class ShapeDirective<T extends Shape, C extends CodegenContext<S, ?, ?>, S>
         extends ContextualDirective<C, S> {
 
     private final T shape;
diff --git a/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/IntegrationTopologicalSortTest.java b/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/IntegrationTopologicalSortTest.java
index 53eaf61c801..7b76820d4de 100644
--- a/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/IntegrationTopologicalSortTest.java
+++ b/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/IntegrationTopologicalSortTest.java
@@ -34,7 +34,7 @@ private static final class MySettings {}
     private static final class MyIntegration implements SmithyIntegration<
             MySettings,
             MySimpleWriter,
-            CodegenContext<MySettings, MySimpleWriter>
+            CodegenContext<MySettings, MySimpleWriter, MyIntegration>
     > {
         private final String name;
         private final byte priority;
diff --git a/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/SmithyIntegrationTest.java b/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/SmithyIntegrationTest.java
index 8bccfbce57c..f3f80d8e2fa 100644
--- a/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/SmithyIntegrationTest.java
+++ b/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/SmithyIntegrationTest.java
@@ -29,7 +29,7 @@ public class SmithyIntegrationTest {
 
     private static final class MySettings {}
 
-    private static final class MyContext implements CodegenContext<MySettings, MySimpleWriter> {
+    private static final class MyContext implements CodegenContext<MySettings, MySimpleWriter, MyIntegration> {
         @Override
         public Model model() {
             return null;
@@ -54,6 +54,9 @@ public FileManifest fileManifest() {
         public WriterDelegator<MySimpleWriter> writerDelegator() {
             return null;
         }
+
+        @Override
+        public List<MyIntegration> integrations() { return null; }
     }
 
     private static final class MyIntegration implements SmithyIntegration<MySettings, MySimpleWriter, MyContext> {
diff --git a/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/CodegenDirectorTest.java b/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/CodegenDirectorTest.java
index 0c42136bedd..0a4eaf1570b 100644
--- a/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/CodegenDirectorTest.java
+++ b/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/CodegenDirectorTest.java
@@ -24,7 +24,6 @@
 import org.junit.jupiter.api.Test;
 import software.amazon.smithy.build.FileManifest;
 import software.amazon.smithy.build.MockManifest;
-import software.amazon.smithy.codegen.core.SmithyIntegration;
 import software.amazon.smithy.codegen.core.Symbol;
 import software.amazon.smithy.codegen.core.SymbolProvider;
 import software.amazon.smithy.codegen.core.WriterDelegator;
@@ -35,8 +34,6 @@
 
 public class CodegenDirectorTest {
 
-    interface TestIntegration extends SmithyIntegration<TestSettings, TestWriter, TestContext> {}
-
     private static final class TestDirected implements DirectedCodegen<TestContext, TestSettings, TestIntegration> {
         public final List<ShapeId> generatedShapes = new ArrayList<>();
 
diff --git a/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/TestContext.java b/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/TestContext.java
index 9f07245641c..4307aa2d5ef 100644
--- a/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/TestContext.java
+++ b/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/TestContext.java
@@ -15,6 +15,8 @@
 
 package software.amazon.smithy.codegen.core.directed;
 
+import java.util.ArrayList;
+import java.util.List;
 import software.amazon.smithy.build.FileManifest;
 import software.amazon.smithy.build.MockManifest;
 import software.amazon.smithy.codegen.core.CodegenContext;
@@ -25,7 +27,7 @@
 import software.amazon.smithy.model.shapes.ServiceShape;
 import software.amazon.smithy.model.shapes.ShapeId;
 
-final class TestContext implements CodegenContext<TestSettings, TestWriter> {
+final class TestContext implements CodegenContext<TestSettings, TestWriter, TestIntegration> {
 
     private final Model model;
     private final TestSettings settings;
@@ -86,6 +88,11 @@ public WriterDelegator<TestWriter> writerDelegator() {
         return delegator;
     }
 
+    @Override
+    public List<TestIntegration> integrations() {
+        return new ArrayList<>();
+    }
+
     public ServiceShape service() {
         return service;
     }
diff --git a/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/TestIntegration.java b/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/TestIntegration.java
index ba9406fd185..de3cd6231ec 100644
--- a/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/TestIntegration.java
+++ b/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/TestIntegration.java
@@ -1,2 +1,8 @@
-package software.amazon.smithy.codegen.core.directed;public class TestIntegration {
+package software.amazon.smithy.codegen.core.directed;
+
+import software.amazon.smithy.codegen.core.SmithyIntegration;
+
+interface TestIntegration extends SmithyIntegration<TestSettings, TestWriter, TestContext> {
+
 }
+