From 68680df8058d0a374f3ff33ac4b4dd182276bfa8 Mon Sep 17 00:00:00 2001
From: Yannick Cholette <yannick.cholette@ubisoft.com>
Date: Tue, 13 Feb 2024 15:52:39 -0500
Subject: [PATCH] Set task as Canceled based on token

---
 src/JsonRpc/ResponseRouter.cs | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/src/JsonRpc/ResponseRouter.cs b/src/JsonRpc/ResponseRouter.cs
index d66afc528..dfff2c42e 100644
--- a/src/JsonRpc/ResponseRouter.cs
+++ b/src/JsonRpc/ResponseRouter.cs
@@ -101,10 +101,10 @@ public async Task<TResponse> Returning<TResponse>(CancellationToken cancellation
                 var tcs = new TaskCompletionSource<JToken>();
                 _router.Requests.TryAdd(nextId, ( _method, tcs ));
 
-                cancellationToken.ThrowIfCancellationRequested();
-
                 try
                 {
+                    cancellationToken.ThrowIfCancellationRequested();
+
                     _router.OutputHandler.Value.Send(
                         new OutgoingRequest
                         {
@@ -117,6 +117,13 @@ public async Task<TResponse> Returning<TResponse>(CancellationToken cancellation
                         () =>
                         {
                             if (tcs.Task.IsCompleted) return;
+
+                            if (cancellationToken.IsCancellationRequested)
+                            {
+                                tcs.SetCanceled();
+                                return;
+                            }
+
                             _router.CancelRequest(new CancelParams { Id = nextId });
                         }
                     );