From 4a126b797161d2b476aa0ffe4b2236f1d8c659f5 Mon Sep 17 00:00:00 2001
From: John DiSanti <jdisanti@amazon.com>
Date: Fri, 18 Aug 2023 08:18:25 -0700
Subject: [PATCH] Fix the re-export in middleware mode

---
 .../customize/RequiredCustomizations.kt       | 23 +++++++++++++------
 .../ServerRequiredCustomizations.kt           |  5 ++--
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customize/RequiredCustomizations.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customize/RequiredCustomizations.kt
index bf46d496d1..f2228bba17 100644
--- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customize/RequiredCustomizations.kt
+++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customize/RequiredCustomizations.kt
@@ -24,6 +24,7 @@ import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationCus
 import software.amazon.smithy.rust.codegen.client.smithy.generators.ServiceRuntimePluginCustomization
 import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ConfigCustomization
 import software.amazon.smithy.rust.codegen.core.rustlang.Feature
+import software.amazon.smithy.rust.codegen.core.rustlang.rust
 import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
@@ -94,18 +95,26 @@ class RequiredCustomizations : ClientCodegenDecorator {
             pubUseSmithyPrimitives(codegenContext, codegenContext.model)(this)
         }
         rustCrate.withModule(ClientRustModule.Error) {
+            // TODO(enableNewSmithyRuntimeCleanup): Change SdkError to a `pub use` after changing the generic's default
+            rust("/// Error type returned by the client.")
+            if (codegenContext.smithyRuntimeMode.generateOrchestrator) {
+                rustTemplate(
+                    "pub type SdkError<E> = #{SdkError}<E, #{R}>;",
+                    "SdkError" to RuntimeType.sdkError(rc),
+                    "R" to RuntimeType.smithyRuntimeApi(rc).resolve("client::orchestrator::HttpResponse"),
+                )
+            } else {
+                rustTemplate(
+                    "pub type SdkError<E> = #{SdkError}<E, #{R}>;",
+                    "SdkError" to RuntimeType.sdkError(rc),
+                    "R" to RuntimeType.smithyHttp(rc).resolve("operation::Response"),
+                )
+            }
             rustTemplate(
                 """
-                pub type SdkError<E> = #{SdkError}<E, #{SdkErrorResponse}>;
                 pub use #{DisplayErrorContext};
                 pub use #{ProvideErrorMetadata};
                 """,
-                "SdkError" to RuntimeType.smithyHttp(rc).resolve("result::SdkError"),
-                "SdkErrorResponse" to if (codegenContext.smithyRuntimeMode.generateOrchestrator) {
-                    RuntimeType.smithyRuntimeApi(rc).resolve("client::orchestrator::HttpResponse")
-                } else {
-                    RuntimeType.HttpResponse
-                },
                 "DisplayErrorContext" to RuntimeType.smithyTypes(rc).resolve("error::display::DisplayErrorContext"),
                 "ProvideErrorMetadata" to RuntimeType.smithyTypes(rc).resolve("error::metadata::ProvideErrorMetadata"),
             )
diff --git a/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/customizations/ServerRequiredCustomizations.kt b/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/customizations/ServerRequiredCustomizations.kt
index b79ff3b287..a8698ae91d 100644
--- a/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/customizations/ServerRequiredCustomizations.kt
+++ b/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/customizations/ServerRequiredCustomizations.kt
@@ -42,13 +42,14 @@ class ServerRequiredCustomizations : ServerCodegenDecorator {
 
         rustCrate.withModule(ServerRustModule.Types) {
             pubUseSmithyPrimitives(codegenContext, codegenContext.model)(this)
+            // TODO(enableNewSmithyRuntimeCleanup): Remove re-export of SdkError in server and add changelog entry
             rustTemplate(
                 """
-                pub type SdkError<E> = #{SdkError}<E, #{SdkErrorResponse}>;
+                pub type SdkError<E> = #{SdkError}<E, #{Response}>;
                 pub use #{DisplayErrorContext};
                 """,
                 "SdkError" to RuntimeType.smithyHttp(rc).resolve("result::SdkError"),
-                "SdkErrorResponse" to RuntimeType.HttpResponse,
+                "Response" to RuntimeType.smithyHttp(rc).resolve("operation::Response"),
                 "DisplayErrorContext" to RuntimeType.smithyTypes(rc).resolve("error::display::DisplayErrorContext"),
             )
         }