Handle additionalInput without modifying the actual input model

aws/aws-service-bundle/src/main/java/software/amazon/smithy/java/aws/servicebundle/provider/AwsServiceBundle.java

@@ -25,8 +25,8 @@
 import software.amazon.smithy.java.client.core.endpoint.EndpointResolver;
 import software.amazon.smithy.java.client.core.interceptors.CallHook;
 import software.amazon.smithy.java.client.core.interceptors.ClientInterceptor;
-import software.amazon.smithy.java.core.serde.document.Document;
 import software.amazon.smithy.java.dynamicclient.DynamicClient;
+import software.amazon.smithy.java.server.ProxyService;
 import software.amazon.smithy.model.knowledge.ServiceIndex;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.modelbundle.api.BundlePlugin;
@@ -87,9 +87,7 @@ private record AwsServiceClientInterceptor(AwsServiceMetadata serviceMetadata, A
 
         @Override
         public ClientConfig modifyBeforeCall(CallHook<?, ?> hook) {
-            if (!(hook.input() instanceof Document d)) {
-                throw new IllegalArgumentException("Input must be a Document");
-            }
+            var d = hook.config().context().get(ProxyService.PROXY_INPUT);
             var input = d.asShape(PreRequest.builder());
 
             var endpoint = URI.create(Objects.requireNonNull(serviceMetadata.getEndpoints().get(input.getAwsRegion()),

client/client-core/src/main/java/software/amazon/smithy/java/client/core/Client.java

@@ -80,12 +80,21 @@ protected <I extends SerializableStruct, O extends SerializableStruct> Completab
         IdentityResolvers callIdentityResolvers = identityResolvers;
         ClientInterceptor callInterceptor = interceptor;
 
-        // First apply overrides from interceptors.
-        ClientConfig callConfig = callInterceptor.modifyBeforeCall(new CallHook<>(operation, config, input));
-        // Overrides given per/operation take precedence over interceptors.
+        //If there is an override config first apply that before sending to interceptors.
+        ClientConfig callConfig = config;
         if (overrideConfig != null) {
             callConfig = callConfig.withRequestOverride(overrideConfig);
         }
+        ClientConfig afterInterceptionConfig =
+                callInterceptor.modifyBeforeCall(new CallHook<>(operation, callConfig, input));
+        if (afterInterceptionConfig != null && afterInterceptionConfig != callConfig) {
+            if (overrideConfig != null) {
+                callConfig = afterInterceptionConfig.withRequestOverride(overrideConfig);
+            } else {
+                callConfig = afterInterceptionConfig;
+            }
+
+        }
 
         // Rebuild the pipeline, resolvers, etc if the config changed.
         if (callConfig != config) {

client/client-core/src/test/java/software/amazon/smithy/java/client/core/ClientTest.java

@@ -169,6 +169,8 @@ public void requestOverridesPerCallTakePrecedence() throws URISyntaxException {
                 .addPlugin(config -> config.addInterceptor(new ClientInterceptor() {
                     @Override
                     public ClientConfig modifyBeforeCall(CallHook<?, ?> hook) {
+                        //RequestOverrides config should be visible here.
+                        assertThat(hook.config().context().get(CallContext.APPLICATION_ID), equalTo(id));
                         // Note that the overrides given to the call itself will override interceptors.
                         var override = RequestOverrideConfig.builder()
                                 .putConfig(CallContext.APPLICATION_ID, "foo")

model-bundle/model-bundle-api/src/main/java/software/amazon/smithy/modelbundle/api/ModelBundles.java

@@ -5,10 +5,14 @@
 
 package software.amazon.smithy.modelbundle.api;
 
+import software.amazon.smithy.java.server.ProxyOperationTrait;
 import software.amazon.smithy.java.server.ProxyService;
 import software.amazon.smithy.java.server.Service;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.loader.ModelAssembler;
+import software.amazon.smithy.model.shapes.MemberShape;
+import software.amazon.smithy.model.shapes.OperationShape;
+import software.amazon.smithy.model.shapes.ServiceShape;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.shapes.StructureShape;
 import software.amazon.smithy.model.traits.StreamingTrait;
@@ -50,59 +54,84 @@ private static Model getModel(SmithyBundle bundle) {
         var b = model.toBuilder();
 
         // mix in the generic arg members
-        for (var op : model.getOperationShapes()) {
-            boolean skipOperation = false;
-            if (op.getOutput().isPresent()) {
-                for (var member : model.expectShape(op.getOutputShape(), StructureShape.class).members()) {
-                    if (model.expectShape(member.getTarget()).hasTrait(StreamingTrait.class)) {
-                        b.removeShape(op.toShapeId());
-                        skipOperation = true;
-                        break;
-                    }
-                }
-            }
-
-            if (skipOperation) {
-                continue;
-            }
-
-            if (op.getInput().isEmpty() && additionalInputShape != null) {
-                b.addShape(op.toBuilder()
-                        .input(additionalInputShape)
-                        .build());
-            } else {
-                var shape = model.expectShape(op.getInputShape(), StructureShape.class);
-                for (var member : shape.members()) {
-                    if (model.expectShape(member.getTarget()).hasTrait(StreamingTrait.class)) {
-                        b.removeShape(op.toShapeId());
-                        skipOperation = true;
-                        break;
+        for (var service : model.getServiceShapes()) {
+            var serviceShape = service.asServiceShape().get();
+            var serviceBuilder = serviceShape.toBuilder();
+            for (var opId : serviceShape.getAllOperations()) {
+                var op = model.expectShape(opId, OperationShape.class);
+                boolean skipOperation = false;
+                if (op.getOutput().isPresent()) {
+                    for (var member : model.expectShape(op.getOutputShape(), StructureShape.class).members()) {
+                        if (model.expectShape(member.getTarget()).hasTrait(StreamingTrait.class)) {
+                            b.removeShape(op.toShapeId());
+                            skipOperation = true;
+                            break;
+                        }
                     }
                 }
 
                 if (skipOperation) {
                     continue;
                 }
 
-                if (additionalInputShape != null) {
-                    var input = shape.toBuilder();
-                    for (var member : additionalInputShape.members()) {
-                        input.addMember(member.toBuilder()
-                                .id(ShapeId.from(input.getId().toString() + "$" + member.getMemberName()))
-                                .build());
+                if (op.getInput().isEmpty() && additionalInputShape != null) {
+                    addProxyOperationWithAdditionalInput(op, additionalInputShape, b, serviceBuilder, model);
+                } else {
+                    var shape = model.expectShape(op.getInputShape(), StructureShape.class);
+                    for (var member : shape.members()) {
+                        if (model.expectShape(member.getTarget()).hasTrait(StreamingTrait.class)) {
+                            b.removeShape(op.toShapeId());
+                            skipOperation = true;
+                            break;
+                        }
+                    }
+
+                    if (skipOperation) {
+                        continue;
+                    }
+
+                    if (additionalInputShape != null) {
+                        addProxyOperationWithAdditionalInput(op, additionalInputShape, b, serviceBuilder, model);
                     }
-                    b.addShape(input.build());
                 }
             }
-        }
 
-        for (var service : model.getServiceShapes()) {
-            b.addShape(service.toBuilder()
+            b.addShape(serviceBuilder
                     // trim the endpoint rules because they're huge and we don't need them
                     .removeTrait(ShapeId.from("smithy.rules#endpointRuleSet"))
                     .removeTrait(ShapeId.from("smithy.rules#endpointTests"))
                     .build());
         }
         return b.build();
     }
+
+    private static void addProxyOperationWithAdditionalInput(
+            OperationShape op,
+            StructureShape additionalInput,
+            Model.Builder builder,
+            ServiceShape.Builder serviceBuilder,
+            Model model
+    ) {
+        var input = op.getInput();
+        StructureShape finalInput;
+        if (op.getInput().isEmpty()) {
+            finalInput = additionalInput;
+        } else {
+            var inputBuilder = model.expectShape(input.get(), StructureShape.class).toBuilder();
+            inputBuilder.addMember(MemberShape.builder()
+                    .id(ShapeId.from(inputBuilder.getId().toString() + "$additionalInput"))
+                    .target(additionalInput.getId())
+                    .build());
+            finalInput = inputBuilder.id(ShapeId.from(inputBuilder.getId().toString()) + "Proxy").build();
+        }
+        builder.addShape(finalInput);
+        var newOperation = op.toBuilder()
+                .id(ShapeId.from(op.getId().toString() + "Proxy"))
+                .input(finalInput)
+                .output(op.getOutputShape())
+                .addTrait(new ProxyOperationTrait(op.getId()))
+                .build();
+        builder.addShape(newOperation);
+        serviceBuilder.addOperation(newOperation).build();
+    }
 }

server/server-proxy/src/main/java/software/amazon/smithy/java/server/ProxyOperationTrait.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.java.server;
+
+import software.amazon.smithy.model.node.Node;
+import software.amazon.smithy.model.shapes.ShapeId;
+import software.amazon.smithy.model.traits.Trait;
+
+/**
+ * A trait that marks operations as proxies to delegate operations.
+ * 
+ * <p>Operations annotated with this trait are treated as proxies to the 
+ * delegate operation specified by this trait. Proxy operations enable 
+ * additional input processing while maintaining compatibility with the 
+ * original operation interface.</p>
+ * 
+ * <h3>Proxy Operation Input Handling</h3>
+ * <p>Proxy operation inputs are expected to be a superset of the original 
+ * operation input. For operations with existing input, a new input structure 
+ * is created that contains all original input members plus an 
+ * {@code additionalInput} field containing the additional parameters. For 
+ * operations with no input, the additional input structure is used directly.</p>
+ * 
+ * <h3>Additional Input Access</h3>
+ * <p>The additional input data can be accessed in Dynamic client interceptors 
+ * using the {@code ProxyService.PROXY_INPUT} context key. This enables 
+ * interceptors to process the extra data before the request is forwarded 
+ * to the delegate service.</p>
+ * 
+ * <h3>Input Stripping</h3>
+ * <p>The proxy service automatically strips out the additional input before 
+ * sending the request to the delegate service, ensuring that only the 
+ * expected input parameters are forwarded to the original operation.</p>
+ * 
+ * @see ProxyService#PROXY_INPUT
+ */
+public class ProxyOperationTrait implements Trait {
+
+    private static final ShapeId SHAPE_ID = ShapeId.from("smithy.server.api#proxyOperation");
+    private final ShapeId delegateOperation;
+
+    public ProxyOperationTrait(ShapeId delegateOperation) {
+        this.delegateOperation = delegateOperation;
+    }
+
+    @Override
+    public ShapeId toShapeId() {
+        return SHAPE_ID;
+    }
+
+    @Override
+    public Node toNode() {
+        return null;
+    }
+
+    public ShapeId getDelegateOperation() {
+        return delegateOperation;
+    }
+}