From a95cba1225e41d003abad892c477dd2e4b87cd56 Mon Sep 17 00:00:00 2001
From: Jordan <Cyberboss@users.noreply.github.com>
Date: Mon, 19 Jun 2023 11:23:54 -0400
Subject: [PATCH] Fix TaskExtensions.WithToken

---
 src/Tgstation.Server.Host/Extensions/TaskExtensions.cs   | 6 +++++-
 tests/Tgstation.Server.Host.Tests/Jobs/TestJobHandler.cs | 2 +-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/Tgstation.Server.Host/Extensions/TaskExtensions.cs b/src/Tgstation.Server.Host/Extensions/TaskExtensions.cs
index aeac45a0cc..e2dbb8f8ff 100644
--- a/src/Tgstation.Server.Host/Extensions/TaskExtensions.cs
+++ b/src/Tgstation.Server.Host/Extensions/TaskExtensions.cs
@@ -50,8 +50,12 @@ namespace Tgstation.Server.Host.Extensions
 			ArgumentNullException.ThrowIfNull(task);
 
 			var cancelTcs = new TaskCompletionSource();
+			Task completedTask;
 			using (cancellationToken.Register(() => cancelTcs.SetCanceled(cancellationToken)))
-				await Task.WhenAny(task, cancelTcs.Task);
+				completedTask = await Task.WhenAny(task, cancelTcs.Task);
+
+			if (completedTask != task)
+				await cancelTcs.Task;
 
 			return await task;
 		}
diff --git a/tests/Tgstation.Server.Host.Tests/Jobs/TestJobHandler.cs b/tests/Tgstation.Server.Host.Tests/Jobs/TestJobHandler.cs
index 3f8e8663d8..ed8660ec0e 100644
--- a/tests/Tgstation.Server.Host.Tests/Jobs/TestJobHandler.cs
+++ b/tests/Tgstation.Server.Host.Tests/Jobs/TestJobHandler.cs
@@ -41,7 +41,7 @@ namespace Tgstation.Server.Host.Jobs.Tests
 				using var handler = new JobHandler(TestJob);
 				await Assert.ThrowsExceptionAsync<InvalidOperationException>(() => handler.Wait(cts.Token));
 				handler.Start();
-				await Assert.ThrowsExceptionAsync<OperationCanceledException>(() => handler.Wait(cts.Token));
+				await Assert.ThrowsExceptionAsync<TaskCanceledException>(() => handler.Wait(cts.Token));
 				tcs.SetResult();
 				await handler.Wait(default);
 			}