Fix event order and callback calling during container creation

src/Aspire.Hosting/ApplicationModel/ArgumentsExecutionConfigurationGatherer.cs

@@ -13,18 +13,24 @@ internal class ArgumentsExecutionConfigurationGatherer : IExecutionConfiguration
     /// <inheritdoc/>
     public async ValueTask GatherAsync(IExecutionConfigurationGathererContext context, IResource resource, ILogger resourceLogger, DistributedApplicationExecutionContext executionContext, CancellationToken cancellationToken = default)
     {
-        if (resource.TryGetAnnotationsOfType<CommandLineArgsCallbackAnnotation>(out var callbacks))
+        if (resource.TryGetAnnotationsOfType<CommandLineArgsCallbackAnnotation>(out var argumentAnnotations))
         {
-            var callbackContext = new CommandLineArgsCallbackContext(context.Arguments, resource, cancellationToken)
+            IList<object> args = [.. context.Arguments];
+            var callbackContext = new CommandLineArgsCallbackContext(args, resource, cancellationToken)
             {
                 Logger = resourceLogger,
                 ExecutionContext = executionContext
             };
 
-            foreach (var callback in callbacks)
+            foreach (var ann in argumentAnnotations)
             {
-                await callback.Callback(callbackContext).ConfigureAwait(false);
+                // Each annotation operates on a shared context.
+                args = await ann.AsCallbackAnnotation().EvaluateOnceAsync(callbackContext).ConfigureAwait(false);
             }
+
+            // Take the final result and apply to the gatherer context.
+            context.Arguments.Clear();
+            context.Arguments.AddRange(args);
         }
     }
 }

src/Aspire.Hosting/ApplicationModel/CommandLineArgsCallbackAnnotation.cs

@@ -1,16 +1,22 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Collections.Immutable;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Logging.Abstractions;
 
 namespace Aspire.Hosting.ApplicationModel;
 
+using IArgCallbackAnnotation = ICallbackResourceAnnotation<CommandLineArgsCallbackContext, IList<object>>;
+
 /// <summary>
 /// Represents an annotation that provides a callback to be executed with a list of command-line arguments when an executable resource is started.
 /// </summary>
-public class CommandLineArgsCallbackAnnotation : IResourceAnnotation
+public class CommandLineArgsCallbackAnnotation : IResourceAnnotation, IArgCallbackAnnotation
 {
+    private Task<IList<object>>? _callbackTask;
+    private readonly object _lock = new();
+
     /// <summary>
     /// Initializes a new instance of the <see cref="CommandLineArgsCallbackAnnotation"/> class with the specified callback action.
     /// </summary>
@@ -41,6 +47,35 @@ public CommandLineArgsCallbackAnnotation(Action<IList<object>> callback)
     /// Gets the callback action to be executed when the executable arguments are parsed.
     /// </summary>
     public Func<CommandLineArgsCallbackContext, Task> Callback { get; }
+
+    internal IArgCallbackAnnotation AsCallbackAnnotation() => this;
+
+    Task<IList<object>> IArgCallbackAnnotation.EvaluateOnceAsync(CommandLineArgsCallbackContext context)
+    {
+        lock(_lock)
+        {
+            if (_callbackTask is null)
+            {
+                _callbackTask = ExecuteCallbackAsync(context);
+            }
+            return _callbackTask;
+        }
+    }
+
+    void IArgCallbackAnnotation.ForgetCachedResult()
+    {
+        lock(_lock)
+        {
+            _callbackTask = null;
+        }
+    }
+
+    private async Task<IList<object>> ExecuteCallbackAsync(CommandLineArgsCallbackContext context)
+    {
+        await Callback(context).ConfigureAwait(false);
+        var result = context.Args.ToImmutableList();
+        return result;
+    }
 }
 
 /// <summary>

src/Aspire.Hosting/ApplicationModel/EnvironmentCallbackAnnotation.cs

@@ -5,13 +5,17 @@
 
 namespace Aspire.Hosting.ApplicationModel;
 
+using IEnvCallbackAnnotation = ICallbackResourceAnnotation<EnvironmentCallbackContext, Dictionary<string, object>>;
+
 /// <summary>
 /// Represents an annotation that provides a callback to modify the environment variables of an application.
 /// </summary>
 [DebuggerDisplay("{DebuggerToString(),nq}")]
-public class EnvironmentCallbackAnnotation : IResourceAnnotation
+public class EnvironmentCallbackAnnotation : IResourceAnnotation, IEnvCallbackAnnotation
 {
     private readonly string? _name;
+    private Task<Dictionary<string, object>>? _callbackTask;
+    private readonly object _lock = new();
 
     /// <summary>
     /// Initializes a new instance of the <see cref="EnvironmentCallbackAnnotation"/> class with the specified name and callback function.
@@ -77,6 +81,35 @@ public EnvironmentCallbackAnnotation(Func<EnvironmentCallbackContext, Task> call
     /// </summary>
     public Func<EnvironmentCallbackContext, Task> Callback { get; private set; }
 
+    internal IEnvCallbackAnnotation AsCallbackAnnotation() => this;
+
+    Task<Dictionary<string, object>> IEnvCallbackAnnotation.EvaluateOnceAsync(EnvironmentCallbackContext context)
+    {
+        lock(_lock)
+        {
+            if (_callbackTask is null)
+            {
+                _callbackTask = ExecuteCallbackAsync(context);
+            }
+            return _callbackTask;
+        }
+    }
+
+    void IEnvCallbackAnnotation.ForgetCachedResult()
+    {
+        lock(_lock)
+        {
+            _callbackTask = null;
+        }
+    }
+
+    private async Task<Dictionary<string, object>> ExecuteCallbackAsync(EnvironmentCallbackContext context)
+    {
+        await Callback(context).ConfigureAwait(false);
+        var result = new Dictionary<string, object>(context.EnvironmentVariables);
+        return result;
+    }
+
     private string DebuggerToString()
     {
         var text = $@"Type = {GetType().Name}";

src/Aspire.Hosting/ApplicationModel/EnvironmentVariablesConfigurationGatherer.cs

@@ -13,16 +13,24 @@ internal class EnvironmentVariablesExecutionConfigurationGatherer : IExecutionCo
     /// <inheritdoc/>
     public async ValueTask GatherAsync(IExecutionConfigurationGathererContext context, IResource resource, ILogger resourceLogger, DistributedApplicationExecutionContext executionContext, CancellationToken cancellationToken = default)
     {
-        if (resource.TryGetEnvironmentVariables(out var callbacks))
+        if (resource.TryGetEnvironmentVariables(out var envVarAnnotations))
         {
-            var callbackContext = new EnvironmentCallbackContext(executionContext, resource, context.EnvironmentVariables, cancellationToken)
+            var envVars = new Dictionary<string, object>(context.EnvironmentVariables);
+            var callbackContext = new EnvironmentCallbackContext(executionContext, resource, envVars, cancellationToken: cancellationToken)
             {
                 Logger = resourceLogger,
             };
 
-            foreach (var callback in callbacks)
+            foreach (var ann in envVarAnnotations)
             {
-                await callback.Callback(callbackContext).ConfigureAwait(false);
+                // Each annotation operates on a shared context.
+                envVars = await ann.AsCallbackAnnotation().EvaluateOnceAsync(callbackContext).ConfigureAwait(false);
+            }
+
+            // Take the final result and apply to the gatherer context.
+            foreach (var kvp in envVars)
+            {
+                context.EnvironmentVariables[kvp.Key] = kvp.Value;
             }
         }
     }

src/Aspire.Hosting/ApplicationModel/ICallbackResourceAnnotation.cs

@@ -0,0 +1,29 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace Aspire.Hosting.ApplicationModel;
+
+/// <summary>
+/// Represents a resource annotation whose callback should be evaluated at most once,
+/// with the result cached for subsequent retrievals.
+/// </summary>
+/// <typeparam name="TContext">The type of the context passed to the callback.</typeparam>
+/// <typeparam name="TResult">The type of the result produced by the callback.</typeparam>
+internal interface ICallbackResourceAnnotation<TContext, TResult>
+{
+    /// <summary>
+    /// Evaluates the callback if it has not been evaluated yet, caching the result.
+    /// Subsequent calls return the cached result regardless of the context passed.
+    /// </summary>
+    /// <param name="context">The context for the callback evaluation. Only used on the first call.</param>
+    /// <returns>The cached result of the callback evaluation.</returns>
+    Task<TResult> EvaluateOnceAsync(TContext context);
+
+    /// <summary>
+    /// Clears the cached result so that the next call to <see cref="EvaluateOnceAsync"/> will re-execute the callback. 
+    ///</summary>
+    /// <remarks>
+    /// Use <see cref="ForgetCachedResult"/> when a resource decorated with this callback annotation is restarted.
+    /// </remarks>
+    void ForgetCachedResult();
+}

src/Aspire.Hosting/ApplicationModel/ResourceDependencyDiscoveryMode.cs

@@ -6,19 +6,26 @@ namespace Aspire.Hosting.ApplicationModel;
 /// <summary>
 /// Specifies how resource dependencies are discovered.
 /// </summary>
+[Flags]
 public enum ResourceDependencyDiscoveryMode
 {
     /// <summary>
     /// Discover the full transitive closure of all dependencies.
     /// This includes direct dependencies and all dependencies of those dependencies, recursively.
     /// </summary>
-    Recursive,
+    Recursive = 1,
 
     /// <summary>
     /// Discover only direct dependencies.
     /// This includes dependencies from annotations (parent, wait, connection string redirect)
     /// and from environment variables and command-line arguments, but does not recurse
     /// into the dependencies of those dependencies.
     /// </summary>
-    DirectOnly
+    DirectOnly = 2,
+
+    /// <summary>
+    /// When set, unresolved values from annotation callbacks will be cached and reused 
+    /// on subsequent evaluations of the same annotation, rather than re-evaluating the callback each time.
+    /// </summary>
+    CacheAnnotationCallbackResults = 4
 }

src/Aspire.Hosting/ApplicationModel/ResourceExtensions.cs

@@ -1256,7 +1256,7 @@ internal static ILogger GetLogger(this IResource resource, IServiceProvider serv
     /// </summary>
     /// <param name="resource">The resource to compute dependencies for.</param>
     /// <param name="executionContext">The execution context for resolving environment variables and arguments.</param>
-    /// <param name="mode">Specifies whether to discover only direct dependencies or the full transitive closure.</param>
+    /// <param name="mode">Specifies dependency discovery mode.</param>
     /// <param name="cancellationToken">A cancellation token to observe while computing dependencies.</param>
     /// <returns>A set of all resources that the specified resource depends on.</returns>
     /// <remarks>
@@ -1292,7 +1292,7 @@ public static Task<IReadOnlySet<IResource>> GetResourceDependenciesAsync(
     /// </summary>
     /// <param name="resources">The source set of resources to compute dependencies for.</param>
     /// <param name="executionContext">The execution context for resolving environment variables and arguments.</param>
-    /// <param name="mode">Specifies whether to discover only direct dependencies or the full transitive closure.</param>
+    /// <param name="mode">Specifies dependency discovery mode.</param>
     /// <param name="cancellationToken">A cancellation token to observe while computing dependencies.</param>
     /// <returns>A set of all resources that the specified resource depends on.</returns>
     /// <remarks>
@@ -1320,20 +1320,26 @@ internal static async Task<IReadOnlySet<IResource>> GetDependenciesAsync(
         ResourceDependencyDiscoveryMode mode = ResourceDependencyDiscoveryMode.Recursive,
         CancellationToken cancellationToken = default)
     {
+        var modeOK = mode.HasFlag(ResourceDependencyDiscoveryMode.Recursive) ^ mode.HasFlag(ResourceDependencyDiscoveryMode.DirectOnly);
+        if (!modeOK)
+        {
+            throw new ArgumentException($"Exactly one of {nameof(ResourceDependencyDiscoveryMode.Recursive)} or {nameof(ResourceDependencyDiscoveryMode.DirectOnly)} must be set in the mode parameter.", nameof(mode));
+        }
+
         var dependencies = new HashSet<IResource>();
         var newDependencies = new HashSet<IResource>();
         var toProcess = new Queue<IResource>();
 
         foreach (var resource in resources)
         {
             newDependencies.Clear();
-            await GatherDirectDependenciesAsync(resource, dependencies, newDependencies, executionContext, cancellationToken).ConfigureAwait(false);
+            await GatherDirectDependenciesAsync(resource, dependencies, newDependencies, executionContext, mode, cancellationToken).ConfigureAwait(false);
 
-            if (mode == ResourceDependencyDiscoveryMode.Recursive)
+            if (mode.HasFlag(ResourceDependencyDiscoveryMode.Recursive))
             {
                 // Compute transitive closure by recursively processing dependencies
 
-                foreach(var nd in newDependencies)
+                foreach (var nd in newDependencies)
                 {
                     toProcess.Enqueue(nd);
                 }
@@ -1343,7 +1349,7 @@ internal static async Task<IReadOnlySet<IResource>> GetDependenciesAsync(
                     var dep = toProcess.Dequeue();
                     newDependencies.Clear();
 
-                    await GatherDirectDependenciesAsync(dep, dependencies, newDependencies, executionContext, cancellationToken).ConfigureAwait(false);
+                    await GatherDirectDependenciesAsync(dep, dependencies, newDependencies, executionContext, mode, cancellationToken).ConfigureAwait(false);
 
                     foreach (var newDep in newDependencies)
                     {
@@ -1372,12 +1378,14 @@ internal static async Task<IReadOnlySet<IResource>> GetDependenciesAsync(
     /// <param name="dependencies">The set of dependencies (where dependency resources will be placed).</param>
     /// <param name="newDependencies">The set of newly discovered dependencies in this invocation (not present in <paramref name="dependencies"/> at the moment of invocation).</param>
     /// <param name="executionContext">The execution context for resolving environment variables and arguments.</param>
+    /// <param name="mode">Specifies dependency discovery mode.</param>
     /// <param name="cancellationToken">A cancellation token to observe while gathering dependencies.</param>
     private static async Task GatherDirectDependenciesAsync(
         IResource resource,
         HashSet<IResource> dependencies,
         HashSet<IResource> newDependencies,
         DistributedApplicationExecutionContext executionContext,
+        ResourceDependencyDiscoveryMode mode,
         CancellationToken cancellationToken)
     {
         var visited = new HashSet<object>();
@@ -1386,7 +1394,7 @@ private static async Task GatherDirectDependenciesAsync(
         CollectAnnotationDependencies(resource, dependencies, newDependencies);
 
         // Collect raw (unresolved) environment variable and argument values
-        var rawValues = await GatherRawEnvironmentAndArgumentValuesAsync(resource, executionContext, cancellationToken).ConfigureAwait(false);
+        var rawValues = await GatherRawEnvironmentAndArgumentValuesAsync(resource, executionContext, mode, cancellationToken).ConfigureAwait(false);
 
         foreach (var value in rawValues)
         {
@@ -1400,39 +1408,61 @@ private static async Task GatherDirectDependenciesAsync(
     private static async Task<List<object>> GatherRawEnvironmentAndArgumentValuesAsync(
         IResource resource,
         DistributedApplicationExecutionContext executionContext,
+        ResourceDependencyDiscoveryMode mode,
         CancellationToken cancellationToken)
     {
         var rawValues = new List<object>();
 
         // Gather environment variable values
-        if (resource.TryGetEnvironmentVariables(out var environmentCallbacks))
+        if (resource.TryGetEnvironmentVariables(out var envAnnotations))
         {
             var envVars = new Dictionary<string, object>();
-            var context = new EnvironmentCallbackContext(executionContext, resource, envVars, cancellationToken);
-
-            foreach (var callback in environmentCallbacks)
+            var context = new EnvironmentCallbackContext(executionContext, resource, envVars, cancellationToken: cancellationToken);
+            
+            if (mode.HasFlag(ResourceDependencyDiscoveryMode.CacheAnnotationCallbackResults))
             {
-                await callback.Callback(context).ConfigureAwait(false);
+                foreach (var ann in envAnnotations)
+                {
+                    var resultingVars = await ann.AsCallbackAnnotation().EvaluateOnceAsync(context).ConfigureAwait(false);
+                    rawValues.AddRange(resultingVars.Values);
+                }
+
+            }
+            else
+            {
+                foreach (var ann in envAnnotations)
+                {
+                    await ann.Callback(context).ConfigureAwait(false);
+                }
+                rawValues.AddRange(envVars.Values);
             }
-
-            rawValues.AddRange(envVars.Values);
         }
 
         // Gather command-line argument values
-        if (resource.TryGetAnnotationsOfType<CommandLineArgsCallbackAnnotation>(out var argsCallbacks))
+        if (resource.TryGetAnnotationsOfType<CommandLineArgsCallbackAnnotation>(out var argAnnotations))
         {
             var args = new List<object>();
             var context = new CommandLineArgsCallbackContext(args, resource, cancellationToken)
             {
                 ExecutionContext = executionContext
             };
 
-            foreach (var callback in argsCallbacks)
+            if (mode.HasFlag(ResourceDependencyDiscoveryMode.CacheAnnotationCallbackResults))
             {
-                await callback.Callback(context).ConfigureAwait(false);
+                foreach (var ann in argAnnotations)
+                {
+                    var resultingArgs = await ann.AsCallbackAnnotation().EvaluateOnceAsync(context).ConfigureAwait(false);
+                    rawValues.AddRange(resultingArgs);
+                }
+            }
+            else
+            {
+                foreach (var ann in argAnnotations)
+                {
+                    await ann.Callback(context).ConfigureAwait(false);
+                }
+                rawValues.AddRange(args);
             }
-
-            rawValues.AddRange(args);
         }
 
         return rawValues;