Support server event streams

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

@@ -155,7 +155,7 @@ class AwsInputPresignedMethod(
     }
 
     private fun RustWriter.writeInputPresignedMethod(section: OperationSection.InputImpl) {
-        val operationError = operationShape.errorSymbol(symbolProvider)
+        val operationError = operationShape.errorSymbol(coreCodegenContext.model, symbolProvider, coreCodegenContext.target)
         val presignableOp = PRESIGNABLE_OPERATIONS.getValue(operationShape.id)
 
         val makeOperationOp = if (presignableOp.hasModelTransforms()) {

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

@@ -15,20 +15,21 @@ 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.server.python.smithy.PythonServerCargoDependency
-import software.amazon.smithy.rust.codegen.server.smithy.generators.ServerCombinedErrorGenerator
 import software.amazon.smithy.rust.codegen.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.smithy.RustSymbolProvider
+import software.amazon.smithy.rust.codegen.smithy.generators.CodegenTarget
+import software.amazon.smithy.rust.codegen.smithy.generators.error.ServerCombinedErrorGenerator
 import software.amazon.smithy.rust.codegen.smithy.generators.error.errorSymbol
 
 /**
  * Generates a unified error enum for [operation]. It depends on [ServerCombinedErrorGenerator]
  * to generate the errors from the model and adds the Rust implementation `From<pyo3::PyErr>`.
  */
 class PythonServerCombinedErrorGenerator(
-    model: Model,
+    private val model: Model,
     private val symbolProvider: RustSymbolProvider,
     private val operation: OperationShape
-) : ServerCombinedErrorGenerator(model, symbolProvider, operation) {
+) : ServerCombinedErrorGenerator(model, symbolProvider, symbolProvider.toSymbol(operation), listOf()) {
 
     private val operationIndex = OperationIndex.of(model)
     private val errors = operationIndex.getErrors(operation)
@@ -53,7 +54,7 @@ class PythonServerCombinedErrorGenerator(
 
             """,
             "pyo3" to PythonServerCargoDependency.PyO3.asType(),
-            "Error" to operation.errorSymbol(symbolProvider),
+            "Error" to operation.errorSymbol(model, symbolProvider, CodegenTarget.SERVER),
             "From" to RuntimeType.From,
             "CastPyErrToRustError" to castPyErrToRustError()
         )

codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerOperationHandlerGenerator.kt

@@ -16,6 +16,7 @@ import software.amazon.smithy.rust.codegen.server.smithy.ServerRuntimeType
 import software.amazon.smithy.rust.codegen.server.smithy.protocols.ServerHttpBoundProtocolGenerator
 import software.amazon.smithy.rust.codegen.smithy.CoreCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.RuntimeType
+import software.amazon.smithy.rust.codegen.smithy.generators.CodegenTarget
 import software.amazon.smithy.rust.codegen.smithy.generators.error.errorSymbol
 import software.amazon.smithy.rust.codegen.smithy.transformers.operationErrors
 import software.amazon.smithy.rust.codegen.util.hasStreamingMember
@@ -139,7 +140,7 @@ open class ServerOperationHandlerGenerator(
             "Fun: FnOnce($inputName) -> Fut + Clone + Send + 'static,"
         }
         val outputType = if (operation.operationErrors(model).isNotEmpty()) {
-            "Result<${symbolProvider.toSymbol(operation.outputShape(model)).fullName}, ${operation.errorSymbol(symbolProvider).fullyQualifiedName()}>"
+            "Result<${symbolProvider.toSymbol(operation.outputShape(model)).fullName}, ${operation.errorSymbol(model, symbolProvider, CodegenTarget.SERVER).fullyQualifiedName()}>"
         } else {
             symbolProvider.toSymbol(operation.outputShape(model)).fullName
         }

codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerOperationRegistryGenerator.kt

@@ -27,6 +27,7 @@ import software.amazon.smithy.rust.codegen.smithy.Errors
 import software.amazon.smithy.rust.codegen.smithy.Inputs
 import software.amazon.smithy.rust.codegen.smithy.Outputs
 import software.amazon.smithy.rust.codegen.smithy.RuntimeType
+import software.amazon.smithy.rust.codegen.smithy.generators.CodegenTarget
 import software.amazon.smithy.rust.codegen.smithy.generators.error.errorSymbol
 import software.amazon.smithy.rust.codegen.smithy.protocols.Protocol
 import software.amazon.smithy.rust.codegen.util.getTrait
@@ -91,11 +92,11 @@ class ServerOperationRegistryGenerator(
         }
 
         writer.rustTemplate(
-            """
+"""
 ##[allow(clippy::tabs_in_doc_comments)]
 /// The `$operationRegistryName` is the place where you can register
 /// your service's operation implementations.
-/// 
+///
 /// Use [`$operationRegistryBuilderName`] to construct the
 /// `$operationRegistryName`. For each of the [operations] modeled in
 /// your Smithy service, you need to provide an implementation in the
@@ -116,9 +117,9 @@ class ServerOperationRegistryGenerator(
 /// type implementing [`tower::make::MakeService`], a _service
 /// factory_. You can feed this value to a [Hyper server], and the
 /// server will instantiate and [`serve`] your service.
-/// 
+///
 /// Here's a full example to get you started:
-/// 
+///
 /// ```rust
 /// use std::net::SocketAddr;
 $inputOutputErrorsImport
@@ -346,9 +347,9 @@ ${operationImplementationStubs(operations)}
             } else ""
             ret +
                 """
-                    /// ${it.signature()} {
-                    ///     todo!()
-                    /// }
+                /// ${it.signature()} {
+                ///     todo!()
+                /// }
                 """.trimIndent()
         }
 
@@ -358,7 +359,7 @@ ${operationImplementationStubs(operations)}
     private fun OperationShape.signature(): String {
         val inputSymbol = symbolProvider.toSymbol(inputShape(model))
         val outputSymbol = symbolProvider.toSymbol(outputShape(model))
-        val errorSymbol = errorSymbol(symbolProvider)
+        val errorSymbol = errorSymbol(model, symbolProvider, CodegenTarget.SERVER)
 
         val inputT = "${Inputs.namespace}::${inputSymbol.name}"
         val t = "${Outputs.namespace}::${outputSymbol.name}"

codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerServiceGenerator.kt

@@ -74,7 +74,7 @@ open class ServerServiceGenerator(
 
     // Render combined errors.
     open fun renderCombinedErrors(writer: RustWriter, operation: OperationShape) {
-        ServerCombinedErrorGenerator(coreCodegenContext.model, coreCodegenContext.symbolProvider, operation).render(writer)
+        /* Subclasses can override */
     }
 
     // Render operations handler.

codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/protocols/ServerHttpBoundProtocolGenerator.kt

@@ -44,6 +44,7 @@ import software.amazon.smithy.rust.codegen.server.smithy.generators.http.ServerR
 import software.amazon.smithy.rust.codegen.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.smithy.ServerCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.extractSymbolFromOption
+import software.amazon.smithy.rust.codegen.smithy.generators.CodegenTarget
 import software.amazon.smithy.rust.codegen.smithy.generators.StructureGenerator
 import software.amazon.smithy.rust.codegen.smithy.generators.builderSymbol
 import software.amazon.smithy.rust.codegen.smithy.generators.error.errorSymbol
@@ -215,7 +216,7 @@ private class ServerHttpBoundProtocolTraitImplGenerator(
         // Implement `into_response` for output types.
 
         val outputName = "${operationName}${ServerHttpBoundProtocolGenerator.OPERATION_OUTPUT_WRAPPER_SUFFIX}"
-        val errorSymbol = operationShape.errorSymbol(symbolProvider)
+        val errorSymbol = operationShape.errorSymbol(model, symbolProvider, CodegenTarget.SERVER)
 
         if (operationShape.operationErrors(model).isNotEmpty()) {
             // The output of fallible operations is a `Result` which we convert into an
@@ -417,7 +418,7 @@ private class ServerHttpBoundProtocolTraitImplGenerator(
 
     private fun serverSerializeError(operationShape: OperationShape): RuntimeType {
         val fnName = "serialize_${operationShape.id.name.toSnakeCase()}_error"
-        val errorSymbol = operationShape.errorSymbol(symbolProvider)
+        val errorSymbol = operationShape.errorSymbol(model, symbolProvider, CodegenTarget.SERVER)
         return RuntimeType.forInlineFun(fnName, operationSerModule) {
             Attribute.Custom("allow(clippy::unnecessary_wraps)").render(it)
             it.rustBlockTemplate(

codegen-server/src/test/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerCombinedErrorGeneratorTest.kt

@@ -10,6 +10,7 @@ import software.amazon.smithy.model.shapes.StructureShape
 import software.amazon.smithy.rust.codegen.rustlang.RustModule
 import software.amazon.smithy.rust.codegen.server.smithy.testutil.serverTestSymbolProvider
 import software.amazon.smithy.rust.codegen.smithy.generators.CodegenTarget
+import software.amazon.smithy.rust.codegen.smithy.generators.error.ServerCombinedErrorGenerator
 import software.amazon.smithy.rust.codegen.smithy.transformers.OperationNormalizer
 import software.amazon.smithy.rust.codegen.testutil.TestWorkspace
 import software.amazon.smithy.rust.codegen.testutil.asSmithyModel
@@ -52,7 +53,8 @@ class ServerCombinedErrorGeneratorTest {
             listOf("FooException", "ComplexError", "InvalidGreeting").forEach {
                 model.lookup<StructureShape>("error#$it").renderWithModelBuilder(model, symbolProvider, writer, CodegenTarget.SERVER)
             }
-            val generator = ServerCombinedErrorGenerator(model, symbolProvider, model.lookup("error#Greeting"))
+            val errors = listOf("FooException", "ComplexError", "InvalidGreeting").map { model.lookup<StructureShape>("error#$it") }
+            val generator = ServerCombinedErrorGenerator(model, symbolProvider, symbolProvider.toSymbol(model.lookup("error#Greeting")), errors)
             generator.render(writer)
 
             writer.unitTest(

codegen-server/src/test/kotlin/software/amazon/smithy/rust/codegen/server/smithy/protocols/EventStreamTestTools.kt

@@ -18,7 +18,6 @@ import software.amazon.smithy.model.shapes.UnionShape
 import software.amazon.smithy.model.traits.ErrorTrait
 import software.amazon.smithy.rust.codegen.rustlang.RustModule
 import software.amazon.smithy.rust.codegen.rustlang.RustWriter
-import software.amazon.smithy.rust.codegen.server.smithy.generators.ServerCombinedErrorGenerator
 import software.amazon.smithy.rust.codegen.server.smithy.testutil.serverTestSymbolProvider
 import software.amazon.smithy.rust.codegen.smithy.CoreCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.RustSymbolProvider
@@ -27,6 +26,7 @@ import software.amazon.smithy.rust.codegen.smithy.generators.CodegenTarget
 import software.amazon.smithy.rust.codegen.smithy.generators.StructureGenerator
 import software.amazon.smithy.rust.codegen.smithy.generators.UnionGenerator
 import software.amazon.smithy.rust.codegen.smithy.generators.error.CombinedErrorGenerator
+import software.amazon.smithy.rust.codegen.smithy.generators.error.ServerCombinedErrorGenerator
 import software.amazon.smithy.rust.codegen.smithy.generators.implBlock
 import software.amazon.smithy.rust.codegen.smithy.generators.renderUnknownVariant
 import software.amazon.smithy.rust.codegen.smithy.transformers.EventStreamNormalizer
@@ -40,6 +40,7 @@ import software.amazon.smithy.rust.codegen.util.hasTrait
 import software.amazon.smithy.rust.codegen.util.lookup
 import software.amazon.smithy.rust.codegen.util.outputShape
 import java.util.stream.Stream
+import kotlin.streams.toList
 
 private fun fillInBaseModel(
     protocolName: String,
@@ -342,10 +343,15 @@ object EventStreamTestTools {
             CodegenTarget.SERVER -> serverTestSymbolProvider(model)
         }
         val project = TestWorkspace.testProject(symbolProvider)
+        val operationSymbol = symbolProvider.toSymbol(operationShape)
         project.withModule(RustModule.public("error")) {
+            val errors = model.shapes()
+                .filter { shape -> shape.isStructureShape && shape.hasTrait<ErrorTrait>() }
+                .map { it.asStructureShape().get() }
+                .toList()
             when (testCase.target) {
-                CodegenTarget.CLIENT -> CombinedErrorGenerator(model, symbolProvider, operationShape).render(it)
-                CodegenTarget.SERVER -> ServerCombinedErrorGenerator(model, symbolProvider, operationShape).render(it)
+                CodegenTarget.CLIENT -> CombinedErrorGenerator(model, symbolProvider, operationSymbol, errors).render(it)
+                CodegenTarget.SERVER -> ServerCombinedErrorGenerator(model, symbolProvider, operationSymbol, errors).render(it)
             }
             for (shape in model.shapes().filter { shape -> shape.isStructureShape && shape.hasTrait<ErrorTrait>() }) {
                 StructureGenerator(model, symbolProvider, it, shape as StructureShape).render(testCase.target)

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/EventStreamSymbolProvider.kt

@@ -51,7 +51,7 @@ class EventStreamSymbolProvider(
                 val error = if (target == CodegenTarget.SERVER && unionShape.eventStreamErrors().isEmpty()) {
                     RuntimeType("MessageStreamError", smithyEventStream, "aws_smithy_http::event_stream").toSymbol()
                 } else {
-                    unionShape.eventStreamErrorSymbol(this).toSymbol()
+                    unionShape.eventStreamErrorSymbol(model, this, target).toSymbol()
                 }
                 val errorFmt = error.rustType().render(fullyQualified = true)
                 val innerFmt = initial.rustType().stripOuter<RustType.Option>().render(fullyQualified = true)

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

@@ -83,7 +83,7 @@ class PaginatorGenerator private constructor(
 
     private val inputType = symbolProvider.toSymbol(operation.inputShape(model))
     private val outputType = operation.outputShape(model)
-    private val errorType = operation.errorSymbol(symbolProvider)
+    private val errorType = operation.errorSymbol(model, symbolProvider, CodegenTarget.CLIENT)
 
     private fun paginatorType(): RuntimeType = RuntimeType.forInlineFun(
         paginatorName,

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

@@ -12,7 +12,6 @@ import software.amazon.smithy.rust.codegen.smithy.ClientCodegenContext
 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.config.ServiceConfigGenerator
-import software.amazon.smithy.rust.codegen.smithy.generators.error.CombinedErrorGenerator
 import software.amazon.smithy.rust.codegen.smithy.generators.error.TopLevelErrorGenerator
 import software.amazon.smithy.rust.codegen.smithy.generators.protocol.ProtocolGenerator
 import software.amazon.smithy.rust.codegen.smithy.generators.protocol.ProtocolSupport
@@ -55,10 +54,6 @@ class ServiceGenerator(
                     ProtocolTestGenerator(clientCodegenContext, protocolSupport, operation, operationWriter).render()
                 }
             }
-            // Render a service-level error enum containing every error that the service can emit
-            rustCrate.withModule(RustModule.Error) { writer ->
-                CombinedErrorGenerator(clientCodegenContext.model, clientCodegenContext.symbolProvider, operation).render(writer)
-            }
         }
 
         TopLevelErrorGenerator(clientCodegenContext, operations).render(rustCrate)

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

@@ -47,6 +47,7 @@ import software.amazon.smithy.rust.codegen.smithy.customize.RustCodegenDecorator
 import software.amazon.smithy.rust.codegen.smithy.customize.Section
 import software.amazon.smithy.rust.codegen.smithy.customize.writeCustomizations
 import software.amazon.smithy.rust.codegen.smithy.expectRustMetadata
+import software.amazon.smithy.rust.codegen.smithy.generators.CodegenTarget
 import software.amazon.smithy.rust.codegen.smithy.generators.LibRsCustomization
 import software.amazon.smithy.rust.codegen.smithy.generators.LibRsSection
 import software.amazon.smithy.rust.codegen.smithy.generators.PaginatorGenerator
@@ -394,7 +395,7 @@ class FluentClientGenerator(
 
                 val output = operation.outputShape(model)
                 val operationOk = symbolProvider.toSymbol(output)
-                val operationErr = operation.errorSymbol(symbolProvider).toSymbol()
+                val operationErr = operation.errorSymbol(model, symbolProvider, CodegenTarget.CLIENT).toSymbol()
 
                 val inputFieldsBody = generateOperationShapeDocs(writer, symbolProvider, operation, model).joinToString("\n") {
                     "///   - $it"
@@ -486,7 +487,7 @@ class FluentClientGenerator(
                 ) {
                     val inputType = symbolProvider.toSymbol(operation.inputShape(model))
                     val outputType = symbolProvider.toSymbol(operation.outputShape(model))
-                    val errorType = operation.errorSymbol(symbolProvider)
+                    val errorType = operation.errorSymbol(model, symbolProvider, CodegenTarget.CLIENT)
                     rustTemplate(
                         """
                         /// Creates a new `${operationSymbol.name}`.
@@ -534,7 +535,7 @@ class FluentClientGenerator(
                         customizations,
                         FluentClientSection.FluentBuilderImpl(
                             operation,
-                            operation.errorSymbol(symbolProvider)
+                            operation.errorSymbol(model, symbolProvider, CodegenTarget.CLIENT)
                         )
                     )
                     input.members().forEach { member ->