Unify ATS withEnvironment exports

playground/polyglot/Java/Aspire.Hosting.Azure/ValidationAppHost/AppHost.java

@@ -73,11 +73,11 @@ void main() throws Exception {
         identity.publishAsExisting("identity-existing", "rg-identity");
         identity.publishAsExistingFromParameters(existingName, existingResourceGroup);
         identity.asExisting(existingName, existingResourceGroup);
-        container.withEnvironmentFromOutput("INFRA_URL", infrastructureOutput);
-        container.withEnvironmentFromKeyVaultSecret("SECRET_FROM_IDENTITY", identity);
+        container.withEnvironment("INFRA_URL", infrastructureOutput);
+        container.withEnvironment("SECRET_FROM_IDENTITY", identity);
         container.withAzureUserAssignedIdentity(identity);
-        executable.withEnvironmentFromOutput("INFRA_URL", infrastructureOutput);
-        executable.withEnvironmentFromKeyVaultSecret("SECRET_FROM_IDENTITY", identity);
+        executable.withEnvironment("INFRA_URL", infrastructureOutput);
+        executable.withEnvironment("SECRET_FROM_IDENTITY", identity);
         executable.withAzureUserAssignedIdentity(identity);
         builder.build().run();
     }

playground/polyglot/Java/Aspire.Hosting/ValidationAppHost/AppHost.java

@@ -19,9 +19,9 @@ void main() throws Exception {
         var builtConnectionString = builder.addConnectionStringBuilder("customcs", (connectionStringBuilder) -> { var _isEmpty = connectionStringBuilder.isEmpty(); connectionStringBuilder.appendLiteral("Host="); connectionStringBuilder.appendValueProvider(endpoint); connectionStringBuilder.appendLiteral(";Key="); connectionStringBuilder.appendValueProvider(secretParam); var _builtExpression = connectionStringBuilder.build(); });
         builtConnectionString.withConnectionProperty("Host", expr);
         builtConnectionString.withConnectionPropertyValue("Mode", "Development");
-        container.withEnvironmentEndpoint("MY_ENDPOINT", endpoint);
-        container.withEnvironmentParameter("MY_PARAM", configParam);
-        container.withEnvironmentConnectionString("MY_CONN", builtConnectionString);
+        container.withEnvironment("MY_ENDPOINT", endpoint);
+        container.withEnvironment("MY_PARAM", configParam);
+        container.withEnvironment("MY_CONN", builtConnectionString);
         builtConnectionString.withConnectionProperty("Endpoint", expr);
         builtConnectionString.withConnectionPropertyValue("Protocol", "https");
         container.excludeFromManifest();

playground/polyglot/Python/Aspire.Hosting.Azure/ValidationAppHost/apphost.py

@@ -72,10 +72,10 @@
     identity.publish_as_existing()
     identity.publish_as_existing_from_parameters()
     identity.as_existing()
-    container.with_environment_from_output()
-    container.with_environment_from_key_vault_secret()
+    container.with_environment("INFRA_URL", infrastructure_output)
+    container.with_environment("SECRET_FROM_IDENTITY", identity)
     container.with_azure_user_assigned_identity()
-    executable.with_environment_from_output()
-    executable.with_environment_from_key_vault_secret()
+    executable.with_environment("INFRA_URL", infrastructure_output)
+    executable.with_environment("SECRET_FROM_IDENTITY", identity)
     executable.with_azure_user_assigned_identity()
     builder.run()

playground/polyglot/Python/Aspire.Hosting/ValidationAppHost/apphost.py

@@ -33,12 +33,12 @@
     built_connection_string = builder.add_connection_string_builder("connection-string", lambda *_args, **_kwargs: None)
     built_connection_string.with_connection_property("Key", "Value")
     built_connection_string.with_connection_property_value("Key", "Value")
-    # withEnvironmentEndpoint
-    container.with_environment_endpoint("KEY", None)
-    # withEnvironmentParameter
-    container.with_environment_parameter("KEY", None)
-    # withEnvironmentConnectionString
-    container.with_environment_connection_string("KEY", None)
+    # withEnvironment - EndpointReference
+    container.with_environment("KEY", endpoint)
+    # withEnvironment - ParameterResource
+    container.with_environment("KEY", builder.add_parameter("param"))
+    # withEnvironment - connection string resource
+    container.with_environment("KEY", built_connection_string)
     # withConnectionProperty — with ReferenceExpression
     built_connection_string.with_connection_property("Key", "Value")
     # withConnectionPropertyValue — with string

playground/polyglot/TypeScript/Aspire.Hosting.Azure/ValidationAppHost/apphost.ts

@@ -89,12 +89,12 @@ await identity.publishAsExisting("identity-existing", "rg-identity");
 await identity.publishAsExistingFromParameters(existingName, existingResourceGroup);
 await identity.asExisting(existingName, existingResourceGroup);
 
-await container.withEnvironmentFromOutput("INFRA_URL", infrastructureOutput);
-await container.withEnvironmentFromKeyVaultSecret("SECRET_FROM_IDENTITY", identity);
+await container.withEnvironment("INFRA_URL", infrastructureOutput);
+await container.withEnvironment("SECRET_FROM_IDENTITY", identity);
 await container.withAzureUserAssignedIdentity(identity);
 
-await executable.withEnvironmentFromOutput("INFRA_URL", infrastructureOutput);
-await executable.withEnvironmentFromKeyVaultSecret("SECRET_FROM_IDENTITY", identity);
+await executable.withEnvironment("INFRA_URL", infrastructureOutput);
+await executable.withEnvironment("SECRET_FROM_IDENTITY", identity);
 await executable.withAzureUserAssignedIdentity(identity);
 
 await builder.build().run();

playground/polyglot/TypeScript/Aspire.Hosting/ValidationAppHost/apphost.ts

@@ -73,14 +73,14 @@ await builtConnectionString.withConnectionPropertyValue("Mode", "Development");
 // ResourceBuilderExtensions.cs — NEW exports on ContainerResource
 // ===================================================================
 
-// withEnvironmentEndpoint
-await container.withEnvironmentEndpoint("MY_ENDPOINT", endpoint);
+// withEnvironment - EndpointReference
+await container.withEnvironment("MY_ENDPOINT", endpoint);
 
-// withEnvironmentParameter
-await container.withEnvironmentParameter("MY_PARAM", configParam);
+// withEnvironment - ParameterResource
+await container.withEnvironment("MY_PARAM", configParam);
 
-// withEnvironmentConnectionString
-await container.withEnvironmentConnectionString("MY_CONN", builtConnectionString);
+// withEnvironment - connection string resource
+await container.withEnvironment("MY_CONN", builtConnectionString);
 
 // withConnectionProperty — with ReferenceExpression
 await builtConnectionString.withConnectionProperty("Endpoint", expr);

src/Aspire.Hosting.Azure.KeyVault/AzureKeyVaultSecretReference.cs

@@ -10,7 +10,7 @@ namespace Aspire.Hosting.Azure;
 /// </summary>
 /// <param name="secretName">The name of the secret.</param>
 /// <param name="azureKeyVaultResource">The Azure Key Vault resource.</param>
-internal sealed class AzureKeyVaultSecretReference(string secretName, AzureKeyVaultResource azureKeyVaultResource) : IAzureKeyVaultSecretReference, IValueProvider, IManifestExpressionProvider
+internal sealed class AzureKeyVaultSecretReference(string secretName, AzureKeyVaultResource azureKeyVaultResource) : IAzureKeyVaultSecretReference, IValueProvider, IManifestExpressionProvider, IValueWithCustomWithEnvironment<AzureKeyVaultSecretReference>
 {
     /// <summary>
     /// Gets the name of the secret.
@@ -38,4 +38,12 @@ internal sealed class AzureKeyVaultSecretReference(string secretName, AzureKeyVa
 
         throw new InvalidOperationException($"Secret '{secretName}' not found in Key Vault '{azureKeyVaultResource.Name}'.");
     }
+
+    static IResourceBuilder<TDestination>? IValueWithCustomWithEnvironment<AzureKeyVaultSecretReference>.TryWithEnvironment<TDestination>(
+        IResourceBuilder<TDestination> builder,
+        string name,
+        AzureKeyVaultSecretReference value)
+    {
+        return ResourceBuilderExtensions.WithEnvironment<TDestination, AzureKeyVaultSecretReference>(builder, name, value);
+    }
 }

src/Aspire.Hosting.Azure.KeyVault/AzureKeyVaultSecretResource.cs

@@ -15,7 +15,7 @@ namespace Aspire.Hosting.Azure;
 /// </remarks>
 [DebuggerDisplay("Type = {GetType().Name,nq}, Name = {Name}, Secret = {SecretName}")]
 public class AzureKeyVaultSecretResource(string name, string secretName, AzureKeyVaultResource parent, object value)
-    : Resource(name), IResourceWithParent<AzureKeyVaultResource>, IAzureKeyVaultSecretReference
+    : Resource(name), IResourceWithParent<AzureKeyVaultResource>, IAzureKeyVaultSecretReference, IValueWithCustomWithEnvironment<AzureKeyVaultSecretResource>
 {
     private readonly IAzureKeyVaultSecretReference _secret = parent.GetSecret(secretName);
 
@@ -53,4 +53,12 @@ public class AzureKeyVaultSecretResource(string name, string secretName, AzureKe
     /// <returns>The secret value.</returns>
     ValueTask<string?> IValueProvider.GetValueAsync(CancellationToken cancellationToken) =>
         _secret.GetValueAsync(cancellationToken);
+
+    static IResourceBuilder<TDestination>? IValueWithCustomWithEnvironment<AzureKeyVaultSecretResource>.TryWithEnvironment<TDestination>(
+        IResourceBuilder<TDestination> builder,
+        string name,
+        AzureKeyVaultSecretResource value)
+    {
+        return ResourceBuilderExtensions.WithEnvironment<TDestination, AzureKeyVaultSecretResource>(builder, name, value);
+    }
 }

src/Aspire.Hosting.Azure/AzureBicepResource.cs

@@ -628,7 +628,7 @@ public string? Value
 /// <param name="name">The name of the output</param>
 /// <param name="resource">The <see cref="AzureBicepResource"/>.</param>
 [AspireExport(ExposeProperties = true)]
-public sealed class BicepOutputReference(string name, AzureBicepResource resource) : IManifestExpressionProvider, IValueProvider, IValueWithReferences, IEquatable<BicepOutputReference>
+public sealed class BicepOutputReference(string name, AzureBicepResource resource) : IManifestExpressionProvider, IValueProvider, IValueWithReferences, IValueWithCustomWithEnvironment<BicepOutputReference>, IEquatable<BicepOutputReference>
 {
     /// <summary>
     /// Name of the output.
@@ -686,4 +686,12 @@ other is not null &&
     /// <inheritdoc/>
     public override int GetHashCode() =>
         HashCode.Combine(Resource, Name);
+
+    static IResourceBuilder<TDestination>? IValueWithCustomWithEnvironment<BicepOutputReference>.TryWithEnvironment<TDestination>(
+        IResourceBuilder<TDestination> builder,
+        string name,
+        BicepOutputReference value)
+    {
+        return ResourceBuilderExtensions.WithEnvironment<TDestination, BicepOutputReference>(builder, name, value);
+    }
 }

src/Aspire.Hosting.Azure/AzureBicepResourceExtensions.cs

@@ -78,16 +78,11 @@ public static BicepSecretOutputReference GetSecretOutput(this IResourceBuilder<A
     /// <param name="name">The name of the environment variable.</param>
     /// <param name="bicepOutputReference">The reference to the bicep output.</param>
     /// <returns>An <see cref="IResourceBuilder{T}"/>.</returns>
-    [AspireExport("withEnvironmentFromOutput", Description = "Sets an environment variable from a Bicep output reference")]
+    [AspireExportIgnore(Reason = "Polyglot app hosts use the internal withEnvironment dispatcher export.")]
     public static IResourceBuilder<T> WithEnvironment<T>(this IResourceBuilder<T> builder, string name, BicepOutputReference bicepOutputReference)
         where T : IResourceWithEnvironment
     {
-        builder.WithReferenceRelationship(bicepOutputReference.Resource);
-
-        return builder.WithEnvironment(ctx =>
-        {
-            ctx.EnvironmentVariables[name] = bicepOutputReference;
-        });
+        return ResourceBuilderExtensions.WithEnvironment<T, BicepOutputReference>(builder, name, bicepOutputReference);
     }
 
     /// <summary>
@@ -116,14 +111,11 @@ public static IResourceBuilder<T> WithEnvironment<T>(this IResourceBuilder<T> bu
     /// <param name="name">The name of the environment variable.</param>
     /// <param name="secretReference">The reference to the key vault secret.</param>
     /// <returns>An <see cref="IResourceBuilder{T}"/>.</returns>
-    [AspireExport("withEnvironmentFromKeyVaultSecret", Description = "Sets an environment variable from an Azure Key Vault secret reference")]
+    [AspireExportIgnore(Reason = "Polyglot app hosts use the internal withEnvironment dispatcher export.")]
     public static IResourceBuilder<T> WithEnvironment<T>(this IResourceBuilder<T> builder, string name, IAzureKeyVaultSecretReference secretReference)
         where T : IResourceWithEnvironment
     {
-        return builder.WithEnvironment(ctx =>
-        {
-            ctx.EnvironmentVariables[name] = secretReference;
-        });
+        return ResourceBuilderExtensions.WithEnvironment<T, IAzureKeyVaultSecretReference>(builder, name, secretReference);
     }
 
     /// <summary>

src/Aspire.Hosting.Azure/IAzureKeyVaultSecretReference.cs

@@ -9,7 +9,7 @@ namespace Aspire.Hosting.Azure;
 /// Represents a reference to a secret in an Azure Key Vault resource.
 /// </summary>
 [AspireExport]
-public interface IAzureKeyVaultSecretReference : IValueProvider, IManifestExpressionProvider, IValueWithReferences
+public interface IAzureKeyVaultSecretReference : IEnvironmentValue, IValueWithReferences
 {
     /// <summary>
     /// Gets the name of the secret.

src/Aspire.Hosting.CodeGeneration.Java/AtsJavaCodeGenerator.cs

@@ -756,6 +756,11 @@ private List<JavaMethodParameter> CreateMethodParameters(IEnumerable<AtsParamete
         }
 
         var wrapperType = MapInputTypeToJava(parameter.Type, parameter.IsOptional || parameter.IsNullable);
+        return GetResourceBuilderWrapperType(wrapperType);
+    }
+
+    private (string? ResourceWrapperType, string? ResourceWrapperParameterType) GetResourceBuilderWrapperType(string wrapperType)
+    {
         if (!wrapperType.StartsWith("I", StringComparison.Ordinal))
         {
             return (null, null);
@@ -1001,6 +1006,12 @@ private void GenerateUnionOverloads(JavaCapabilityReturnInfo returnInfo, string
             }
             WriteLine("    }");
             WriteLine();
+
+            GenerateResourceBuilderOverloads(
+                returnInfo.ReturnType,
+                methodName,
+                CreateUnionMethodParameters(parameters, unionParameter, unionType),
+                returnInfo.HasReturn);
         }
     }
 
@@ -1112,6 +1123,14 @@ private void GenerateUnionOverloadsWithOptions(
             WriteLine("    }");
             WriteLine();
 
+            var bridgeParameters = CreateUnionMethodParameters(requiredParameters, unionParameter, unionType);
+            bridgeParameters.Add(new JavaMethodParameter(optionsClassName, "options"));
+            GenerateResourceBuilderOverloads(
+                returnInfo.ReturnType,
+                methodName,
+                bridgeParameters,
+                returnInfo.HasReturn);
+
             WriteLine($"    public {returnInfo.ReturnType} {methodName}({string.Join(", ", requiredParameters.Select(parameter => ReferenceEquals(parameter, unionParameter) ? $"{MapInputTypeToJava(unionType, unionParameter.IsOptional || unionParameter.IsNullable)} {ToCamelCase(parameter.Name)}" : $"{MapParameterToJava(parameter)} {ToCamelCase(parameter.Name)}"))}) {{");
             if (returnInfo.HasReturn)
             {
@@ -1126,6 +1145,40 @@ private void GenerateUnionOverloadsWithOptions(
         }
     }
 
+    private List<JavaMethodParameter> CreateUnionMethodParameters(
+        List<AtsParameterInfo> parameters,
+        AtsParameterInfo unionParameter,
+        AtsTypeRef unionType)
+    {
+        var result = new List<JavaMethodParameter>(parameters.Count);
+
+        foreach (var parameter in parameters)
+        {
+            var parameterName = ToCamelCase(parameter.Name);
+
+            if (!ReferenceEquals(parameter, unionParameter))
+            {
+                var (parameterResourceWrapperType, parameterResourceWrapperParameterType) = GetResourceBuilderWrapperType(parameter);
+                result.Add(new JavaMethodParameter(
+                    MapParameterToJava(parameter),
+                    parameterName,
+                    parameterResourceWrapperType,
+                    parameterResourceWrapperParameterType));
+                continue;
+            }
+
+            var parameterType = MapInputTypeToJava(unionType, unionParameter.IsOptional || unionParameter.IsNullable);
+            var (resourceWrapperType, resourceWrapperParameterType) = GetResourceBuilderWrapperType(parameterType);
+            result.Add(new JavaMethodParameter(
+                parameterType,
+                parameterName,
+                resourceWrapperType,
+                resourceWrapperParameterType));
+        }
+
+        return result;
+    }
+
     private void GenerateOptionsOverloads(
         AtsCapabilityInfo capability,
         JavaCapabilityReturnInfo returnInfo,