Fix aspire wait for JSON-emitted resource names

src/Aspire.Hosting/ApplicationModel/ResourceNotificationService.cs

@@ -224,7 +224,7 @@ public async Task<ResourceEvent> WaitForResourceHealthyAsync(string resourceName
         _logger.LogDebug("Waiting for resource '{ResourceName}' to enter the '{State}' state.", resourceName, HealthStatus.Healthy);
         var resourceEvent = await WaitForResourceCoreAsync(
             resourceName,
-            re => ShouldYield(waitBehavior, re.Snapshot),
+            re => ShouldYieldHealthyWait(waitBehavior, re.Snapshot),
             $"Resource '{resourceName}' failed to become healthy before the operation was cancelled.",
             cancellationToken: cancellationToken).ConfigureAwait(false);
 
@@ -248,21 +248,20 @@ public async Task<ResourceEvent> WaitForResourceHealthyAsync(string resourceName
         _logger.LogDebug("Finished waiting for resource '{ResourceName}'.", resourceName);
 
         return resourceEvent;
-
-        // Determine if we should yield based on the wait behavior and the snapshot of the resource.
-        static bool ShouldYield(WaitBehavior waitBehavior, CustomResourceSnapshot snapshot) =>
-            waitBehavior switch
-            {
-                WaitBehavior.WaitOnResourceUnavailable => snapshot.HealthStatus == HealthStatus.Healthy,
-                WaitBehavior.StopOnResourceUnavailable => snapshot.HealthStatus == HealthStatus.Healthy ||
-                                                      snapshot.State?.Text == KnownResourceStates.Finished ||
-                                                      snapshot.State?.Text == KnownResourceStates.Exited ||
-                                                      snapshot.State?.Text == KnownResourceStates.FailedToStart ||
-                                                      snapshot.State?.Text == KnownResourceStates.RuntimeUnhealthy,
-                _ => throw new DistributedApplicationException($"Unexpected wait behavior: {waitBehavior}")
-            };
     }
 
+    internal static bool ShouldYieldHealthyWait(WaitBehavior waitBehavior, CustomResourceSnapshot snapshot) =>
+        waitBehavior switch
+        {
+            WaitBehavior.WaitOnResourceUnavailable => snapshot.HealthStatus == HealthStatus.Healthy,
+            WaitBehavior.StopOnResourceUnavailable => snapshot.HealthStatus == HealthStatus.Healthy ||
+                                                  snapshot.State?.Text == KnownResourceStates.Finished ||
+                                                  snapshot.State?.Text == KnownResourceStates.Exited ||
+                                                  snapshot.State?.Text == KnownResourceStates.FailedToStart ||
+                                                  snapshot.State?.Text == KnownResourceStates.RuntimeUnhealthy,
+            _ => throw new DistributedApplicationException($"Unexpected wait behavior: {waitBehavior}")
+        };
+
     private async Task WaitUntilCompletionAsync(IResource resource, IResource dependency, int exitCode, CancellationToken cancellationToken)
     {
         var names = dependency.GetResolvedResourceNames();

src/Aspire.Hosting/Backchannel/AuxiliaryBackchannelRpcTarget.cs

@@ -8,6 +8,7 @@
 using Aspire.Hosting.ApplicationModel;
 using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Diagnostics.HealthChecks;
 using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Logging;
 using ModelContextProtocol.Client;
@@ -247,29 +248,30 @@ public async Task<WaitForResourceResponse> WaitForResourceAsync(WaitForResourceR
     {
         ArgumentNullException.ThrowIfNull(request);
 
-        var appModel = serviceProvider.GetService<DistributedApplicationModel>();
-        if (appModel is not null && !appModel.Resources.Any(r => string.Equals(r.Name, request.ResourceName, StringComparisons.ResourceName)))
+        var notificationService = serviceProvider.GetRequiredService<ResourceNotificationService>();
+        var targetResolution = ResolveWaitTarget(notificationService, request.ResourceName);
+        var targetResource = targetResolution.Target;
+
+        if (targetResource is null)
         {
             return new WaitForResourceResponse
             {
                 Success = false,
-                ResourceNotFound = true,
-                ErrorMessage = $"Resource '{request.ResourceName}' was not found."
+                ResourceNotFound = targetResolution.ResourceNotFound,
+                ErrorMessage = targetResolution.ErrorMessage
             };
         }
 
-        var notificationService = serviceProvider.GetRequiredService<ResourceNotificationService>();
-
         using var timeoutCts = new CancellationTokenSource(TimeSpan.FromSeconds(request.TimeoutSeconds));
         using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutCts.Token);
 
         try
         {
             return request.Status switch
             {
-                "healthy" => await WaitForHealthyAsync(notificationService, request.ResourceName, linkedCts.Token).ConfigureAwait(false),
-                "up" => await WaitForRunningAsync(notificationService, request.ResourceName, linkedCts.Token).ConfigureAwait(false),
-                "down" => await WaitForTerminalAsync(notificationService, request.ResourceName, linkedCts.Token).ConfigureAwait(false),
+                "healthy" => await WaitForHealthyAsync(notificationService, targetResource, linkedCts.Token).ConfigureAwait(false),
+                "up" => await WaitForRunningAsync(notificationService, targetResource, linkedCts.Token).ConfigureAwait(false),
+                "down" => await WaitForTerminalAsync(notificationService, targetResource, linkedCts.Token).ConfigureAwait(false),
                 _ => new WaitForResourceResponse { Success = false, ErrorMessage = $"Unknown status: {request.Status}" }
             };
         }
@@ -283,9 +285,28 @@ public async Task<WaitForResourceResponse> WaitForResourceAsync(WaitForResourceR
         }
     }
 
-    private static async Task<WaitForResourceResponse> WaitForHealthyAsync(ResourceNotificationService notificationService, string resourceName, CancellationToken cancellationToken)
+    private static async Task<WaitForResourceResponse> WaitForHealthyAsync(ResourceNotificationService notificationService, WaitResourceTarget target, CancellationToken cancellationToken)
     {
-        var resourceEvent = await notificationService.WaitForResourceHealthyAsync(resourceName, WaitBehavior.StopOnResourceUnavailable, cancellationToken).ConfigureAwait(false);
+        var resourceEvent = await WaitForResourceEventAsync(
+            notificationService,
+            target,
+            re => ResourceNotificationService.ShouldYieldHealthyWait(WaitBehavior.StopOnResourceUnavailable, re.Snapshot),
+            $"Resource '{target.DisplayName}' failed to become healthy before the operation was cancelled.",
+            cancellationToken).ConfigureAwait(false);
+
+        if (resourceEvent.Snapshot.HealthStatus != HealthStatus.Healthy)
+        {
+            throw new DistributedApplicationException($"Stopped waiting for resource '{target.DisplayName}' to become healthy because it failed to start.");
+        }
+
+        resourceEvent = await WaitForResourceEventAsync(
+            notificationService,
+            new WaitResourceTarget(target.DisplayName, resourceEvent.ResourceId, null),
+            re => re.Snapshot.ResourceReadyEvent is not null,
+            $"Resource '{target.DisplayName}' failed to execute the resource ready event before the operation was cancelled.",
+            cancellationToken).ConfigureAwait(false);
+
+        await resourceEvent.Snapshot.ResourceReadyEvent!.EventTask.WaitAsync(cancellationToken).ConfigureAwait(false);
 
         return new WaitForResourceResponse
         {
@@ -295,11 +316,13 @@ private static async Task<WaitForResourceResponse> WaitForHealthyAsync(ResourceN
         };
     }
 
-    private static async Task<WaitForResourceResponse> WaitForRunningAsync(ResourceNotificationService notificationService, string resourceName, CancellationToken cancellationToken)
+    private static async Task<WaitForResourceResponse> WaitForRunningAsync(ResourceNotificationService notificationService, WaitResourceTarget target, CancellationToken cancellationToken)
     {
-        var resourceEvent = await notificationService.WaitForResourceAsync(
-            resourceName,
+        var resourceEvent = await WaitForResourceEventAsync(
+            notificationService,
+            target,
             re => re.Snapshot.State?.Text == KnownResourceStates.Running || KnownResourceStates.TerminalStates.Contains(re.Snapshot.State?.Text) || re.Snapshot.ExitCode is not null,
+            $"Resource '{target.DisplayName}' failed to reach the target state before the operation was cancelled.",
             cancellationToken).ConfigureAwait(false);
 
         var state = resourceEvent.Snapshot.State?.Text;
@@ -310,15 +333,17 @@ private static async Task<WaitForResourceResponse> WaitForRunningAsync(ResourceN
             Success = isRunning,
             State = state,
             HealthStatus = resourceEvent.Snapshot.HealthStatus?.ToString(),
-            ErrorMessage = isRunning ? null : $"Resource '{resourceName}' failed to reach 'Running' state. Current state: {state ?? "Unknown"}."
+            ErrorMessage = isRunning ? null : $"Resource '{target.DisplayName}' failed to reach 'Running' state. Current state: {state ?? "Unknown"}."
         };
     }
 
-    private static async Task<WaitForResourceResponse> WaitForTerminalAsync(ResourceNotificationService notificationService, string resourceName, CancellationToken cancellationToken)
+    private static async Task<WaitForResourceResponse> WaitForTerminalAsync(ResourceNotificationService notificationService, WaitResourceTarget target, CancellationToken cancellationToken)
     {
-        var resourceEvent = await notificationService.WaitForResourceAsync(
-            resourceName,
+        var resourceEvent = await WaitForResourceEventAsync(
+            notificationService,
+            target,
             re => KnownResourceStates.TerminalStates.Contains(re.Snapshot.State?.Text) || re.Snapshot.ExitCode is not null,
+            $"Resource '{target.DisplayName}' failed to reach the target state before the operation was cancelled.",
             cancellationToken).ConfigureAwait(false);
 
         return new WaitForResourceResponse
@@ -329,6 +354,118 @@ private static async Task<WaitForResourceResponse> WaitForTerminalAsync(Resource
         };
     }
 
+    private static async Task<ResourceEvent> WaitForResourceEventAsync(
+        ResourceNotificationService notificationService,
+        WaitResourceTarget target,
+        Func<ResourceEvent, bool> predicate,
+        string cancellationMessage,
+        CancellationToken cancellationToken)
+    {
+        try
+        {
+            await foreach (var resourceEvent in notificationService.WatchAsync(cancellationToken).ConfigureAwait(false))
+            {
+                if (target.Matches(resourceEvent) && predicate(resourceEvent))
+                {
+                    return resourceEvent;
+                }
+            }
+        }
+        catch (OperationCanceledException ex)
+        {
+            throw new OperationCanceledException(cancellationMessage, ex, ex.CancellationToken);
+        }
+
+        throw new OperationCanceledException(cancellationMessage);
+    }
+
+    private WaitTargetResolutionResult ResolveWaitTarget(ResourceNotificationService notificationService, string requestedResourceName)
+    {
+        var appModel = serviceProvider.GetService<DistributedApplicationModel>();
+        if (notificationService.TryGetCurrentState(requestedResourceName, out var resourceEvent))
+        {
+            return WaitTargetResolutionResult.Success(new WaitResourceTarget(
+                ResolveDisplayName(appModel, requestedResourceName, resourceEvent.ResourceId),
+                resourceEvent.ResourceId,
+                null));
+        }
+
+        // During startup the resource may not have published its first snapshot yet, so fall back to
+        // the app model to resolve the requested logical name or resolved resource id.
+        if (appModel is null)
+        {
+            return WaitTargetResolutionResult.Success(new WaitResourceTarget(requestedResourceName, requestedResourceName, requestedResourceName));
+        }
+
+        var matchingResource = appModel.Resources.SingleOrDefault(resource => string.Equals(resource.Name, requestedResourceName, StringComparisons.ResourceName));
+        if (matchingResource is not null)
+        {
+            var resolvedResourceNames = matchingResource.GetResolvedResourceNames();
+            return resolvedResourceNames.Length switch
+            {
+                1 => WaitTargetResolutionResult.Success(new WaitResourceTarget(requestedResourceName, resolvedResourceNames[0], null)),
+                > 1 => WaitTargetResolutionResult.Ambiguous(requestedResourceName),
+                _ => WaitTargetResolutionResult.NotFound(requestedResourceName)
+            };
+        }
+
+        var resolvedMatches = appModel.Resources
+            .Select(resource => new { Resource = resource, ResolvedResourceNames = resource.GetResolvedResourceNames() })
+            .Where(match => match.ResolvedResourceNames.Any(resourceName => string.Equals(resourceName, requestedResourceName, StringComparisons.ResourceName)))
+            .Take(2)
+            .ToArray();
+
+        return resolvedMatches.Length switch
+        {
+            1 => WaitTargetResolutionResult.Success(new WaitResourceTarget(
+                resolvedMatches[0].ResolvedResourceNames.Length == 1 ? resolvedMatches[0].Resource.Name : requestedResourceName,
+                requestedResourceName,
+                null)),
+            > 1 => WaitTargetResolutionResult.Ambiguous(requestedResourceName),
+            _ => WaitTargetResolutionResult.NotFound(requestedResourceName)
+        };
+    }
+
+    private static string ResolveDisplayName(DistributedApplicationModel? appModel, string requestedResourceName, string resolvedResourceName)
+    {
+        if (appModel is null)
+        {
+            return requestedResourceName;
+        }
+
+        var matchingResource = appModel.Resources
+            .Select(resource => new { Resource = resource, ResolvedResourceNames = resource.GetResolvedResourceNames() })
+            .SingleOrDefault(match => match.ResolvedResourceNames.Any(resourceName => string.Equals(resourceName, resolvedResourceName, StringComparisons.ResourceName)));
+
+        return matchingResource is { ResolvedResourceNames.Length: 1 }
+            ? matchingResource.Resource.Name
+            : requestedResourceName;
+    }
+
+    private sealed record WaitResourceTarget(string DisplayName, string? ResourceId, string? ResourceName)
+    {
+        public bool Matches(ResourceEvent resourceEvent)
+        {
+            return (ResourceId is not null && string.Equals(resourceEvent.ResourceId, ResourceId, StringComparisons.ResourceName))
+                || (ResourceName is not null && string.Equals(resourceEvent.Resource.Name, ResourceName, StringComparisons.ResourceName));
+        }
+    }
+
+    private sealed record WaitTargetResolutionResult(WaitResourceTarget? Target, bool ResourceNotFound, string ErrorMessage)
+    {
+        public static WaitTargetResolutionResult Success(WaitResourceTarget target) => new(target, ResourceNotFound: false, string.Empty);
+
+        public static WaitTargetResolutionResult NotFound(string requestedResourceName) => new(
+            null,
+            ResourceNotFound: true,
+            $"Resource '{requestedResourceName}' was not found.");
+
+        public static WaitTargetResolutionResult Ambiguous(string requestedResourceName) => new(
+            null,
+            ResourceNotFound: false,
+            $"Resource '{requestedResourceName}' is ambiguous because it has multiple replicas. Specify the exact instance name.");
+    }
+
     #endregion
 
     #region V1 API Methods (Legacy - Keep for backward compatibility)