address #1351 - waitfor between script and runspace not respected

src/CommunityToolkit.Aspire.Hosting.PowerShell/DistributedApplicationBuilderExtensions.cs

@@ -4,7 +4,6 @@
 using Aspire.Hosting.ApplicationModel;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Hosting;
-using System.Diagnostics;
 using System.Management.Automation;
 using System.Management.Automation.Runspaces;
 
@@ -16,17 +15,17 @@ namespace CommunityToolkit.Aspire.Hosting.PowerShell;
 public static class DistributedApplicationBuilderExtensions
 {
     /// <summary>
-    /// Adds a PowerShell runspace pool resource to the distributed application.
+    /// Adds a PowerShell runspace pool resource to the distributed application, enabling managed execution of
+    /// PowerShell scripts with configurable language mode and runspace limits.
     /// </summary>
-    /// <param name="builder"></param>
-    /// <param name="name"></param>
-    /// <param name="languageMode"></param>
-    /// <param name="minRunspaces"></param>
-    /// <param name="maxRunspaces"></param>
-    /// <returns></returns>
-    /// <remarks>This overload is not available in polyglot app hosts. Use the string-based overload instead.</remarks>
-    /// <exception cref="ArgumentException"></exception>
-    /// <exception cref="DistributedApplicationException"></exception>
+    /// <remarks>This overload is not ATS-compatible due to the use of PSLanguageMode. For ATS scenarios, use
+    /// the string-based overload instead.</remarks>
+    /// <param name="builder">The distributed application builder to which the PowerShell runspace pool resource will be added.</param>
+    /// <param name="name">The name of the PowerShell runspace pool resource. Cannot be null or whitespace.</param>
+    /// <param name="languageMode">The language mode to use for the PowerShell runspace pool. Defaults to PSLanguageMode.ConstrainedLanguage.</param>
+    /// <param name="minRunspaces">The minimum number of runspaces to maintain in the pool. Must be at least 1.</param>
+    /// <param name="maxRunspaces">The maximum number of runspaces allowed in the pool. Must be greater than or equal to minRunspaces.</param>
+    /// <returns>An IResourceBuilder instance for further configuration of the PowerShell runspace pool resource.</returns>
     [AspireExportIgnore(Reason = "PSLanguageMode is not ATS-compatible. Use the string-based overload instead.")]
     public static IResourceBuilder<PowerShellRunspacePoolResource> AddPowerShell(
         this IDistributedApplicationBuilder builder,
@@ -62,6 +61,15 @@ public static IResourceBuilder<PowerShellRunspacePoolResource> AddPowerShell(
             var sessionState = InitialSessionState.CreateDefault();
             sessionState.UseFullLanguageModeInDebugger = true;
 
+            await notificationService.PublishUpdateAsync(res,
+                state => state with
+                {
+                    State = KnownResourceStates.Starting,
+                    Properties = [
+                        .. state.Properties,
+                    ],
+                });
+
             // This will block until explicit and implied WaitFor calls are completed
             await builder.Eventing.PublishAsync(
                 new BeforeResourceStartedEvent(res, e.Services), ct);
@@ -81,7 +89,8 @@ await builder.Eventing.PublishAsync(
             var poolName = res.Name;
             var poolLogger = loggerService.GetLogger(poolName);
 
-            _ = res.StartAsync(sessionState, notificationService, poolLogger, hostLifetime, ct);
+            // The runspace pool should open rather quickly, so it's ok to await here.
+            await res.StartAsync(sessionState, notificationService, poolLogger, hostLifetime, ct);
         });
 
         return poolBuilder;

src/CommunityToolkit.Aspire.Hosting.PowerShell/PowerShellRunspacePoolResource.cs

@@ -85,21 +85,28 @@ private void ConfigureStateChangeNotifications(ResourceNotificationService notif
                     nameof(poolState), poolState, $"Unexpected runspace pool state {poolState}")
             };
 
-            logger.LogDebug(
-                "Runspace pool '{PoolName}' state mapped to known state '{KnownState}'", Name, knownState);
-
-            await notificationService.PublishUpdateAsync(this,
-                state => state with
-                {
-                    State = knownState,
-                    Properties = [
-                        .. state.Properties,
+            // only publish the update if the state is not NotStarted,
+            // since that's already the initial state. The "Starting" state
+            // is published in OnInitializeResource to allow WaitFor to work.
+            if (knownState != KnownResourceStates.NotStarted &&
+                knownState != KnownResourceStates.Starting)
+            {
+                logger.LogDebug(
+                    "Runspace pool '{PoolName}' state mapped to known state '{KnownState}'", Name, knownState);
+
+                await notificationService.PublishUpdateAsync(this,
+                    state => state with
+                    {
+                        State = knownState,
+                        Properties = [
+                            .. state.Properties,
                         new("RunspacePoolState", poolState.ToString()),
                         new("Reason", reason?.ToString() ?? string.Empty)
-                    ],
-                    StartTimeStamp = knownState == KnownResourceStates.Running ? DateTime.Now : state.StartTimeStamp,
-                    StopTimeStamp = KnownResourceStates.TerminalStates.Contains(knownState) ? DateTime.Now : state.StopTimeStamp,
-                });
+                        ],
+                        StartTimeStamp = knownState == KnownResourceStates.Running ? DateTime.Now : state.StartTimeStamp,
+                        StopTimeStamp = KnownResourceStates.TerminalStates.Contains(knownState) ? DateTime.Now : state.StopTimeStamp,
+                    });
+            }
         };
     }
 

src/CommunityToolkit.Aspire.Hosting.PowerShell/PowerShellRunspacePoolResourceBuilderExtensions.cs

@@ -78,12 +78,23 @@ public static IResourceBuilder<PowerShellScriptResource> AddScript(
 
             try
             {
+                await notificationService.PublishUpdateAsync(res,
+                    state => state with
+                    {
+                        State = KnownResourceStates.Starting,
+                        Properties = [
+                            .. state.Properties,
+                        ],
+                    });
+
                 // this will block until the runspace pool is started, which is implied by the WaitFor call
                 await builder.ApplicationBuilder.Eventing.PublishAsync(
                     new BeforeResourceStartedEvent(res, e.Services), ct);
 
                 scriptLogger.LogInformation("Starting script '{ScriptName}'", scriptName);
 
+                // we don't want to block initialization until the script completes;
+                // it's better to fire and forget here.
                 _ = res.StartAsync(scriptLogger, notificationService, ct);
             }
             catch (Exception ex)

src/CommunityToolkit.Aspire.Hosting.PowerShell/PowerShellScriptResource.cs

@@ -114,13 +114,20 @@ public async Task<bool> BreakAsync()
         }
 
         /// <summary>
-        /// Starts the PowerShell script execution.
+        /// Starts the execution of the PowerShell script asynchronously, publishing state updates and handling script
+        /// arguments as needed.
         /// </summary>
-        /// <param name="scriptLogger"></param>
-        /// <param name="notificationService"></param>
-        /// <param name="cancellationToken"></param>
-        /// <returns></returns>
-        /// <exception cref="ArgumentOutOfRangeException"></exception>
+        /// <remarks>State changes during script execution are published using the provided notification
+        /// service. Script arguments are resolved and added prior to invocation. If the script pipeline is stopped
+        /// intentionally, the resulting exception is ignored. If awaited, this method will complete when the script execution finishes, either successfully,
+        /// with an error, or by being stopped. The method also sets up event handlers to log output and errors from the script
+        /// execution, and to publish state updates using the provided notification service.
+        /// </remarks>
+        /// <param name="scriptLogger">The logger used to record informational and error messages related to the script execution. Cannot be null.</param>
+        /// <param name="notificationService">The service used to publish resource state updates during script execution. Cannot be null.</param>
+        /// <param name="cancellationToken">A cancellation token that can be used to cancel the script execution.</param>
+        /// <returns>A task that represents the asynchronous operation of starting and monitoring the PowerShell script.</returns>
+        /// <exception cref="ArgumentOutOfRangeException">Thrown if the PowerShell invocation state is not recognized when handling state changes.</exception>
         public async Task StartAsync(ILogger scriptLogger,
             ResourceNotificationService notificationService,
             CancellationToken cancellationToken = default)
@@ -151,20 +158,24 @@ public async Task StartAsync(ILogger scriptLogger,
                         "Unknown PowerShell invocation state")
                 };
 
-                scriptLogger.LogDebug("Publishing script {ScriptName} state as known state: {ScriptState}", Name, knownState);
+                // Only publish state updates for known states that are not the initial state.
+                if (knownState != KnownResourceStates.NotStarted)
+                {
+                    scriptLogger.LogDebug("Publishing script {ScriptName} state as known state: {ScriptState}", Name, knownState);
 
-                await notificationService.PublishUpdateAsync(this,
-                    state => state with
-                    {
-                        State = knownState,
-                        Properties = [
-                                .. state.Properties,
+                    await notificationService.PublishUpdateAsync(this,
+                        state => state with
+                        {
+                            State = knownState,
+                            Properties = [
+                                    .. state.Properties,
                                 new( "PSInvocationState", args.InvocationStateInfo.State.ToString() ),
                                 new( "Reason", args.InvocationStateInfo.Reason?.Message ?? string.Empty ),
-                            ],
-                        StartTimeStamp = knownState == KnownResourceStates.Running ? DateTime.Now : state.StartTimeStamp,
-                        StopTimeStamp = KnownResourceStates.TerminalStates.Contains(knownState) ? DateTime.Now : state.StopTimeStamp,
-                    });
+                                ],
+                            StartTimeStamp = knownState == KnownResourceStates.Running ? DateTime.Now : state.StartTimeStamp,
+                            StopTimeStamp = KnownResourceStates.TerminalStates.Contains(knownState) ? DateTime.Now : state.StopTimeStamp,
+                        });
+                }
             };
 
             if (this.TryGetLastAnnotation<PowerShellScriptArgsAnnotation>(out var scriptArgsAnnotation))
@@ -194,8 +205,7 @@ await notificationService.PublishUpdateAsync(this,
             }
             catch (Exception ex)
             {
-                scriptLogger.LogError(ex, "Error invoking PowerShell script: {Message}", ex.Message);
-                throw;
+                scriptLogger.LogError(ex, "Error invoking PowerShell script: {Message}", ex.Message);                
             }
         }
 
@@ -253,7 +263,7 @@ void IDisposable.Dispose()
             _isDisposed = true;
             _ps.Stop();
             _ps.Dispose();
-            _cts?.Dispose();
+            _cts.Dispose();
         }
     }
 }