Split response interceptor stages into afterTransmission and modifyHttpResponse

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/handler/BaseAsyncClientHandler.java

@@ -42,6 +42,7 @@
 import software.amazon.awssdk.core.internal.http.IdempotentAsyncResponseHandler;
 import software.amazon.awssdk.core.internal.http.TransformingAsyncResponseHandler;
 import software.amazon.awssdk.core.internal.http.async.AsyncAfterTransmissionInterceptorCallingResponseHandler;
+import software.amazon.awssdk.core.internal.http.async.AsyncModifyHttpResponseInterceptorCallingResponseHandler;
 import software.amazon.awssdk.core.internal.http.async.AsyncResponseHandler;
 import software.amazon.awssdk.core.internal.http.async.AsyncStreamingResponseHandler;
 import software.amazon.awssdk.core.internal.http.async.CombinedResponseAsyncHttpResponseHandler;
@@ -231,8 +232,7 @@ private <InputT extends SdkRequest, OutputT extends SdkResponse, ReturnT> Comple
                        finalizeSdkHttpRequestContext.asyncRequestBody().orElse(null),
                        inputT,
                        executionContext,
-                       new AsyncAfterTransmissionInterceptorCallingResponseHandler<>(asyncResponseHandler,
-                                                                                     executionContext));
+                       interceptorCallingResponseHandler(asyncResponseHandler, executionContext));
 
             // Captured because the requestBody branch above may reassign 'marshalled', leaving it not effectively final.
             SdkHttpFullRequest requestForMetrics = marshalled;
@@ -259,6 +259,17 @@ public void close() {
         client.close();
     }
 
+    /**
+     * Wraps the response handler with the response interceptor steps, in order: read-only {@code afterTransmission} on
+     * the raw response, then the interceptors that can modify it.
+     */
+    static <T> TransformingAsyncResponseHandler<T> interceptorCallingResponseHandler(
+        TransformingAsyncResponseHandler<T> responseHandler, ExecutionContext executionContext) {
+        return new AsyncAfterTransmissionInterceptorCallingResponseHandler<>(
+            new AsyncModifyHttpResponseInterceptorCallingResponseHandler<>(responseHandler, executionContext),
+            executionContext);
+    }
+
     /**
      * Error responses are never streaming so we always use {@link AsyncResponseHandler}.
      *

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/AmazonSyncHttpClient.java

@@ -46,6 +46,7 @@
 import software.amazon.awssdk.core.internal.http.pipeline.stages.MakeRequestMutableStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.MergeCustomHeadersStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.MergeCustomQueryParamsStage;
+import software.amazon.awssdk.core.internal.http.pipeline.stages.ModifyHttpResponseExecutionInterceptorsStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.QueryParametersToBodyStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.RetryableStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.SigningStage;
@@ -197,6 +198,7 @@ public <OutputT> OutputT execute(HttpResponseHandler<Response<OutputT>> response
                                          .then(BeforeTransmissionExecutionInterceptorsStage::new)
                                          .then(MakeHttpRequestStage::new)
                                          .then(AfterTransmissionExecutionInterceptorsStage::new)
+                                         .then(ModifyHttpResponseExecutionInterceptorsStage::new)
                                          .then(BeforeUnmarshallingExecutionInterceptorsStage::new)
                                          .then(() -> new HandleResponseStage<>(responseHandler))
                                          .wrappedWith(ApiCallAttemptTimeoutTrackingStage::new)

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/async/AsyncAfterTransmissionInterceptorCallingResponseHandler.java

@@ -16,18 +16,17 @@
 package software.amazon.awssdk.core.internal.http.async;
 
 import java.nio.ByteBuffer;
-import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import org.reactivestreams.Publisher;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.core.http.ExecutionContext;
 import software.amazon.awssdk.core.interceptor.InterceptorContext;
 import software.amazon.awssdk.core.internal.http.TransformingAsyncResponseHandler;
-import software.amazon.awssdk.http.SdkHttpFullResponse;
 import software.amazon.awssdk.http.SdkHttpResponse;
 
 /**
- * Async response handler decorator to run interceptors after response is received.
+ * Async response handler decorator that calls the read-only {@code afterTransmission} interceptors with the response
+ * exactly as received.
  *
  * @param <T> the type of the result
  */
@@ -42,29 +41,18 @@ public AsyncAfterTransmissionInterceptorCallingResponseHandler(TransformingAsync
         this.context = context;
     }
 
-    private SdkHttpResponse beforeUnmarshalling(SdkHttpFullResponse response, ExecutionContext context) {
+    @Override
+    public void onHeaders(SdkHttpResponse response) {
         // Update interceptor context to include response
-        InterceptorContext interceptorContext =
-            context.interceptorContext().copy(b -> b.httpResponse(response));
+        InterceptorContext interceptorContext = context.interceptorContext().copy(b -> b.httpResponse(response));
 
         // interceptors.afterTransmission
         context.interceptorChain().afterTransmission(interceptorContext, context.executionAttributes());
 
-        // interceptors.modifyHttpResponse
-        interceptorContext = context.interceptorChain().modifyHttpResponse(interceptorContext, context.executionAttributes());
-
-        // interceptors.beforeUnmarshalling
-        context.interceptorChain().beforeUnmarshalling(interceptorContext, context.executionAttributes());
-
         // Store updated context
         context.interceptorContext(interceptorContext);
 
-        return interceptorContext.httpResponse();
-    }
-
-    @Override
-    public void onHeaders(SdkHttpResponse response) {
-        delegate.onHeaders(beforeUnmarshalling((SdkHttpFullResponse) response, context)); // TODO: Ew
+        delegate.onHeaders(response);
     }
 
     @Override
@@ -74,23 +62,11 @@ public void onError(Throwable error) {
 
     @Override
     public void onStream(Publisher<ByteBuffer> publisher) {
-        Optional<Publisher<ByteBuffer>> newPublisher = context.interceptorChain()
-                                                              .modifyAsyncHttpResponse(context.interceptorContext()
-                                                                                              .toBuilder()
-                                                                                              .responsePublisher(publisher)
-                                                                                              .build(),
-                                                                                       context.executionAttributes())
-                                                              .responsePublisher();
-
-        if (newPublisher.isPresent()) {
-            delegate.onStream(newPublisher.get());
-        } else {
-            delegate.onStream(publisher);
-        }
+        delegate.onStream(publisher);
     }
 
     @Override
     public CompletableFuture<T> prepare() {
         return delegate.prepare();
     }
-}
+}

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/async/AsyncModifyHttpResponseInterceptorCallingResponseHandler.java

@@ -0,0 +1,84 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.core.internal.http.async;
+
+import java.nio.ByteBuffer;
+import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+import org.reactivestreams.Publisher;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.core.http.ExecutionContext;
+import software.amazon.awssdk.core.interceptor.InterceptorContext;
+import software.amazon.awssdk.core.internal.http.TransformingAsyncResponseHandler;
+import software.amazon.awssdk.http.SdkHttpResponse;
+
+/**
+ * Async response handler decorator that calls the {@code modifyHttpResponse} and {@code beforeUnmarshalling} interceptors
+ * when headers arrive, and the {@code modifyAsyncHttpResponse} interceptors when the body arrives.
+ *
+ * @param <T> the type of the result
+ */
+@SdkInternalApi
+public final class AsyncModifyHttpResponseInterceptorCallingResponseHandler<T> implements TransformingAsyncResponseHandler<T> {
+    private final TransformingAsyncResponseHandler<T> delegate;
+    private final ExecutionContext context;
+
+    public AsyncModifyHttpResponseInterceptorCallingResponseHandler(TransformingAsyncResponseHandler<T> delegate,
+                                                                    ExecutionContext context) {
+        this.delegate = delegate;
+        this.context = context;
+    }
+
+    @Override
+    public void onHeaders(SdkHttpResponse response) {
+        // Update interceptor context to include response
+        InterceptorContext interceptorContext = context.interceptorContext().copy(b -> b.httpResponse(response));
+
+        // interceptors.modifyHttpResponse
+        interceptorContext = context.interceptorChain().modifyHttpResponse(interceptorContext, context.executionAttributes());
+
+        // interceptors.beforeUnmarshalling
+        context.interceptorChain().beforeUnmarshalling(interceptorContext, context.executionAttributes());
+
+        // Store updated context
+        context.interceptorContext(interceptorContext);
+
+        delegate.onHeaders(interceptorContext.httpResponse());
+    }
+
+    @Override
+    public void onError(Throwable error) {
+        delegate.onError(error);
+    }
+
+    @Override
+    public void onStream(Publisher<ByteBuffer> publisher) {
+        Optional<Publisher<ByteBuffer>> newPublisher = context.interceptorChain()
+                                                              .modifyAsyncHttpResponse(context.interceptorContext()
+                                                                                              .toBuilder()
+                                                                                              .responsePublisher(publisher)
+                                                                                              .build(),
+                                                                                       context.executionAttributes())
+                                                              .responsePublisher();
+
+        delegate.onStream(newPublisher.orElse(publisher));
+    }
+
+    @Override
+    public CompletableFuture<T> prepare() {
+        return delegate.prepare();
+    }
+}

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/AfterTransmissionExecutionInterceptorsStage.java

@@ -21,48 +21,33 @@
 import software.amazon.awssdk.core.internal.http.InterruptMonitor;
 import software.amazon.awssdk.core.internal.http.RequestExecutionContext;
 import software.amazon.awssdk.core.internal.http.pipeline.RequestPipeline;
-import software.amazon.awssdk.http.AbortableInputStream;
 import software.amazon.awssdk.http.SdkHttpFullRequest;
 import software.amazon.awssdk.http.SdkHttpFullResponse;
 import software.amazon.awssdk.utils.Pair;
 
+/**
+ * Calls the read-only {@code afterTransmission} interceptors with the response exactly as received.
+ */
 @SdkInternalApi
 public class AfterTransmissionExecutionInterceptorsStage
     implements RequestPipeline<Pair<SdkHttpFullRequest, SdkHttpFullResponse>, Pair<SdkHttpFullRequest, SdkHttpFullResponse>> {
     @Override
     public Pair<SdkHttpFullRequest, SdkHttpFullResponse> execute(Pair<SdkHttpFullRequest, SdkHttpFullResponse> input,
                                                                  RequestExecutionContext context) throws Exception {
         InterruptMonitor.checkInterrupted(input.right());
-        InputStream originalResponseBody = input.right().content().orElse(null);
+        InputStream responseBody = input.right().content().orElse(null);
 
         // Update interceptor context
         InterceptorContext interceptorContext =
                 context.executionContext().interceptorContext().copy(b -> b.httpResponse(input.right())
-                                                                           .responseBody(originalResponseBody));
+                                                                           .responseBody(responseBody));
 
         // interceptors.afterTransmission
         context.interceptorChain().afterTransmission(interceptorContext, context.executionAttributes());
 
-        // interceptors.modifyHttpResponse
-        interceptorContext = context.interceptorChain().modifyHttpResponse(interceptorContext, context.executionAttributes());
-
-        // When no replacement is returned, the response retains its original body. Keep the context aligned with the
-        // actual response so later interceptor failures can clean up that body.
-        if (!interceptorContext.responseBody().isPresent() && originalResponseBody != null) {
-            interceptorContext = interceptorContext.copy(b -> b.responseBody(originalResponseBody));
-        }
-
-        // Store updated context
+        // Store updated context, so later stages can clean up the body if they fail
         context.executionContext().interceptorContext(interceptorContext);
 
-        InterruptMonitor.checkInterrupted((SdkHttpFullResponse) interceptorContext.httpResponse());
-
-        SdkHttpFullResponse response = (SdkHttpFullResponse) interceptorContext.httpResponse();
-
-        if (interceptorContext.responseBody().isPresent()) {
-            response = response.toBuilder().content(AbortableInputStream.create(interceptorContext.responseBody().get())).build();
-        }
-
-        return Pair.of(input.left(), response);
+        return input;
     }
 }

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/ModifyHttpResponseExecutionInterceptorsStage.java

@@ -0,0 +1,67 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.core.internal.http.pipeline.stages;
+
+import java.io.InputStream;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.core.interceptor.InterceptorContext;
+import software.amazon.awssdk.core.internal.http.InterruptMonitor;
+import software.amazon.awssdk.core.internal.http.RequestExecutionContext;
+import software.amazon.awssdk.core.internal.http.pipeline.RequestPipeline;
+import software.amazon.awssdk.http.AbortableInputStream;
+import software.amazon.awssdk.http.SdkHttpFullRequest;
+import software.amazon.awssdk.http.SdkHttpFullResponse;
+import software.amazon.awssdk.utils.Pair;
+
+/**
+ * Calls the {@code modifyHttpResponse} interceptors and applies their changes to the response.
+ */
+@SdkInternalApi
+public class ModifyHttpResponseExecutionInterceptorsStage
+    implements RequestPipeline<Pair<SdkHttpFullRequest, SdkHttpFullResponse>, Pair<SdkHttpFullRequest, SdkHttpFullResponse>> {
+    @Override
+    public Pair<SdkHttpFullRequest, SdkHttpFullResponse> execute(Pair<SdkHttpFullRequest, SdkHttpFullResponse> input,
+                                                                 RequestExecutionContext context) throws Exception {
+        InputStream responseBody = input.right().content().orElse(null);
+
+        // Update interceptor context
+        InterceptorContext interceptorContext =
+                context.executionContext().interceptorContext().copy(b -> b.httpResponse(input.right())
+                                                                           .responseBody(responseBody));
+
+        // interceptors.modifyHttpResponse
+        interceptorContext = context.interceptorChain().modifyHttpResponse(interceptorContext, context.executionAttributes());
+
+        // When no replacement is returned, the response retains its original body. Keep the context aligned with the
+        // actual response so later interceptor failures can clean up that body.
+        if (!interceptorContext.responseBody().isPresent() && responseBody != null) {
+            interceptorContext = interceptorContext.copy(b -> b.responseBody(responseBody));
+        }
+
+        // Store updated context
+        context.executionContext().interceptorContext(interceptorContext);
+
+        InterruptMonitor.checkInterrupted((SdkHttpFullResponse) interceptorContext.httpResponse());
+
+        SdkHttpFullResponse response = (SdkHttpFullResponse) interceptorContext.httpResponse();
+
+        if (interceptorContext.responseBody().isPresent()) {
+            response = response.toBuilder().content(AbortableInputStream.create(interceptorContext.responseBody().get())).build();
+        }
+
+        return Pair.of(input.left(), response);
+    }
+}