Fix Federated Event Stream subscription teardown on per-event errors (#9995)

Author: michaelstaib

src/HotChocolate/Fusion/src/Fusion.Execution/Execution/ExecutionState.cs

@@ -24,6 +24,7 @@ internal sealed class ExecutionState
     private ulong[] _failedOrSkippedBitset = [];
 
     private bool _collectTelemetry;
+    private bool _processingCompletedEarly;
     private CancellationTokenSource _cts = default!;
     private byte[] _nodeStates = [];
     private int[] _remainingDependencies = [];
@@ -39,6 +40,20 @@ public void Initialize(bool collectTelemetry, CancellationTokenSource cts)
         _cts = cts;
     }
 
+    /// <summary>
+    /// Sets the CancellationTokenSource that <see cref="CancelProcessing"/> cancels, letting a subscription scope it
+    /// per event instead of to the whole request.
+    /// </summary>
+    public void SetCancellationSource(CancellationTokenSource cts)
+        => _cts = cts;
+
+    /// <summary>
+    /// True when processing stopped early because a field error null-propagated to the root,
+    /// settling the result as <c>null</c>. Lets the caller tell this self-inflicted stop from a real
+    /// cancellation (timeout or abort).
+    /// </summary>
+    public bool ProcessingCompletedEarly => _processingCompletedEarly;
+
     public void Clean()
     {
         Reset();
@@ -141,6 +156,7 @@ public void Reset()
         ClearPendingMerges();
         _mergeFailures?.Clear();
         _activeNodes = 0;
+        _processingCompletedEarly = false;
 
         Traces.Clear();
         Signal.TryResetToIdle();
@@ -223,6 +239,8 @@ public ExecutionNodeResult ApplyPendingMergeFailure(ExecutionNodeResult result)
 
     public void CancelProcessing()
     {
+        _processingCompletedEarly = true;
+
         if (!_cts.IsCancellationRequested)
         {
             _cts.Cancel();

src/HotChocolate/Fusion/src/Fusion.Execution/Execution/OperationPlanExecutor.cs

@@ -1047,13 +1047,12 @@ private static async IAsyncEnumerable<OperationResult> CreateResponseStream(
             static state => Unsafe.As<AsyncAutoResetEvent>(state)!.TryResetToIdle(),
             executionState.Signal);
 
-        // We allocate a single CancellationTokenSource per subscription and reuse it
-        // across all events via TryReset(). The execution token is linked in so that
-        // client-abort / server-shutdown still propagates.
-        var eventCts = new CancellationTokenSource();
-        var eventCtsRegistration = executionCancellationToken.UnsafeRegister(
-            static state => Unsafe.As<CancellationTokenSource>(state)!.Cancel(),
-            eventCts);
+        // We allocate one CancellationTokenSource per subscription and reuse it across
+        // healthy events via TryReset(). If a cancellation is requested because of
+        // null-propagation to the root, we replace the source before the next event.
+        // The execution token is linked in so that client-abort / server-shutdown still
+        // propagates.
+        var (eventCts, eventCtsRegistration) = CreateEventCancellation();
 
         var schemaName = GetSubscriptionSchemaName(context, subscriptionNode);
 
@@ -1140,9 +1139,14 @@ private static async IAsyncEnumerable<OperationResult> CreateResponseStream(
                     // so we throw here to properly cancel the request execution.
                     requestCancellationToken.ThrowIfCancellationRequested();
 
-                    // If the event budget was exhausted, surface it as a cancellation so the
-                    // stream tears down and the caller can observe the timeout.
-                    eventToken.ThrowIfCancellationRequested();
+                    // If the event token was cancelled by a genuine timeout or abort, tear the
+                    // stream down. A root-null halt also cancels the event token, but that is benign
+                    // (the result is a settled {data: null, errors: [...]}), so we keep the stream
+                    // alive and let context.Complete() produce it.
+                    if (!executionState.ProcessingCompletedEarly)
+                    {
+                        eventToken.ThrowIfCancellationRequested();
+                    }
 
                     result = context.Complete(reusable: true);
                 }
@@ -1166,11 +1170,24 @@ private static async IAsyncEnumerable<OperationResult> CreateResponseStream(
                         concurrencyGate!.Release();
                     }
 
-                    // Reset the shared CTS so the next event can start with a fresh budget.
-                    // If TryReset() returns false the source was cancelled (timeout or
-                    // client-abort); the thrown OperationCanceledException has already
-                    // propagated and the enumerator surfaces the teardown.
-                    eventCts.TryReset();
+                    if (executionState.ProcessingCompletedEarly)
+                    {
+                        // A root-null halt cancelled the event source. A cancelled source cannot be
+                        // reset, so swap in a fresh one (and re-link client-abort / shutdown) for the
+                        // next event. This only happens on events that null-propagate to the root.
+                        await eventCtsRegistration.DisposeAsync();
+                        eventCts.Dispose();
+                        (eventCts, eventCtsRegistration) = CreateEventCancellation();
+                    }
+                    else
+                    {
+                        // Healthy event: reset the shared source so the next event starts with a
+                        // fresh timeout budget and no allocation. If TryReset() returns false the
+                        // source was cancelled by a timeout or client-abort; the thrown
+                        // OperationCanceledException has already propagated and the enumerator
+                        // surfaces the teardown.
+                        eventCts.TryReset();
+                    }
                 }
 
                 yield return result;
@@ -1181,6 +1198,18 @@ private static async IAsyncEnumerable<OperationResult> CreateResponseStream(
             await eventCtsRegistration.DisposeAsync();
             eventCts?.Dispose();
         }
+
+        // Creates a fresh event-scoped cancellation source, links client-abort / shutdown into it,
+        // and installs it as the engine's cancellation source for the next event.
+        (CancellationTokenSource Source, CancellationTokenRegistration Registration) CreateEventCancellation()
+        {
+            var cts = new CancellationTokenSource();
+            var registration = executionCancellationToken.UnsafeRegister(
+                static state => Unsafe.As<CancellationTokenSource>(state)!.Cancel(),
+                cts);
+            executionState.SetCancellationSource(cts);
+            return (cts, registration);
+        }
     }
 
     private static string GetSubscriptionSchemaName(

src/HotChocolate/Fusion/src/Fusion.Execution/Execution/Results/FetchResultStore.cs

@@ -1819,7 +1819,10 @@ private static SourceResultElement GetDataElement(SelectionPath sourcePath, Sour
         {
             if (current.ValueKind != JsonValueKind.Object)
             {
-                return default;
+                // An intermediate value on the path is not an object. If it is null, the target
+                // value null-propagates from here, so we surface the null element (not Undefined)
+                // and let value completion integrate it together with any source error.
+                return current.ValueKind is JsonValueKind.Null ? current : default;
             }
 
             var segment = sourcePath[i];
@@ -1867,6 +1870,13 @@ private static SourceResultElement GetDataElement(SelectionPath sourcePath, Sour
         {
             var segment = sourcePath[i];
 
+            // Source schema error paths only carry response names and list indices, never type
+            // conditions, so inline-fragment segments have no corresponding level in the trie.
+            if (segment.Kind is SelectionPathSegmentKind.InlineFragment)
+            {
+                continue;
+            }
+
             if (!current.TryGetValue(segment.Name, out current))
             {
                 return null;

src/HotChocolate/Fusion/test/Fusion.AspNetCore.Tests/EventStreamEnrichmentErrorOverHttpTests.cs

@@ -0,0 +1,228 @@
+using System.Text;
+using DotNet.Testcontainers.Builders;
+using DotNet.Testcontainers.Containers;
+using HotChocolate.Execution;
+using HotChocolate.Fusion.Subscriptions.NATS;
+using HotChocolate.Transport.Http;
+using HotChocolate.Types.Composite;
+using HotChocolate.Types.Relay;
+using Microsoft.Extensions.DependencyInjection;
+using NATS.Client.Core;
+using NATS.Client.JetStream;
+using NATS.Client.JetStream.Models;
+using OperationRequest = HotChocolate.Transport.OperationRequest;
+using OperationResult = HotChocolate.Transport.OperationResult;
+
+namespace HotChocolate.Fusion;
+
+/// <summary>
+/// End-to-end coverage for a Federated Event Stream whose selected fields require an enrichment
+/// lookup that the source schema rejects. Models the fusion-demo-sfo Reviews subgraph: a relay
+/// Node <c>Review</c> (so <c>Review.id</c> is a global id and the enrichment runs through
+/// <c>node(id:)</c>), an <c>@eventStream</c> message that only carries the id, and a client that
+/// also selects <c>body</c>. When an event carries an id that is not a valid global id the lookup
+/// errors; the subscription must surface that as a per-event error result (subgraph error plus
+/// standard non-null propagation) and stay alive, not silently close.
+/// </summary>
+public class EventStreamEnrichmentErrorOverHttpTests : FusionTestBase
+{
+    private const string BrokerName = "nats";
+    private const string Topic = "onCreateReview";
+    private const string GatewayUrl = "http://localhost:5000/graphql";
+
+    [Fact]
+    public async Task Subscribe_Should_ReturnErrorAndStayAlive_When_EnrichmentLookupRejectsEventId()
+    {
+        // arrange
+        await using var nats = await JetStreamNatsFixture.StartAsync();
+        var stream = "S" + Guid.NewGuid().ToString("N");
+        var durable = "D" + Guid.NewGuid().ToString("N");
+        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(60));
+
+        await CreateStreamAsync(nats.Url, stream, Topic, cts.Token);
+
+        using var events = CreateSourceSchema(
+            "EVENTS",
+            b => b
+                .AddQueryType<ReviewsApi.Query>()
+                .AddSubscriptionType<ReviewsApi.Subscription>()
+                .AddGlobalObjectIdentification(o => o.MarkNodeFieldAsLookup = true)
+                .ModifyOptions(o => o.StrictValidation = false));
+
+        using var gateway = await CreateCompositeSchemaAsync(
+            [("EVENTS", events)],
+            configureServices: services => services.AddNatsEventStreamBroker(
+                BrokerName,
+                o =>
+                {
+                    o.Url = nats.Url;
+                    o.JetStream = new NatsJetStreamOptions { Stream = stream, DurableConsumer = durable };
+                }),
+            configureGatewayBuilder: b => b.ModifyRequestOptions(o => o.AllowOperationPlanRequests = false));
+
+        using var client = GraphQLHttpClient.Create(gateway.CreateClient());
+
+        var request = new OperationRequest(
+            """
+            subscription {
+              onCreateReview {
+                review {
+                  id
+                  body
+                }
+                cursor
+              }
+            }
+            """);
+
+        // The encoded global id the second (valid) event publishes, matching what Review.id is.
+        var schema = await events.Services.GetSchemaAsync("EVENTS", cts.Token);
+        var validId = schema.Services
+            .GetRequiredService<INodeIdSerializerAccessor>()
+            .Serializer
+            .Format("Review", 1);
+
+        // act
+        // The first event carries the raw database key (not a valid global id), so the node lookup
+        // that fetches `body` fails. The second event carries the encoded global id and resolves.
+        using var response = await client.PostAsync(request, new Uri(GatewayUrl), cts.Token);
+        await PublishAsync(nats.Url, Topic, """{"review":{"id":1}}""", cts.Token);
+        await PublishAsync(nats.Url, Topic, "{\"review\":{\"id\":\"" + validId + "\"}}", cts.Token);
+
+        var results = new List<OperationResult>();
+        await foreach (var result in response.ReadAsResultStreamAsync().WithCancellation(cts.Token))
+        {
+            results.Add(result);
+
+            if (results.Count == 2)
+            {
+                break;
+            }
+        }
+
+        // assert
+        // First event: the message carried a raw int id, so the relay `node(id: ID!)` enrichment
+        // lookup is fed an int where a string global id is required. That error is integrated and
+        // `body` (non-null) null-propagates to the non-null root; the stream is NOT torn down.
+        // Second event: the stream recovered and delivered the enriched review.
+        Assert.Equal(2, results.Count);
+        results.MatchInlineSnapshots(
+            [
+                """
+                {
+                  "data": null,
+                  "errors": [
+                    {
+                      "message": "The argument literal representation is `HotChocolate.Language.IntValueNode` which is not compatible with the request literal type `HotChocolate.Language.StringValueNode`.",
+                      "path": [
+                        "onCreateReview",
+                        "review",
+                        "body"
+                      ],
+                      "extensions": {
+                        "fieldName": "node",
+                        "argumentName": "id",
+                        "requestedType": "HotChocolate.Language.StringValueNode",
+                        "actualType": "HotChocolate.Language.IntValueNode"
+                      }
+                    }
+                  ]
+                }
+                """,
+                """
+                {
+                  "data": {
+                    "onCreateReview": {
+                      "review": {
+                        "id": "UmV2aWV3OjE=",
+                        "body": "A great read"
+                      },
+                      "cursor": "Mg=="
+                    }
+                  }
+                }
+                """
+            ]);
+    }
+
+    private static async Task CreateStreamAsync(string url, string stream, string subject, CancellationToken ct)
+    {
+        await using var connection = new NatsConnection(new NatsOpts { Url = url });
+        var js = new NatsJSContext(connection);
+        await js.CreateStreamAsync(new StreamConfig { Name = stream, Subjects = [subject] }, ct);
+    }
+
+    private static async Task PublishAsync(string url, string subject, string body, CancellationToken ct)
+    {
+        await using var connection = new NatsConnection(new NatsOpts { Url = url });
+        var js = new NatsJSContext(connection);
+        await js.PublishAsync(subject, Encoding.UTF8.GetBytes(body), cancellationToken: ct);
+    }
+
+    public static class ReviewsApi
+    {
+        public class Query
+        {
+            public string Version => "1.0.0";
+
+            // Mirrors the demo's `[Lookup, NodeResolver] GetReviewByIdAsync(int id)`.
+            [Lookup, NodeResolver]
+            public Review? GetReviewById(int id)
+                => id == 1 ? new Review(1, "A great read") : null;
+        }
+
+        public class Subscription
+        {
+            [EventStream("review { id }", Topic = Topic, Broker = BrokerName)]
+            public ReviewCreated OnCreateReview([EventCursor] string? after)
+                => EventStream.Create<ReviewCreated>(after);
+        }
+
+        public record ReviewCreated(Review Review, [property: EventCursor] string Cursor);
+
+        public record Review(int Id, string Body);
+    }
+
+    private sealed class JetStreamNatsFixture : IAsyncDisposable
+    {
+        private readonly IContainer _container;
+
+        private JetStreamNatsFixture(IContainer container) => _container = container;
+
+        public string Url => $"nats://localhost:{_container.GetMappedPublicPort(4222)}";
+
+        public static async Task<JetStreamNatsFixture> StartAsync()
+        {
+            var fixture = new JetStreamNatsFixture(
+                new ContainerBuilder("nats:2.10-alpine")
+                    .WithPortBinding(4222, assignRandomHostPort: true)
+                    .WithCommand("-js")
+                    .WithWaitStrategy(Wait.ForUnixContainer().UntilInternalTcpPortIsAvailable(4222))
+                    .Build());
+
+            await fixture._container.StartAsync();
+            await fixture.WaitForConnectionAsync();
+            return fixture;
+        }
+
+        private async Task WaitForConnectionAsync()
+        {
+            var attempt = 0;
+            while (true)
+            {
+                try
+                {
+                    await using var connection = new NatsConnection(new NatsOpts { Url = Url });
+                    await connection.ConnectAsync();
+                    return;
+                }
+                catch (NatsException) when (++attempt < 20)
+                {
+                    await Task.Delay(TimeSpan.FromMilliseconds(250));
+                }
+            }
+        }
+
+        public async ValueTask DisposeAsync() => await _container.DisposeAsync();
+    }
+}