Add Publish Timeout on Disconnect

src/NATS.Client.Core/NatsConnection.Publish.cs

@@ -67,7 +67,17 @@ public ValueTask PublishAsync<T>(in NatsMsg<T> msg, INatsSerialize<T>? serialize
 
     private async ValueTask ConnectAndPublishAsync<T>(string subject, T? data, NatsHeaders? headers, string? replyTo, INatsSerialize<T> serializer, CancellationToken cancellationToken)
     {
-        await ConnectAsync().AsTask().WaitAsync(cancellationToken).ConfigureAwait(false);
+        if (Opts.PublishTimeoutOnDisconnected)
+        {
+            using var cts1 = new CancellationTokenSource(Opts.CommandTimeout);
+            using var cts2 = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, cts1.Token);
+            await ConnectAsync().AsTask().WaitAsync(cts2.Token).ConfigureAwait(false);
+        }
+        else
+        {
+            await ConnectAsync().AsTask().WaitAsync(cancellationToken).ConfigureAwait(false);
+        }
+
         await CommandWriter.PublishAsync(subject, data, headers, replyTo, serializer, cancellationToken).ConfigureAwait(false);
     }
 }

src/NATS.Client.Core/NatsOpts.cs

@@ -206,6 +206,14 @@ public sealed record NatsOpts
     /// </remarks>
     public bool RetryOnInitialConnect { get; init; }
 
+    /// <summary>
+    /// Gets or sets a value indicating whether publish would throw an exception
+    /// when the connection is disconnected and the <see cref="CommandTimeout"/> is reached.
+    /// The default is <c>false</c>, meaning publish will not throw on disconnected state
+    /// and will wait to publish the message until reconnected.
+    /// </summary>
+    public bool PublishTimeoutOnDisconnected { get; init; } = false;
+
     internal NatsUri[] GetSeedUris(bool suppressRandomization = false)
     {
         var urls = Url.Split(',');

tests/NATS.Client.Core2.Tests/CancellationTest.cs

@@ -89,4 +89,93 @@ await Assert.ThrowsAsync<TaskCanceledException>(async () =>
             }
         });
     }
+
+    [Fact]
+    public async Task PublishTimeoutOnDisconnected_WhenEnabled_ThrowsOnTimeout()
+    {
+        var server = await NatsServerProcess.StartAsync();
+
+        await using var conn = new NatsConnection(new NatsOpts
+        {
+            Url = server.Url,
+            CommandTimeout = TimeSpan.FromMilliseconds(500),
+            PublishTimeoutOnDisconnected = true,
+        });
+        await conn.ConnectAsync();
+
+        // Kill the server
+        await server.DisposeAsync();
+
+        var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10));
+
+        // Wait for reconnect loop to kick in
+        while (conn.ConnectionState != NatsConnectionState.Reconnecting)
+        {
+            await Task.Delay(1, cts.Token);
+        }
+
+        // PublishAsync should throw TaskCanceledException due to CommandTimeout
+        await Assert.ThrowsAsync<TaskCanceledException>(() => conn.PublishAsync("test", "data").AsTask());
+    }
+
+    [Fact]
+    public async Task PublishTimeoutOnDisconnected_WhenDisabled_WaitsIndefinitely()
+    {
+        var server = await NatsServerProcess.StartAsync();
+
+        await using var conn = new NatsConnection(new NatsOpts
+        {
+            Url = server.Url,
+            PublishTimeoutOnDisconnected = false, // default
+        });
+        await conn.ConnectAsync();
+
+        // Kill the server
+        await server.DisposeAsync();
+
+        var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10));
+
+        // Wait for reconnect loop to kick in
+        while (conn.ConnectionState != NatsConnectionState.Reconnecting)
+        {
+            await Task.Delay(1, cts.Token);
+        }
+
+        // With option disabled, publish waits indefinitely - user cancellation should work
+        var userCts = new CancellationTokenSource(TimeSpan.FromMilliseconds(500));
+
+        // Should throw TaskCanceledException (user-initiated cancellation)
+        await Assert.ThrowsAsync<TaskCanceledException>(() => conn.PublishAsync("test", "data", cancellationToken: userCts.Token).AsTask());
+    }
+
+    [Fact]
+    public async Task PublishTimeoutOnDisconnected_UserCancellation_StillWorks()
+    {
+        var server = await NatsServerProcess.StartAsync();
+
+        await using var conn = new NatsConnection(new NatsOpts
+        {
+            Url = server.Url,
+            CommandTimeout = TimeSpan.FromSeconds(30), // Long timeout
+            PublishTimeoutOnDisconnected = true,
+        });
+        await conn.ConnectAsync();
+
+        // Kill the server
+        await server.DisposeAsync();
+
+        var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10));
+
+        // Wait for reconnect loop to kick in
+        while (conn.ConnectionState != NatsConnectionState.Reconnecting)
+        {
+            await Task.Delay(1, cts.Token);
+        }
+
+        // User cancellation should still work even with PublishTimeoutOnDisconnected enabled
+        var userCts = new CancellationTokenSource(TimeSpan.FromMilliseconds(500));
+
+        // Should throw TaskCanceledException (user-initiated cancellation) before CommandTimeout
+        await Assert.ThrowsAsync<TaskCanceledException>(() => conn.PublishAsync("test", "data", cancellationToken: userCts.Token).AsTask());
+    }
 }