Add annotation to specify custom base images for generated Dockerfiles

src/Aspire.Hosting.NodeJs/NodeExtensions.cs

@@ -126,11 +126,14 @@ public static IResourceBuilder<NodeAppResource> AddNodeApp(this IDistributedAppl
 
                 c.WithDockerfileBuilder(resource.WorkingDirectory, dockerfileContext =>
                 {
-                    var logger = dockerfileContext.Services.GetService<ILogger<NodeAppResource>>() ?? NullLogger<NodeAppResource>.Instance;
-                    var nodeVersion = DetectNodeVersion(appDirectory, logger) ?? DefaultNodeVersion;
+                    var defaultBaseImage = new Lazy<string>(() => GetDefaultBaseImage(appDirectory, "alpine", dockerfileContext.Services));
 
+                    // Get custom base image from annotation, if present
+                    dockerfileContext.Resource.TryGetLastAnnotation<DockerfileBaseImageAnnotation>(out var baseImageAnnotation);
+
+                    var baseBuildImage = baseImageAnnotation?.BuildImage ?? defaultBaseImage.Value;
                     var builderStage = dockerfileContext.Builder
-                        .From($"node:{nodeVersion}-alpine", "build")
+                        .From(baseBuildImage, "build")
                         .EmptyLine()
                         .WorkDir("/app")
                         .Copy(".", ".")
@@ -157,8 +160,9 @@ public static IResourceBuilder<NodeAppResource> AddNodeApp(this IDistributedAppl
                         }
                     }
 
+                    var baseRuntimeImage = baseImageAnnotation?.RuntimeImage ?? defaultBaseImage.Value;
                     var runtimeBuilder = dockerfileContext.Builder
-                        .From($"node:{nodeVersion}-alpine", "runtime")
+                        .From(baseRuntimeImage, "runtime")
                             .EmptyLine()
                             .WorkDir("/app")
                             .CopyFrom("build", "/app", "/app")
@@ -313,10 +317,12 @@ private static IResourceBuilder<TResource> CreateDefaultJavaScriptAppBuilder<TRe
                 {
                     if (c.Resource.TryGetLastAnnotation<JavaScriptPackageManagerAnnotation>(out var packageManager))
                     {
-                        var logger = dockerfileContext.Services.GetService<ILogger<JavaScriptAppResource>>() ?? NullLogger<JavaScriptAppResource>.Instance;
-                        var nodeVersion = DetectNodeVersion(appDirectory, logger) ?? DefaultNodeVersion;
+                        // Get custom base image from annotation, if present
+                        dockerfileContext.Resource.TryGetLastAnnotation<DockerfileBaseImageAnnotation>(out var baseImageAnnotation);
+                        var baseImage = baseImageAnnotation?.BuildImage ?? GetDefaultBaseImage(appDirectory, "slim", dockerfileContext.Services);
+
                         var dockerBuilder = dockerfileContext.Builder
-                            .From($"node:{nodeVersion}-slim")
+                            .From(baseImage)
                             .WorkDir("/app")
                             .Copy(".", ".");
 
@@ -603,6 +609,13 @@ private static void AddInstaller<TResource>(IResourceBuilder<TResource> resource
         }
     }
 
+    private static string GetDefaultBaseImage(string appDirectory, string defaultSuffix, IServiceProvider serviceProvider)
+    {
+        var logger = serviceProvider.GetService<ILogger<JavaScriptAppResource>>() ?? NullLogger<JavaScriptAppResource>.Instance;
+        var nodeVersion = DetectNodeVersion(appDirectory, logger) ?? DefaultNodeVersion;
+        return $"node:{nodeVersion}-{defaultSuffix}";
+    }
+
     /// <summary>
     /// Detects the Node.js version to use for a project by checking common configuration files.
     /// </summary>

src/Aspire.Hosting.Python/PythonAppResourceBuilderExtensions.cs

@@ -469,8 +469,13 @@ private static IResourceBuilder<T> AddPythonAppCore<T>(
                     var uvLockPath = Path.Combine(resource.WorkingDirectory, "uv.lock");
                     var hasUvLock = File.Exists(uvLockPath);
 
+                    // Get custom base images from annotation, if present
+                    context.Resource.TryGetLastAnnotation<DockerfileBaseImageAnnotation>(out var baseImageAnnotation);
+                    var buildImage = baseImageAnnotation?.BuildImage ?? $"ghcr.io/astral-sh/uv:python{pythonVersion}-bookworm-slim";
+                    var runtimeImage = baseImageAnnotation?.RuntimeImage ?? $"python:{pythonVersion}-slim-bookworm";
+
                     var builderStage = context.Builder
-                        .From($"ghcr.io/astral-sh/uv:python{pythonVersion}-bookworm-slim", "builder")
+                        .From(buildImage, "builder")
                         .EmptyLine()
                         .Comment("Enable bytecode compilation and copy mode for the virtual environment")
                         .Env("UV_COMPILE_BYTECODE", "1")
@@ -518,7 +523,7 @@ private static IResourceBuilder<T> AddPythonAppCore<T>(
                     }
 
                     var runtimeBuilder = context.Builder
-                        .From($"python:{pythonVersion}-slim-bookworm", "app")
+                        .From(runtimeImage, "app")
                         .EmptyLine()
                         .AddContainerFiles(context.Resource, "/app")
                         .Comment("------------------------------")

src/Aspire.Hosting/ApplicationModel/DockerfileBaseImageAnnotation.cs

@@ -0,0 +1,37 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Diagnostics.CodeAnalysis;
+
+namespace Aspire.Hosting.ApplicationModel;
+
+/// <summary>
+/// Represents an annotation for specifying custom base images in generated Dockerfiles.
+/// </summary>
+/// <remarks>
+/// This annotation allows developers to override the default base images used when generating
+/// Dockerfiles for resources. It supports specifying separate build-time and runtime base images
+/// for multi-stage builds.
+/// </remarks>
+[Experimental("ASPIREDOCKERFILEBUILDER001", UrlFormat = "https://aka.ms/aspire/diagnostics/{0}")]
+public class DockerfileBaseImageAnnotation : IResourceAnnotation
+{
+    /// <summary>
+    /// Gets or sets the base image to use for the build stage in multi-stage Dockerfiles.
+    /// </summary>
+    /// <remarks>
+    /// This image is used during the build phase where dependencies are installed and
+    /// the application is compiled or prepared. If not specified, the default build image
+    /// for the resource type will be used.
+    /// </remarks>
+    public string? BuildImage { get; set; }
+
+    /// <summary>
+    /// Gets or sets the base image to use for the runtime stage in multi-stage Dockerfiles.
+    /// </summary>
+    /// <remarks>
+    /// This image is used for the final runtime stage where the application actually runs.
+    /// If not specified, the default runtime image for the resource type will be used.
+    /// </remarks>
+    public string? RuntimeImage { get; set; }
+}

src/Aspire.Hosting/ContainerResourceBuilderExtensions.cs

@@ -1427,6 +1427,50 @@ public static IResourceBuilder<T> WithDockerfileBuilder<T>(this IResourceBuilder
             return Task.CompletedTask;
         }, stage);
     }
+
+    /// <summary>
+    /// Configures custom base images for generated Dockerfiles.
+    /// </summary>
+    /// <typeparam name="T">The type of resource.</typeparam>
+    /// <param name="builder">The resource builder.</param>
+    /// <param name="buildImage">The base image to use for the build stage. If null, uses the default build image.</param>
+    /// <param name="runtimeImage">The base image to use for the runtime stage. If null, uses the default runtime image.</param>
+    /// <returns>The <see cref="IResourceBuilder{T}"/>.</returns>
+    /// <remarks>
+    /// <para>
+    /// This extension method allows customization of the base images used in generated Dockerfiles.
+    /// For multi-stage Dockerfiles (e.g., Python with UV), you can specify separate build and runtime images.
+    /// </para>
+    /// <example>
+    /// Specify custom base images for a Python application:
+    /// <code language="csharp">
+    /// var builder = DistributedApplication.CreateBuilder(args);
+    ///
+    /// builder.AddPythonApp("myapp", "path/to/app", "main.py")
+    ///        .WithDockerfileBaseImage(
+    ///            buildImage: "ghcr.io/astral-sh/uv:python3.12-bookworm-slim",
+    ///            runtimeImage: "python:3.12-slim-bookworm");
+    ///
+    /// builder.Build().Run();
+    /// </code>
+    /// </example>
+    /// </remarks>
+    [Experimental("ASPIREDOCKERFILEBUILDER001", UrlFormat = "https://aka.ms/aspire/diagnostics/{0}")]
+    public static IResourceBuilder<T> WithDockerfileBaseImage<T>(this IResourceBuilder<T> builder, string? buildImage = null, string? runtimeImage = null) where T : IResource
+    {
+        ArgumentNullException.ThrowIfNull(builder);
+
+        if (buildImage is null && runtimeImage is null)
+        {
+            throw new ArgumentException($"At least one of {nameof(buildImage)} or {nameof(runtimeImage)} must be specified.", nameof(buildImage));
+        }
+
+        return builder.WithAnnotation(new DockerfileBaseImageAnnotation
+        {
+            BuildImage = buildImage,
+            RuntimeImage = runtimeImage
+        }, ResourceAnnotationMutationBehavior.Replace);
+    }
 }
 
 internal static class IListExtensions

tests/Aspire.Hosting.NodeJs.Tests/AddNodeAppTests.cs

@@ -1,6 +1,8 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+#pragma warning disable ASPIREDOCKERFILEBUILDER001 // Type is for evaluation purposes only
+
 using Aspire.Hosting.ApplicationModel;
 using Aspire.Hosting.Tests.Utils;
 using Aspire.Hosting.Utils;
@@ -174,6 +176,30 @@ USER node
         Assert.Equal(expectedDockerfile, dockerfileContents);
     }
 
+    [Fact]
+    public async Task VerifyDockerfileWithCustomBaseImage()
+    {
+        using var tempDir = new TempDirectory();
+        using var builder = TestDistributedApplicationBuilder.Create(DistributedApplicationOperation.Publish, outputPath: tempDir.Path).WithResourceCleanUp(true);
+
+        var appDir = Path.Combine(tempDir.Path, "js");
+        Directory.CreateDirectory(appDir);
+        File.WriteAllText(Path.Combine(appDir, "package.json"), "{}");
+
+        var customBuildImage = "node:22-mySpecialBuildImage";
+        var customRuntimeImage = "node:22-mySpecialRuntimeImage";
+        var nodeApp = builder.AddNodeApp("js", appDir, "app.js")
+            .WithNpm(install: true)
+            .WithDockerfileBaseImage(customBuildImage, customRuntimeImage);
+
+        await ManifestUtils.GetManifest(nodeApp.Resource, tempDir.Path);
+
+        // Verify the Dockerfile contains the custom base image
+        var dockerfileContents = File.ReadAllText(Path.Combine(tempDir.Path, "js.Dockerfile"));
+        Assert.Contains($"FROM {customBuildImage}", dockerfileContents);
+        Assert.Contains($"FROM {customRuntimeImage}", dockerfileContents);
+    }
+
     [Fact]
     public void AddNodeApp_DoesNotAddNpmWhenNoPackageJson()
     {

tests/Aspire.Hosting.NodeJs.Tests/AddViteAppTests.cs

@@ -1,6 +1,8 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+#pragma warning disable ASPIREDOCKERFILEBUILDER001 // Type is for evaluation purposes only
+
 using Aspire.Hosting.ApplicationModel;
 using Aspire.Hosting.Utils;
 
@@ -201,4 +203,25 @@ public async Task VerifyDockerfileHandlesVariousVersionFormats(string versionStr
 
         Assert.Contains($"FROM {expectedImage}", dockerfileContents);
     }
+
+    [Fact]
+    public async Task VerifyCustomBaseImage()
+    {
+        using var tempDir = new TempDirectory();
+        using var builder = TestDistributedApplicationBuilder.Create(DistributedApplicationOperation.Publish, outputPath: tempDir.Path).WithResourceCleanUp(true);
+
+        var customImage = "node:22-myspecialimage";
+        var nodeApp = builder.AddViteApp("vite", tempDir.Path)
+            .WithNpm(install: true)
+            .WithDockerfileBaseImage(buildImage: customImage);
+
+        var manifest = await ManifestUtils.GetManifest(nodeApp.Resource, tempDir.Path);
+
+        // Verify the manifest structure
+        Assert.Equal("container.v1", manifest["type"]?.ToString());
+
+        // Verify the Dockerfile contains the custom base image
+        var dockerfileContents = File.ReadAllText(Path.Combine(tempDir.Path, "vite.Dockerfile"));
+        Assert.Contains($"FROM {customImage}", dockerfileContents);
+    }
 }

tests/Aspire.Hosting.Python.Tests/AddPythonAppTests.cs

@@ -3,6 +3,7 @@
 
 #pragma warning disable CS0612
 #pragma warning disable CS0618 // Type or member is obsolete
+#pragma warning disable ASPIREDOCKERFILEBUILDER001 // Type is for evaluation purposes only
 
 using Microsoft.Extensions.DependencyInjection;
 using Aspire.Hosting.Utils;
@@ -1313,5 +1314,63 @@ public async Task PythonApp_DoesNotSetPythonUtf8EnvironmentVariable_InPublishMod
         // PYTHONUTF8 should not be set in Publish mode, even on Windows
         Assert.False(environmentVariables.ContainsKey("PYTHONUTF8"));
     }
+
+    [Fact]
+    public async Task WithUvEnvironment_CustomBaseImages_GeneratesDockerfileWithCustomImages()
+    {
+        using var sourceDir = new TempDirectory();
+        using var outputDir = new TempDirectory();
+        var projectDirectory = sourceDir.Path;
+
+        // Create a UV-based Python project with pyproject.toml and uv.lock
+        var pyprojectContent = """
+            [project]
+            name = "test-app"
+            version = "0.1.0"
+            requires-python = ">=3.12"
+            dependencies = []
+
+            [build-system]
+            requires = ["hatchling"]
+            build-backend = "hatchling.build"
+            """;
+
+        var uvLockContent = """
+            version = 1
+            requires-python = ">=3.12"
+            """;
+
+        var scriptContent = """
+            print("Hello from UV project with custom images!")
+            """;
+
+        File.WriteAllText(Path.Combine(projectDirectory, "pyproject.toml"), pyprojectContent);
+        File.WriteAllText(Path.Combine(projectDirectory, "uv.lock"), uvLockContent);
+        File.WriteAllText(Path.Combine(projectDirectory, "main.py"), scriptContent);
+
+        using var builder = TestDistributedApplicationBuilder.Create(DistributedApplicationOperation.Publish, outputDir.Path, step: "publish-manifest");
+
+        // Add Python resource with custom base images
+        builder.AddPythonScript("custom-images-app", projectDirectory, "main.py")
+            .WithUvEnvironment()
+            .WithDockerfileBaseImage(
+                buildImage: "ghcr.io/astral-sh/uv:python3.13-bookworm",
+                runtimeImage: "python:3.13-slim");
+
+        var app = builder.Build();
+        app.Run();
+
+        // Verify that Dockerfile was generated
+        var dockerfilePath = Path.Combine(outputDir.Path, "custom-images-app.Dockerfile");
+        Assert.True(File.Exists(dockerfilePath), "Dockerfile should be generated");
+
+        var dockerfileContent = File.ReadAllText(dockerfilePath);
+
+        // Verify the custom build image is used
+        Assert.Contains("FROM ghcr.io/astral-sh/uv:python3.13-bookworm AS builder", dockerfileContent);
+
+        // Verify the custom runtime image is used
+        Assert.Contains("FROM python:3.13-slim AS app", dockerfileContent);
+    }
 }