From bc13c82f86684f7f5ef7f3581743c330db8a67d9 Mon Sep 17 00:00:00 2001 From: MillyWei Date: Wed, 7 Jan 2026 18:23:17 +0800 Subject: [PATCH 1/3] Fix: Honor CancellationToken in DownloadModelWithProgressAsync - Add ct.ThrowIfCancellationRequested() check at the start of download loop - Ensures cancellation is checked on every iteration for immediate response - Add unit test to verify OperationCanceledException is thrown on cancellation - Resolves issue where downloads continued to 100% despite cancel requests --- sdk/cs/src/FoundryLocalManager.cs | 3 + .../FoundryLocalManagerTest.cs | 56 +++++++++++++++++++ 2 files changed, 59 insertions(+) diff --git a/sdk/cs/src/FoundryLocalManager.cs b/sdk/cs/src/FoundryLocalManager.cs index 737c965..27fe5ff 100644 --- a/sdk/cs/src/FoundryLocalManager.cs +++ b/sdk/cs/src/FoundryLocalManager.cs @@ -429,6 +429,9 @@ public async IAsyncEnumerable DownloadModelWithProgressAs while (!completed && (line = await reader.ReadLineAsync(ct)) is not null) { + // Check for cancellation at the start of each iteration + ct.ThrowIfCancellationRequested(); + // Check if this line contains download percentage if (line.StartsWith("Total", StringComparison.CurrentCultureIgnoreCase) && line.Contains("Downloading") && line.Contains('%')) { diff --git a/sdk/cs/test/FoundryLocal.Tests/FoundryLocalManagerTest.cs b/sdk/cs/test/FoundryLocal.Tests/FoundryLocalManagerTest.cs index f8f4003..fa1624e 100644 --- a/sdk/cs/test/FoundryLocal.Tests/FoundryLocalManagerTest.cs +++ b/sdk/cs/test/FoundryLocal.Tests/FoundryLocalManagerTest.cs @@ -499,6 +499,62 @@ public async Task DownloadModelWithProgressAsync_Error() Assert.Equal("Download error occurred.", last.ErrorMessage); } + [Fact] + public async Task DownloadModelWithProgressAsync_CancellationToken_ThrowsOperationCanceledException() + { + // GIVEN + MockCatalog(includeCuda: true); + MockLocalModels(); // empty to force a download + + var stream = new MemoryStream(); + // Simulate multiple progress updates + for (int i = 0; i <= 50; i += 10) + { + var progressLine = Encoding.UTF8.GetBytes($"Total {i}.00% Downloading model.onnx.data\n"); + stream.Write(progressLine, 0, progressLine.Length); + } + + // Add more lines to ensure the loop continues + for (int i = 60; i <= 100; i += 10) + { + var progressLine = Encoding.UTF8.GetBytes($"Total {i}.00% Downloading model.onnx.data\n"); + stream.Write(progressLine, 0, progressLine.Length); + } + + var doneLine = Encoding.UTF8.GetBytes("[DONE] All Completed!\n"); + stream.Write(doneLine, 0, doneLine.Length); + using (var writer = new Utf8JsonWriter(stream)) + { + JsonSerializer.Serialize(writer, new { success = true, errorMessage = (string?)null }); + } + stream.Position = 0; + + _mockHttp.When("/openai/download").Respond("application/json", stream); + + // WHEN + using var cts = new System.Threading.CancellationTokenSource(); + var progressList = new List(); + + // THEN + await Assert.ThrowsAsync(async () => + { + await foreach (var p in _manager.DownloadModelWithProgressAsync("model-3", ct: cts.Token)) + { + progressList.Add(p); + + // Cancel after receiving a few progress updates (simulating user clicking cancel button) + if (progressList.Count >= 3) + { + cts.Cancel(); + } + } + }); + + // Verify we received some progress before cancellation + Assert.True(progressList.Count >= 3, $"Expected at least 3 progress updates, but got {progressList.Count}"); + Assert.True(progressList.All(p => p.Percentage < 100), "Should not reach 100% after cancellation"); + } + [Fact] public async Task LoadModelAsync_Succeeds_AndPassesEpOverrideWhenCudaPresent() { From 1ff8201c8b18260c6b8c5ed9f654f398e8fa065a Mon Sep 17 00:00:00 2001 From: MillyWei Date: Wed, 7 Jan 2026 18:37:55 +0800 Subject: [PATCH 2/3] Fix cancellation token handling in DownloadModelWithProgressAsync - Add cancellation check after yielding progress for faster response - Improve test to simulate realistic mid-download cancellation - Add test for pre-cancelled token scenario - Update assertions to account for async pipeline behavior Fixes issue where downloads continue to 100% despite cancellation request --- sdk/cs/src/FoundryLocalManager.cs | 2 + .../FoundryLocalManagerTest.cs | 65 +++++++++++++------ 2 files changed, 46 insertions(+), 21 deletions(-) diff --git a/sdk/cs/src/FoundryLocalManager.cs b/sdk/cs/src/FoundryLocalManager.cs index 27fe5ff..6147c50 100644 --- a/sdk/cs/src/FoundryLocalManager.cs +++ b/sdk/cs/src/FoundryLocalManager.cs @@ -440,6 +440,8 @@ public async IAsyncEnumerable DownloadModelWithProgressAs if (double.TryParse(percentStr, out var percentage)) { yield return ModelDownloadProgress.Progress(percentage); + // Check cancellation after yielding progress to ensure timely response + ct.ThrowIfCancellationRequested(); } } else if (line.Contains("[DONE]") || line.Contains("All Completed")) diff --git a/sdk/cs/test/FoundryLocal.Tests/FoundryLocalManagerTest.cs b/sdk/cs/test/FoundryLocal.Tests/FoundryLocalManagerTest.cs index fa1624e..87b034e 100644 --- a/sdk/cs/test/FoundryLocal.Tests/FoundryLocalManagerTest.cs +++ b/sdk/cs/test/FoundryLocal.Tests/FoundryLocalManagerTest.cs @@ -502,36 +502,26 @@ public async Task DownloadModelWithProgressAsync_Error() [Fact] public async Task DownloadModelWithProgressAsync_CancellationToken_ThrowsOperationCanceledException() { - // GIVEN + // GIVEN - Simulate an in-progress download that will be cancelled mid-stream MockCatalog(includeCuda: true); MockLocalModels(); // empty to force a download var stream = new MemoryStream(); - // Simulate multiple progress updates + // Simulate partial download progress (NOT complete) + // This mimics a real download that's in progress when user clicks cancel for (int i = 0; i <= 50; i += 10) { var progressLine = Encoding.UTF8.GetBytes($"Total {i}.00% Downloading model.onnx.data\n"); stream.Write(progressLine, 0, progressLine.Length); } - // Add more lines to ensure the loop continues - for (int i = 60; i <= 100; i += 10) - { - var progressLine = Encoding.UTF8.GetBytes($"Total {i}.00% Downloading model.onnx.data\n"); - stream.Write(progressLine, 0, progressLine.Length); - } - - var doneLine = Encoding.UTF8.GetBytes("[DONE] All Completed!\n"); - stream.Write(doneLine, 0, doneLine.Length); - using (var writer = new Utf8JsonWriter(stream)) - { - JsonSerializer.Serialize(writer, new { success = true, errorMessage = (string?)null }); - } + // DO NOT add [DONE] or completion JSON - this simulates incomplete download + // The stream will end without completion, which is realistic for cancelled operations stream.Position = 0; _mockHttp.When("/openai/download").Respond("application/json", stream); - // WHEN + // WHEN - User cancels after receiving some progress updates using var cts = new System.Threading.CancellationTokenSource(); var progressList = new List(); @@ -542,17 +532,50 @@ await Assert.ThrowsAsync(async () => { progressList.Add(p); - // Cancel after receiving a few progress updates (simulating user clicking cancel button) - if (progressList.Count >= 3) + // Cancel after receiving 2 progress updates (simulating user clicking cancel button) + // Using 2 instead of 3 to ensure we cancel before all mock data is consumed + if (progressList.Count >= 2) { cts.Cancel(); } } }); - // Verify we received some progress before cancellation - Assert.True(progressList.Count >= 3, $"Expected at least 3 progress updates, but got {progressList.Count}"); - Assert.True(progressList.All(p => p.Percentage < 100), "Should not reach 100% after cancellation"); + // Verify cancellation was responsive and happened during download + Assert.InRange(progressList.Count, 2, 4); // Allow for 1-2 items in async pipeline + Assert.All(progressList, p => Assert.False(p.IsCompleted, "Should not reach completion status")); + Assert.All(progressList, p => Assert.True(p.Percentage < 100, $"Progress {p.Percentage}% should be less than 100%")); + } + + [Fact] + public async Task DownloadModelWithProgressAsync_AlreadyCancelledToken_ThrowsImmediately() + { + // GIVEN - Token is already cancelled before download starts + MockCatalog(includeCuda: true); + MockLocalModels(); // empty to force a download + + var stream = new MemoryStream(); + var progressLine = Encoding.UTF8.GetBytes("Total 10.00% Downloading model.onnx.data\n"); + stream.Write(progressLine, 0, progressLine.Length); + stream.Position = 0; + + _mockHttp.When("/openai/download").Respond("application/json", stream); + + using var cts = new System.Threading.CancellationTokenSource(); + cts.Cancel(); // Cancel BEFORE starting download + var progressList = new List(); + + // WHEN/THEN - Should throw immediately without processing any progress + await Assert.ThrowsAsync(async () => + { + await foreach (var p in _manager.DownloadModelWithProgressAsync("model-3", ct: cts.Token)) + { + progressList.Add(p); + } + }); + + // Should not receive any progress updates since token was already cancelled + Assert.Empty(progressList); } [Fact] From a44d5b1591d66cbcf61879d6dbee6f3e1b77acbe Mon Sep 17 00:00:00 2001 From: MillyWei Date: Wed, 7 Jan 2026 19:24:45 +0800 Subject: [PATCH 3/3] Fix test: Change expected exception from OperationCanceledException to TaskCanceledException --- sdk/cs/test/FoundryLocal.Tests/FoundryLocalManagerTest.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cs/test/FoundryLocal.Tests/FoundryLocalManagerTest.cs b/sdk/cs/test/FoundryLocal.Tests/FoundryLocalManagerTest.cs index 87b034e..0f74cfa 100644 --- a/sdk/cs/test/FoundryLocal.Tests/FoundryLocalManagerTest.cs +++ b/sdk/cs/test/FoundryLocal.Tests/FoundryLocalManagerTest.cs @@ -566,7 +566,7 @@ public async Task DownloadModelWithProgressAsync_AlreadyCancelledToken_ThrowsImm var progressList = new List(); // WHEN/THEN - Should throw immediately without processing any progress - await Assert.ThrowsAsync(async () => + await Assert.ThrowsAsync(async () => { await foreach (var p in _manager.DownloadModelWithProgressAsync("model-3", ct: cts.Token)) {