Model bundle-specific args in the bundle itself

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

@@ -33,11 +33,11 @@
 import software.amazon.smithy.modelbundle.api.BundleClientPluginProvider;
 import software.amazon.smithy.modelbundle.api.StaticAuthSchemePlugin;
 
-final class AwsAuthProvider implements BundleClientPluginProvider {
+final class AwsServiceBundle implements BundleClientPluginProvider {
     private final AwsServiceMetadata serviceMetadata;
     private final AwsServicePlugin plugin;
 
-    AwsAuthProvider(AwsServiceMetadata serviceMetadata) {
+    AwsServiceBundle(AwsServiceMetadata serviceMetadata) {
         this.serviceMetadata = serviceMetadata;
         this.plugin = new AwsServicePlugin();
     }

aws/aws-service-bundle/src/main/java/software/amazon/smithy/java/aws/servicebundle/provider/AwsAuthProviderFactory.java → aws/aws-service-bundle/src/main/java/software/amazon/smithy/java/aws/servicebundle/provider/AwsServiceBundlePluginFactory.java

@@ -10,8 +10,8 @@
 import software.amazon.smithy.modelbundle.api.BundleClientPluginProvider;
 import software.amazon.smithy.modelbundle.api.BundleClientPluginProviderFactory;
 
-public final class AwsAuthProviderFactory implements BundleClientPluginProviderFactory {
-    public AwsAuthProviderFactory() {
+public final class AwsServiceBundlePluginFactory implements BundleClientPluginProviderFactory {
+    public AwsServiceBundlePluginFactory() {
 
     }
 
@@ -21,7 +21,7 @@ public String identifier() {
     }
 
     @Override
-    public BundleClientPluginProvider createAuthFactory(Document input) {
-        return new AwsAuthProvider(input.asShape(AwsServiceMetadata.builder()));
+    public BundleClientPluginProvider createPluginProvider(Document input) {
+        return new AwsServiceBundle(input.asShape(AwsServiceMetadata.builder()));
     }
 }

aws/aws-service-bundle/src/main/resources/META-INF/services/software.amazon.smithy.modelbundle.api.BundleClientPluginProviderFactory

@@ -1 +1 @@
-software.amazon.smithy.java.aws.servicebundle.provider.AwsAuthProviderFactory
+software.amazon.smithy.java.aws.servicebundle.provider.AwsServiceBundlePluginFactory

aws/aws-service-bundler/src/main/java/software/amazon/smithy/java/aws/servicebundle/bundler/AwsServiceBundler.java

@@ -17,6 +17,7 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
+import java.util.stream.Collectors;
 import software.amazon.smithy.aws.traits.auth.SigV4Trait;
 import software.amazon.smithy.awsmcp.model.AwsServiceMetadata;
 import software.amazon.smithy.awsmcp.model.PreRequest;
@@ -27,10 +28,10 @@
 import software.amazon.smithy.model.node.ObjectNode;
 import software.amazon.smithy.model.shapes.ModelSerializer;
 import software.amazon.smithy.model.shapes.ShapeId;
-import software.amazon.smithy.model.shapes.StructureShape;
 import software.amazon.smithy.model.traits.EndpointTrait;
 import software.amazon.smithy.modelbundle.api.Bundler;
 import software.amazon.smithy.modelbundle.api.model.Bundle;
+import software.amazon.smithy.modelbundle.api.model.GenericArguments;
 import software.amazon.smithy.modelbundle.api.model.Model;
 
 final class AwsServiceBundler implements Bundler {
@@ -63,32 +64,6 @@ final class AwsServiceBundler implements Bundler {
         this.resolver = resolver;
     }
 
-    private static software.amazon.smithy.model.Model adapt(software.amazon.smithy.model.Model model) {
-        var template = model.expectShape(PreRequest.$ID).asStructureShape().get();
-        var b = model.toBuilder();
-
-        // mix in the PreRequest structure members
-        for (var op : model.getOperationShapes()) {
-            var input = model.expectShape(op.getInput().get(), StructureShape.class).toBuilder();
-            for (var member : template.members()) {
-                input.addMember(member.toBuilder()
-                        .id(ShapeId.from(input.getId().toString() + "$" + member.getMemberName()))
-                        .build());
-            }
-            b.addShape(input.build());
-        }
-
-        for (var service : model.getServiceShapes()) {
-            b.addShape(service.toBuilder()
-                    // trim the endpoint rules because they're huge and we don't need them
-                    .removeTrait(ShapeId.from("smithy.rules#endpointRuleSet"))
-                    .removeTrait(ENDPOINT_TESTS)
-                    .build());
-        }
-
-        return b.build();
-    }
-
     @Override
     public Bundle bundle() {
         try {
@@ -122,15 +97,36 @@ public Bundle bundle() {
                     .configType("aws")
                     .serviceName(model.getServiceShapes().iterator().next().getId().toString())
                     .model(Model.builder()
-                            .smithyModel(
-                                    ObjectNode.printJson(ModelSerializer.builder().build().serialize(adapt(model))))
+                            .smithyModel(serializeModel(model))
+                            .build())
+                    .requestArguments(GenericArguments.builder()
+                            .identifier(PreRequest.$ID.toString())
+                            .model(Model.builder()
+                                    .smithyModel(loadModel("/META-INF/smithy/bundle.smithy"))
+                                    .build())
                             .build())
                     .build();
         } catch (Exception e) {
             throw new RuntimeException("Failed to bundle " + serviceName, e);
         }
     }
 
+    private static String serializeModel(software.amazon.smithy.model.Model model) {
+        return ObjectNode.printJson(ModelSerializer.builder()
+                .build()
+                .serialize(model));
+    }
+
+    private static String loadModel(String path) {
+        try (var reader = new BufferedReader(new InputStreamReader(
+                Objects.requireNonNull(AwsServiceBundler.class.getResourceAsStream(path)),
+                StandardCharsets.UTF_8))) {
+            return reader.lines().collect(Collectors.joining());
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
     Map<String, String> parseEndpoints(ObjectNode endpointTests) {
         var testCases = endpointTests.expectArrayMember("testCases");
         var endpoints = new HashMap<String, String>();

model-bundler/bundle-api/build.gradle.kts

@@ -47,11 +47,3 @@ tasks.sourcesJar {
 tasks.processResources {
     dependsOn("compileJava")
 }
-
-sourceSets {
-    main {
-        java {
-            srcDir("model")
-        }
-    }
-}

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

@@ -8,18 +8,7 @@
 import software.amazon.smithy.java.client.core.ClientPlugin;
 
 /**
- * A ConfigProvider is used to parse a bundle of service information (model, auth configuration, endpoints, etc.) and
- * configure outgoing client calls as necessary.
- *
- * <p>Implementations of this interface can define a wrapper type that adds additional parameters to vended MCP tools.
- * For example, an AWS auth provider can make a wrapper that adds the region and AWS credential profile name as
- * arguments to tools generated for AWS APIs. A wrapper type does not need to be defined if no per-request parameters
- * need to be injected.
- *
- * <p>The ConfigProvider is responsible for configuring outbound client calls with endpoint, identity, and auth resolver
- * mechanisms. The default implementation of {@link #adaptConfig(T)} orchestrates the calls to all other ConfigProvider
- * APIs and should not be overridden. If an override is needed, the {@code super} method should be called and the
- * returned RequestOverrideConfig.Builder should be modified.
+ * A factory for creating {@link ClientPlugin} instances from a service bundle.
  */
 public interface BundleClientPluginProvider {
     /**

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

@@ -10,5 +10,5 @@
 public interface BundleClientPluginProviderFactory {
     String identifier();
 
-    BundleClientPluginProvider createAuthFactory(Document input);
+    BundleClientPluginProvider createPluginProvider(Document input);
 }

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

@@ -22,7 +22,7 @@ public BundleClientPluginProvider getProvider(String identifier, Document input)
             throw new NullPointerException("no auth provider named " + identifier);
         }
 
-        return provider.createAuthFactory(input);
+        return provider.createPluginProvider(input);
     }
 
     public static Builder builder() {

model-bundler/bundle-api/src/main/resources/META-INF/smithy/bundle.smithy

@@ -0,0 +1,37 @@
+$version: "2"
+
+namespace software.amazon.smithy.modelbundle.api
+
+structure Bundle {
+    /// unique identifier for the configuration type. used to resolve the appropriate Bundler.
+    @required
+    configType: String
+
+    /// fully-qualified ShapeId of the service
+    @required
+    serviceName: String
+
+    /// Bundle-specific configuration. If this bundle does not require configuration, this
+    /// field may be omitted.
+    config: Document
+
+    /// model that describes the service. The service given in `serviceName` must be present.
+    @required
+    model: Model
+
+    /// model describing the generic arguments that must be present in every request. If this
+    /// bundle does not require generic arguments, this field may be omitted.
+    requestArguments: GenericArguments
+}
+
+union Model {
+    smithyModel: String
+}
+
+structure GenericArguments {
+    @required
+    identifier: String
+
+    @required
+    model: Model
+}

model-bundler/bundle-api/src/main/resources/META-INF/smithy/manifest

@@ -0,0 +1 @@
+bundle.smithy

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

@@ -33,6 +33,7 @@
 import software.amazon.smithy.model.shapes.ServiceShape;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.shapes.ShapeType;
+import software.amazon.smithy.model.shapes.StructureShape;
 import software.amazon.smithy.model.shapes.ToShapeId;
 import software.amazon.smithy.modelbundle.api.PluginProviders;
 import software.amazon.smithy.modelbundle.api.model.Bundle;
@@ -50,8 +51,46 @@ public final class ProxyService implements Service {
     private final List<Operation<? extends SerializableStruct, ? extends SerializableStruct>> allOperations;
     private final Schema schema;
 
+    private static software.amazon.smithy.model.Model adapt(Builder builder) {
+        if (builder.bundle == null || builder.bundle.getRequestArguments() == null) {
+            return builder.model;
+        }
+
+        var args = builder.bundle.getRequestArguments();
+        var model = new ModelAssembler()
+                .addModel(builder.model)
+                .addModel(args.getModel().getValue())
+                .assemble()
+                .unwrap();
+        var template = model.expectShape(ShapeId.from(args.getIdentifier()))
+                .asStructureShape()
+                .get();
+        var b = model.toBuilder();
+
+        // mix in the generic arg members
+        for (var op : model.getOperationShapes()) {
+            var input = model.expectShape(op.getInput().get(), StructureShape.class).toBuilder();
+            for (var member : template.members()) {
+                input.addMember(member.toBuilder()
+                        .id(ShapeId.from(input.getId().toString() + "$" + member.getMemberName()))
+                        .build());
+            }
+            b.addShape(input.build());
+        }
+
+        for (var service : model.getServiceShapes()) {
+            b.addShape(service.toBuilder()
+                    // 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 ProxyService(Builder builder) {
-        this.model = builder.model;
+        this.model = adapt(builder);
         DynamicClient.Builder clientBuilder = DynamicClient.builder()
                 .service(builder.service)
                 .model(model);