diff --git a/sdk/cs/src/FoundryLocalManager.cs b/sdk/cs/src/FoundryLocalManager.cs index 737c965..6147c50 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('%')) { @@ -437,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 f8f4003..0f74cfa 100644 --- a/sdk/cs/test/FoundryLocal.Tests/FoundryLocalManagerTest.cs +++ b/sdk/cs/test/FoundryLocal.Tests/FoundryLocalManagerTest.cs @@ -499,6 +499,85 @@ public async Task DownloadModelWithProgressAsync_Error() Assert.Equal("Download error occurred.", last.ErrorMessage); } + [Fact] + public async Task DownloadModelWithProgressAsync_CancellationToken_ThrowsOperationCanceledException() + { + // 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 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); + } + + // 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 - User cancels after receiving some progress updates + 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 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 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] public async Task LoadModelAsync_Succeeds_AndPassesEpOverrideWhenCudaPresent() {