Correctly load client or server specific decorators from classpath (#1592)

The current approach that attempts to downcast never worked; all Rust
decorators were being loaded, and the cast was doing nothing, because
the generic type parameter is erased at runtime.

Attempting to downcast a generic class `C<T>` to `C<U>` where `U: T` is
not possible to do in Kotlin (and presumably all JVM-based languages)
_at runtime_. Not even when using reified type parameters of inline
functions. See https://kotlinlang.org/docs/generics.html#type-erasure
for details.

This commit thus goes for another approach, suggested in the linked
Stack Overflow question [0]: add a method to the loaded classes that
signals at runtime the generic type parameter (`ClientCodegenContext` or
`ServerCodegenContext`) they can work with, in order to filter them.

This commit also simplifies the way the Python server project loads the Python
server-specific decorators, by deleting the combined decorator
`PythonServerCodegenDecorator`, which was being loaded from the classpath, and
instead directly using `CombinedCodegenDecorator` and passing it the Python
server-specific decorators in the `extras` parameter.

[0]: https://stackoverflow.com/questions/5451734/loading-generic-service-implementations-via-java-util-serviceloader

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsEndpointDecorator.kt

@@ -73,6 +73,9 @@ class AwsEndpointDecorator : RustCodegenDecorator<ClientCodegenContext> {
     ): List<LibRsCustomization> {
         return baseCustomizations + PubUseEndpoint(codegenContext.runtimeConfig)
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }
 
 class EndpointConfigCustomization(private val coreCodegenContext: CoreCodegenContext, private val endpointData: ObjectNode) :

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsFluentClientDecorator.kt

@@ -111,6 +111,9 @@ class AwsFluentClientDecorator : RustCodegenDecorator<ClientCodegenContext> {
             }
         }
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }
 
 private class AwsFluentClientExtensions(types: Types) {

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsPresigningDecorator.kt

@@ -117,6 +117,9 @@ class AwsPresigningDecorator internal constructor(
         return presignableTransforms.fold(intermediate) { m, t -> t.transform(m) }
     }
 
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
+
     private fun addSyntheticOperations(model: Model): Model {
         val presignableOps = model.shapes()
             .filter { shape -> shape is OperationShape && presignableOperations.containsKey(shape.id) }

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsReadmeDecorator.kt

@@ -11,6 +11,7 @@ import org.jsoup.nodes.TextNode
 import software.amazon.smithy.model.traits.DocumentationTrait
 import software.amazon.smithy.rust.codegen.rustlang.raw
 import software.amazon.smithy.rust.codegen.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.smithy.CoreCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.RustCrate
 import software.amazon.smithy.rust.codegen.smithy.customize.RustCodegenDecorator
 import software.amazon.smithy.rust.codegen.smithy.generators.ManifestCustomizations
@@ -110,6 +111,9 @@ class AwsReadmeDecorator : RustCodegenDecorator<ClientCodegenContext> {
         }
     }
 
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
+
     /**
      * Strips HTML from the description and makes it human-readable Markdown.
      */

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/CrateLicenseDecorator.kt

@@ -7,6 +7,7 @@ package software.amazon.smithy.rustsdk
 
 import software.amazon.smithy.rust.codegen.rustlang.raw
 import software.amazon.smithy.rust.codegen.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.smithy.CoreCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.RustCrate
 import software.amazon.smithy.rust.codegen.smithy.customize.RustCodegenDecorator
 
@@ -21,4 +22,7 @@ class CrateLicenseDecorator : RustCodegenDecorator<ClientCodegenContext> {
             it.raw(license)
         }
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/CredentialProviders.kt

@@ -12,6 +12,7 @@ import software.amazon.smithy.rust.codegen.rustlang.rust
 import software.amazon.smithy.rust.codegen.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.rustlang.writable
 import software.amazon.smithy.rust.codegen.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.smithy.CoreCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.RuntimeConfig
 import software.amazon.smithy.rust.codegen.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.smithy.customize.OperationCustomization
@@ -47,6 +48,9 @@ class CredentialsProviderDecorator : RustCodegenDecorator<ClientCodegenContext>
     ): List<LibRsCustomization> {
         return baseCustomizations + PubUseCredentials(codegenContext.runtimeConfig)
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }
 
 /**

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/HttpRequestChecksumDecorator.kt

@@ -14,6 +14,7 @@ import software.amazon.smithy.rust.codegen.rustlang.asType
 import software.amazon.smithy.rust.codegen.rustlang.rust
 import software.amazon.smithy.rust.codegen.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.smithy.CoreCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.RuntimeConfig
 import software.amazon.smithy.rust.codegen.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.smithy.customize.OperationCustomization
@@ -51,6 +52,9 @@ class HttpRequestChecksumDecorator : RustCodegenDecorator<ClientCodegenContext>
     ): List<OperationCustomization> {
         return baseCustomizations + HttpRequestChecksumCustomization(codegenContext, operation)
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }
 
 private fun HttpChecksumTrait.requestAlgorithmMember(

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/HttpResponseChecksumDecorator.kt

@@ -11,6 +11,7 @@ import software.amazon.smithy.model.shapes.OperationShape
 import software.amazon.smithy.rust.codegen.rustlang.Writable
 import software.amazon.smithy.rust.codegen.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.smithy.CoreCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.customize.OperationCustomization
 import software.amazon.smithy.rust.codegen.smithy.customize.OperationSection
 import software.amazon.smithy.rust.codegen.smithy.customize.RustCodegenDecorator
@@ -38,6 +39,9 @@ class HttpResponseChecksumDecorator : RustCodegenDecorator<ClientCodegenContext>
     ): List<OperationCustomization> {
         return baseCustomizations + HttpResponseChecksumCustomization(codegenContext, operation)
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }
 
 // This generator was implemented based on this spec:

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/IntegrationTestDependencies.kt

@@ -13,6 +13,7 @@ import software.amazon.smithy.rust.codegen.rustlang.DependencyScope
 import software.amazon.smithy.rust.codegen.rustlang.Writable
 import software.amazon.smithy.rust.codegen.rustlang.writable
 import software.amazon.smithy.rust.codegen.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.smithy.CoreCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.RuntimeConfig
 import software.amazon.smithy.rust.codegen.smithy.customize.RustCodegenDecorator
 import software.amazon.smithy.rust.codegen.smithy.generators.LibRsCustomization
@@ -49,6 +50,9 @@ class IntegrationTestDecorator : RustCodegenDecorator<ClientCodegenContext> {
             baseCustomizations
         }
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }
 
 class IntegrationTestDependencies(

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/RegionDecorator.kt

@@ -95,6 +95,9 @@ class RegionDecorator : RustCodegenDecorator<ClientCodegenContext> {
     ): List<LibRsCustomization> {
         return baseCustomizations + PubUseRegion(codegenContext.runtimeConfig)
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }
 
 class RegionProviderConfig(coreCodegenContext: CoreCodegenContext) : ConfigCustomization() {

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/RetryPolicyDecorator.kt

@@ -10,6 +10,7 @@ import software.amazon.smithy.rust.codegen.rustlang.asType
 import software.amazon.smithy.rust.codegen.rustlang.rust
 import software.amazon.smithy.rust.codegen.rustlang.writable
 import software.amazon.smithy.rust.codegen.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.smithy.CoreCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.RuntimeConfig
 import software.amazon.smithy.rust.codegen.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.smithy.customize.OperationCustomization
@@ -27,6 +28,9 @@ class RetryPolicyDecorator : RustCodegenDecorator<ClientCodegenContext> {
     ): List<OperationCustomization> {
         return baseCustomizations + RetryPolicyFeature(codegenContext.runtimeConfig)
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }
 
 class RetryPolicyFeature(private val runtimeConfig: RuntimeConfig) : OperationCustomization() {

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/SdkConfigDecorator.kt

@@ -11,6 +11,7 @@ import software.amazon.smithy.rust.codegen.rustlang.asType
 import software.amazon.smithy.rust.codegen.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.rustlang.writable
 import software.amazon.smithy.rust.codegen.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.smithy.CoreCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.RuntimeConfig
 import software.amazon.smithy.rust.codegen.smithy.RustCrate
 import software.amazon.smithy.rust.codegen.smithy.customize.RustCodegenDecorator
@@ -66,6 +67,9 @@ class SdkConfigDecorator : RustCodegenDecorator<ClientCodegenContext> {
             )
         }
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }
 
 class NewFromShared(runtimeConfig: RuntimeConfig) : ConfigCustomization() {

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/ServiceConfigDecorator.kt

@@ -9,6 +9,7 @@ import software.amazon.smithy.rust.codegen.rustlang.Writable
 import software.amazon.smithy.rust.codegen.rustlang.docs
 import software.amazon.smithy.rust.codegen.rustlang.writable
 import software.amazon.smithy.rust.codegen.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.smithy.CoreCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.customize.RustCodegenDecorator
 import software.amazon.smithy.rust.codegen.smithy.generators.config.ConfigCustomization
 import software.amazon.smithy.rust.codegen.smithy.generators.config.ServiceConfig
@@ -21,6 +22,9 @@ class ServiceConfigDecorator : RustCodegenDecorator<ClientCodegenContext> {
         codegenContext: ClientCodegenContext,
         baseCustomizations: List<ConfigCustomization>,
     ): List<ConfigCustomization> = baseCustomizations + SharedConfigDocsCustomization()
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }
 
 class SharedConfigDocsCustomization : ConfigCustomization() {

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/SigV4SigningDecorator.kt

@@ -75,6 +75,9 @@ class SigV4SigningDecorator : RustCodegenDecorator<ClientCodegenContext> {
             )
         }
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }
 
 class SigV4SigningConfig(

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/UserAgentDecorator.kt

@@ -13,6 +13,7 @@ import software.amazon.smithy.rust.codegen.rustlang.rust
 import software.amazon.smithy.rust.codegen.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.rustlang.writable
 import software.amazon.smithy.rust.codegen.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.smithy.CoreCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.RuntimeConfig
 import software.amazon.smithy.rust.codegen.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.smithy.customize.OperationCustomization
@@ -55,6 +56,9 @@ class UserAgentDecorator : RustCodegenDecorator<ClientCodegenContext> {
     ): List<OperationCustomization> {
         return baseCustomizations + UserAgentFeature(codegenContext.runtimeConfig)
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }
 
 /**

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/apigateway/ApiGatewayDecorator.kt

@@ -34,6 +34,9 @@ class ApiGatewayDecorator : RustCodegenDecorator<ClientCodegenContext> {
             it + ApiGatewayAddAcceptHeader()
         }
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }
 
 class ApiGatewayAddAcceptHeader : OperationCustomization() {

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/auth/DisabledAuthDecorator.kt

@@ -12,6 +12,7 @@ import software.amazon.smithy.model.shapes.ShapeId
 import software.amazon.smithy.model.traits.AuthTrait
 import software.amazon.smithy.model.transform.ModelTransformer
 import software.amazon.smithy.rust.codegen.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.smithy.CoreCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.customize.RustCodegenDecorator
 
 private fun String.shapeId() = ShapeId.from(this)
@@ -46,4 +47,7 @@ class DisabledAuthDecorator : RustCodegenDecorator<ClientCodegenContext> {
             }
         }
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/ec2/Ec2Decorator.kt

@@ -9,6 +9,7 @@ import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.ServiceShape
 import software.amazon.smithy.model.shapes.ShapeId
 import software.amazon.smithy.rust.codegen.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.smithy.CoreCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.customize.RustCodegenDecorator
 import software.amazon.smithy.rust.codegen.smithy.letIf
 
@@ -28,4 +29,7 @@ class Ec2Decorator : RustCodegenDecorator<ClientCodegenContext> {
             BoxPrimitiveShapes::processModel,
         )
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/glacier/GlacierDecorator.kt

@@ -37,4 +37,7 @@ class GlacierDecorator : RustCodegenDecorator<ClientCodegenContext> {
         }
         return baseCustomizations + extras
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/route53/Route53Decorator.kt

@@ -16,6 +16,7 @@ import software.amazon.smithy.rust.codegen.rustlang.Writable
 import software.amazon.smithy.rust.codegen.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.rustlang.writable
 import software.amazon.smithy.rust.codegen.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.smithy.CoreCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.smithy.customize.OperationCustomization
 import software.amazon.smithy.rust.codegen.smithy.customize.OperationSection
@@ -59,6 +60,9 @@ class Route53Decorator : RustCodegenDecorator<ClientCodegenContext> {
         } else baseCustomizations
     }
 
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
+
     private fun isResourceId(shape: Shape): Boolean {
         return (shape is MemberShape && resourceShapes.contains(shape.target)) && shape.hasTrait<HttpLabelTrait>()
     }

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/s3/S3Decorator.kt

@@ -56,6 +56,9 @@ class S3Decorator : RustCodegenDecorator<ClientCodegenContext> {
     ): List<LibRsCustomization> = baseCustomizations.letIf(applies(codegenContext.serviceShape.id)) {
         it + S3PubUse()
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }
 
 class S3(coreCodegenContext: CoreCodegenContext) : RestXml(coreCodegenContext) {

aws/sdk-codegen/src/test/kotlin/software/amazon/smithy/rustsdk/EndpointConfigCustomizationTest.kt

@@ -12,6 +12,7 @@ import software.amazon.smithy.rust.codegen.rustlang.asType
 import software.amazon.smithy.rust.codegen.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.smithy.ClientCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.CodegenVisitor
+import software.amazon.smithy.rust.codegen.smithy.CoreCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.RustCrate
 import software.amazon.smithy.rust.codegen.smithy.customizations.AllowLintsGenerator
 import software.amazon.smithy.rust.codegen.smithy.customize.CombinedCodegenDecorator
@@ -147,6 +148,9 @@ internal class EndpointConfigCustomizationTest {
                     test(rustCrate)
                 }
             }
+
+            override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+                clazz.isAssignableFrom(ClientCodegenContext::class.java)
         }
         val customization = CombinedCodegenDecorator(listOf(RequiredCustomizations(), codegenDecorator))
         CodegenVisitor(context, customization).execute()

codegen-server/python/src/main/kotlin/software/amazon/smithy/rust/codegen/server/python/smithy/PythonCodegenServerPlugin.kt

@@ -11,6 +11,7 @@ import software.amazon.smithy.codegen.core.ReservedWordSymbolProvider
 import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.ServiceShape
 import software.amazon.smithy.rust.codegen.rustlang.RustReservedWordSymbolProvider
+import software.amazon.smithy.rust.codegen.server.python.smithy.customizations.DECORATORS
 import software.amazon.smithy.rust.codegen.server.python.smithy.generators.PythonServerSymbolProvider
 import software.amazon.smithy.rust.codegen.server.smithy.customizations.ServerRequiredCustomizations
 import software.amazon.smithy.rust.codegen.smithy.BaseSymbolMetadataProvider
@@ -45,7 +46,10 @@ class PythonCodegenServerPlugin : SmithyBuildPlugin {
         // - context (e.g. the of the operation)
         // - writer: The active RustWriter at the given location
         val codegenDecorator: CombinedCodegenDecorator<ServerCodegenContext> =
-            CombinedCodegenDecorator.fromClasspath(context, ServerRequiredCustomizations())
+            CombinedCodegenDecorator.fromClasspath(
+                context,
+                CombinedCodegenDecorator(DECORATORS + ServerRequiredCustomizations()),
+            )
 
         // PythonServerCodegenVisitor is the main driver of code generation that traverses the model and generates code
         logger.info("Loaded plugin to generate Rust/Python bindings for the server SSDK for projection ${context.projectionName}")

codegen-server/python/src/main/kotlin/software/amazon/smithy/rust/codegen/server/python/smithy/customizations/PythonServerCodegenDecorator.kt

@@ -14,9 +14,9 @@ import software.amazon.smithy.rust.codegen.rustlang.writable
 import software.amazon.smithy.rust.codegen.server.python.smithy.PythonServerRuntimeType
 import software.amazon.smithy.rust.codegen.server.python.smithy.generators.PythonServerModuleGenerator
 import software.amazon.smithy.rust.codegen.server.smithy.customizations.AddInternalServerErrorToAllOperationsDecorator
+import software.amazon.smithy.rust.codegen.smithy.CoreCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.RustCrate
 import software.amazon.smithy.rust.codegen.smithy.ServerCodegenContext
-import software.amazon.smithy.rust.codegen.smithy.customize.CombinedCodegenDecorator
 import software.amazon.smithy.rust.codegen.smithy.customize.RustCodegenDecorator
 import software.amazon.smithy.rust.codegen.smithy.generators.LibRsCustomization
 import software.amazon.smithy.rust.codegen.smithy.generators.LibRsSection
@@ -44,6 +44,9 @@ class CdylibManifestDecorator : RustCodegenDecorator<ServerCodegenContext> {
                 "crate-type" to listOf("cdylib"),
             ),
         )
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ServerCodegenContext::class.java)
 }
 
 /**
@@ -76,6 +79,9 @@ class PythonExportModuleDecorator : RustCodegenDecorator<ServerCodegenContext> {
         val serviceShapes = Walker(codegenContext.model).walkShapes(service)
         PythonServerModuleGenerator(codegenContext, rustCrate, serviceShapes).render()
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ServerCodegenContext::class.java)
 }
 
 /**
@@ -91,6 +97,9 @@ class PubUsePythonTypesDecorator : RustCodegenDecorator<ServerCodegenContext> {
     ): List<LibRsCustomization> {
         return baseCustomizations + PubUsePythonTypes(codegenContext)
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ServerCodegenContext::class.java)
 }
 
 val DECORATORS = listOf(
@@ -106,9 +115,3 @@ val DECORATORS = listOf(
     // Render the Python shared library export.
     PythonExportModuleDecorator(),
 )
-
-// Combined codegen decorator for Python services.
-class PythonServerCodegenDecorator : CombinedCodegenDecorator<ServerCodegenContext>(DECORATORS) {
-    override val name: String = "PythonServerCodegenDecorator"
-    override val order: Byte = -1
-}

codegen-server/python/src/main/resources/META-INF/services/software.amazon.smithy.rust.codegen.smithy.customize.RustCodegenDecorator

@@ -1,6 +0,0 @@
-#
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-
-software.amazon.smithy.rust.codegen.server.python.smithy.customizations.PythonServerCodegenDecorator

codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/customizations/AdditionalErrorsDecorator.kt

@@ -13,6 +13,7 @@ import software.amazon.smithy.model.shapes.StructureShape
 import software.amazon.smithy.model.traits.ErrorTrait
 import software.amazon.smithy.model.traits.RequiredTrait
 import software.amazon.smithy.model.transform.ModelTransformer
+import software.amazon.smithy.rust.codegen.smithy.CoreCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.ServerCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.customize.RustCodegenDecorator
 import software.amazon.smithy.rust.codegen.smithy.transformers.allErrors
@@ -37,6 +38,9 @@ class AddInternalServerErrorToInfallibleOperationsDecorator : RustCodegenDecorat
 
     override fun transformModel(service: ServiceShape, model: Model): Model =
         addErrorShapeToModelOperations(service, model) { shape -> shape.allErrors(model).isEmpty() }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ServerCodegenContext::class.java)
 }
 
 /**
@@ -63,6 +67,9 @@ class AddInternalServerErrorToAllOperationsDecorator : RustCodegenDecorator<Serv
 
     override fun transformModel(service: ServiceShape, model: Model): Model =
         addErrorShapeToModelOperations(service, model) { true }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ServerCodegenContext::class.java)
 }
 
 fun addErrorShapeToModelOperations(service: ServiceShape, model: Model, opSelector: (OperationShape) -> Boolean): Model {

codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/customizations/ServerRequiredCustomizations.kt

@@ -6,6 +6,7 @@
 package software.amazon.smithy.rust.codegen.server.smithy.customizations
 
 import software.amazon.smithy.rust.codegen.rustlang.Feature
+import software.amazon.smithy.rust.codegen.smithy.CoreCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.RustCrate
 import software.amazon.smithy.rust.codegen.smithy.ServerCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.customizations.AllowLintsGenerator
@@ -35,4 +36,7 @@ class ServerRequiredCustomizations : RustCodegenDecorator<ServerCodegenContext>
         // Add rt-tokio feature for `ByteStream::from_path`
         rustCrate.mergeFeature(Feature("rt-tokio", true, listOf("aws-smithy-http/rt-tokio")))
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ServerCodegenContext::class.java)
 }

codegen-server/src/test/kotlin/software/amazon/smithy/rust/codegen/smithy/customize/CombinedCodegenDecoratorTest.kt

@@ -0,0 +1,35 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.rust.codegen.smithy.customize
+
+import io.kotest.matchers.collections.shouldContainExactly
+import org.junit.jupiter.api.Test
+import software.amazon.smithy.rust.codegen.server.smithy.customizations.ServerRequiredCustomizations
+import software.amazon.smithy.rust.codegen.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.smithy.ServerCodegenContext
+
+internal class CombinedCodegenDecoratorTest {
+    private val clientDecorator: RustCodegenDecorator<ClientCodegenContext> = RequiredCustomizations()
+    private val serverDecorator: RustCodegenDecorator<ServerCodegenContext> = ServerRequiredCustomizations()
+
+    @Test
+    fun filterClientDecorators() {
+        val filteredDecorators = CombinedCodegenDecorator.filterDecorators<ClientCodegenContext>(
+            listOf(clientDecorator, serverDecorator),
+        ).toList()
+
+        filteredDecorators.shouldContainExactly(clientDecorator)
+    }
+
+    @Test
+    fun filterServerDecorators() {
+        val filteredDecorators = CombinedCodegenDecorator.filterDecorators<ServerCodegenContext>(
+            listOf(clientDecorator, serverDecorator),
+        ).toList()
+
+        filteredDecorators.shouldContainExactly(serverDecorator)
+    }
+}

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/customizations/ClientCustomizations.kt

@@ -6,6 +6,7 @@
 package software.amazon.smithy.rust.codegen.smithy.customizations
 
 import software.amazon.smithy.rust.codegen.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.smithy.CoreCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.customize.RustCodegenDecorator
 import software.amazon.smithy.rust.codegen.smithy.generators.LibRsCustomization
 
@@ -20,4 +21,7 @@ class ClientCustomizations : RustCodegenDecorator<ClientCodegenContext> {
         codegenContext: ClientCodegenContext,
         baseCustomizations: List<LibRsCustomization>,
     ): List<LibRsCustomization> = baseCustomizations + ClientDocsGenerator()
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/customizations/DocsRsMetadataDecorator.kt

@@ -6,6 +6,7 @@
 package software.amazon.smithy.rust.codegen.smithy.customizations
 
 import software.amazon.smithy.rust.codegen.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.smithy.CoreCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.customize.RustCodegenDecorator
 import software.amazon.smithy.rust.codegen.smithy.generators.ManifestCustomizations
 
@@ -56,4 +57,7 @@ class DocsRsMetadataDecorator(private val docsRsMetadataSettings: DocsRsMetadata
     override fun crateManifestCustomizations(codegenContext: ClientCodegenContext): ManifestCustomizations {
         return docsRsMetadataSettings.asMap()
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/customizations/RetryConfigDecorator.kt

@@ -84,6 +84,9 @@ class RetryConfigDecorator : RustCodegenDecorator<ClientCodegenContext> {
     ): List<LibRsCustomization> {
         return baseCustomizations + PubUseRetryConfig(codegenContext.runtimeConfig)
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }
 
 class RetryConfigProviderConfig(coreCodegenContext: CoreCodegenContext) : ConfigCustomization() {

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/customizations/SleepImplDecorator.kt

@@ -126,6 +126,9 @@ class SleepImplDecorator : RustCodegenDecorator<ClientCodegenContext> {
     ): List<ConfigCustomization> {
         return baseCustomizations + SleepImplProviderConfig(codegenContext)
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }
 
 class SleepImplProviderConfig(coreCodegenContext: CoreCodegenContext) : ConfigCustomization() {

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/customizations/TimeoutConfigDecorator.kt

@@ -113,6 +113,9 @@ class TimeoutConfigDecorator : RustCodegenDecorator<ClientCodegenContext> {
     ): List<ConfigCustomization> {
         return baseCustomizations + TimeoutConfigProviderConfig(codegenContext)
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }
 
 class TimeoutConfigProviderConfig(coreCodegenContext: CoreCodegenContext) : ConfigCustomization() {

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/customize/NoOpEventStreamSigningDecorator.kt

@@ -40,6 +40,8 @@ open class NoOpEventStreamSigningDecorator<C : CoreCodegenContext> : RustCodegen
             codegenContext.runtimeConfig,
         )
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>) = true
 }
 
 class NoOpEventStreamSigningConfig(

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/customize/RequiredCustomizations.kt

@@ -8,6 +8,7 @@ package software.amazon.smithy.rust.codegen.smithy.customize
 import software.amazon.smithy.model.shapes.OperationShape
 import software.amazon.smithy.rust.codegen.rustlang.Feature
 import software.amazon.smithy.rust.codegen.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.smithy.CoreCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.RustCrate
 import software.amazon.smithy.rust.codegen.smithy.customizations.AllowLintsGenerator
 import software.amazon.smithy.rust.codegen.smithy.customizations.CrateVersionGenerator
@@ -48,4 +49,7 @@ class RequiredCustomizations : RustCodegenDecorator<ClientCodegenContext> {
         // Add rt-tokio feature for `ByteStream::from_path`
         rustCrate.mergeFeature(Feature("rt-tokio", true, listOf("aws-smithy-http/rt-tokio")))
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/customize/RustCodegenDecorator.kt

@@ -68,6 +68,8 @@ interface RustCodegenDecorator<C : CoreCodegenContext> {
         currentProtocols
 
     fun transformModel(service: ServiceShape, model: Model): Model = model
+
+    fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean
 }
 
 /**
@@ -75,7 +77,8 @@ interface RustCodegenDecorator<C : CoreCodegenContext> {
  *
  * This makes the actual concrete codegen simpler by not needing to deal with multiple separate decorators.
  */
-open class CombinedCodegenDecorator<C : CoreCodegenContext>(decorators: List<RustCodegenDecorator<C>>) : RustCodegenDecorator<C> {
+open class CombinedCodegenDecorator<C : CoreCodegenContext>(decorators: List<RustCodegenDecorator<C>>) :
+    RustCodegenDecorator<C> {
     private val orderedDecorators = decorators.sortedBy { it.order }
     override val name: String
         get() = "MetaDecorator"
@@ -137,6 +140,10 @@ open class CombinedCodegenDecorator<C : CoreCodegenContext>(decorators: List<Rus
         }
     }
 
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        // `CombinedCodegenDecorator` can work with all types of codegen context.
+        CoreCodegenContext::class.java.isAssignableFrom(clazz)
+
     companion object {
         inline fun <reified T : CoreCodegenContext> fromClasspath(
             context: PluginContext,
@@ -147,31 +154,46 @@ open class CombinedCodegenDecorator<C : CoreCodegenContext>(decorators: List<Rus
                 RustCodegenDecorator::class.java,
                 context.pluginClassLoader.orElse(RustCodegenDecorator::class.java.classLoader),
             )
+
+            val filteredDecorators = filterDecorators<T>(decorators, logger).toList()
+            return CombinedCodegenDecorator(filteredDecorators + extras)
+        }
+
+        /*
+         * This function has been extracted solely for the purposes of easily unit testing the important filtering logic.
+         * Unfortunately, it must be part of the public API because public API inline functions are not allowed to use
+         * non-public-API declarations.
+         * See https://kotlinlang.org/docs/inline-functions.html#restrictions-for-public-api-inline-functions.
+         */
+        inline fun <reified T : CoreCodegenContext> filterDecorators(
+            decorators: Iterable<RustCodegenDecorator<*>>,
+            logger: Logger = Logger.getLogger("RustCodegenSPILoader"),
+        ): Sequence<RustCodegenDecorator<T>> =
+            decorators.asSequence()
+                .onEach {
+                    logger.info("Discovered Codegen Decorator: ${it.javaClass.name}")
+                }
                 // The JVM's `ServiceLoader` is woefully underpowered in that it can not load classes with generic
                 // parameters with _fixed_ parameters (like what we're trying to do here; we only want `RustCodegenDecorator`
                 // classes with code-generation context matching the input `T`).
                 // There are various workarounds: https://stackoverflow.com/questions/5451734/loading-generic-service-implementations-via-java-util-serviceloader
                 // All involve loading _all_ classes from the classpath (i.e. all `RustCodegenDecorator<*>`), and then
-                // filtering them. The most elegant way to filter is arguably by checking if we can cast the loaded
-                // class to what we want.
+                // filtering them.
+                // Note that attempting to downcast a generic class `C<T>` to `C<U>` where `U: T` is not possible to do
+                // in Kotlin (and presumably all JVM-based languages) _at runtime_. Not even when using reified type
+                // parameters of inline functions. See https://kotlinlang.org/docs/generics.html#type-erasure for details.
                 .filter {
-                    try {
-                        it as RustCodegenDecorator<T>
-                        true
-                    } catch (e: ClassCastException) {
-                        false
-                    }
+                    val clazz = T::class.java
+                    it.supportsCodegenContext(clazz)
                 }
                 .onEach {
                     logger.info("Adding Codegen Decorator: ${it.javaClass.name}")
                 }
                 .map {
                     // Cast is safe because of the filter above.
+                    // Not that it really has an effect at runtime, since its unchecked.
                     @Suppress("UNCHECKED_CAST")
                     it as RustCodegenDecorator<T>
                 }
-                .toList()
-            return CombinedCodegenDecorator(decorators + extras)
-        }
     }
 }

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/generators/client/FluentClientDecorator.kt

@@ -99,6 +99,9 @@ class FluentClientDecorator : RustCodegenDecorator<ClientCodegenContext> {
             }
         }
     }
+
+    override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+        clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }
 
 sealed class FluentClientSection(name: String) : Section(name) {

codegen/src/test/kotlin/software/amazon/smithy/rust/codegen/customizations/HttpVersionListGeneratorTest.kt

@@ -13,6 +13,7 @@ import software.amazon.smithy.rust.codegen.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.rustlang.writable
 import software.amazon.smithy.rust.codegen.smithy.ClientCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.CodegenVisitor
+import software.amazon.smithy.rust.codegen.smithy.CoreCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.RuntimeConfig
 import software.amazon.smithy.rust.codegen.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.smithy.RustCrate
@@ -84,6 +85,9 @@ internal class HttpVersionListGeneratorTest {
                     )
                 }
             }
+
+            override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+                clazz.isAssignableFrom(ClientCodegenContext::class.java)
         }
         val combinedCodegenDecorator: CombinedCodegenDecorator<ClientCodegenContext> =
             CombinedCodegenDecorator.fromClasspath(ctx, RequiredCustomizations()).withDecorator(testWriter)
@@ -146,6 +150,9 @@ internal class HttpVersionListGeneratorTest {
                     )
                 }
             }
+
+            override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+                clazz.isAssignableFrom(ClientCodegenContext::class.java)
         }
 
         val combinedCodegenDecorator: CombinedCodegenDecorator<ClientCodegenContext> =
@@ -228,6 +235,9 @@ internal class HttpVersionListGeneratorTest {
                     )
                 }
             }
+
+            override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+                clazz.isAssignableFrom(ClientCodegenContext::class.java)
         }
 
         val combinedCodegenDecorator: CombinedCodegenDecorator<ClientCodegenContext> =

codegen/src/test/kotlin/software/amazon/smithy/rust/codegen/smithy/generators/EndpointTraitBindingsTest.kt

@@ -15,6 +15,7 @@ import software.amazon.smithy.rust.codegen.rustlang.rust
 import software.amazon.smithy.rust.codegen.rustlang.rustBlock
 import software.amazon.smithy.rust.codegen.smithy.ClientCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.CodegenVisitor
+import software.amazon.smithy.rust.codegen.smithy.CoreCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.RustCrate
 import software.amazon.smithy.rust.codegen.smithy.customize.CombinedCodegenDecorator
 import software.amazon.smithy.rust.codegen.smithy.customize.RequiredCustomizations
@@ -171,6 +172,9 @@ internal class EndpointTraitBindingsTest {
                     )
                 }
             }
+
+            override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+                clazz.isAssignableFrom(ClientCodegenContext::class.java)
         }
         val combinedCodegenDecorator: CombinedCodegenDecorator<ClientCodegenContext> =
             CombinedCodegenDecorator.fromClasspath(ctx, RequiredCustomizations()).withDecorator(codegenDecorator)

codegen/src/test/kotlin/software/amazon/smithy/rust/codegen/smithy/generators/protocol/ProtocolTestGeneratorTest.kt

@@ -231,6 +231,9 @@ class ProtocolTestGeneratorTest {
                 ): ProtocolMap<ClientCodegenContext> =
                     // Intentionally replace the builtin implementation of RestJson1 with our fake protocol
                     mapOf(RestJson1Trait.ID to TestProtocolFactory(httpRequestBuilder, body, correctResponse))
+
+                override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
+                    clazz.isAssignableFrom(ClientCodegenContext::class.java)
             },
         )
         visitor.execute()