From 90bd4a76b673fa61e9dd3b5533654db866b3279b Mon Sep 17 00:00:00 2001 From: Marcos Salamanca Date: Wed, 4 Feb 2026 15:37:22 -0600 Subject: [PATCH 1/2] add cancellation token usage to enumerator scenarios --- .../src/Services/Azure/Tenant/TenantService.cs | 2 +- .../src/Services/CosmosService.cs | 9 +++++---- .../src/Services/EventHubsService.cs | 10 +++++----- .../src/Services/FileSharesService.cs | 6 +++--- .../src/Services/FoundryService.cs | 15 ++++++++------- .../src/Services/FunctionAppService.cs | 8 ++++---- .../src/Services/LoadTestingService.cs | 2 +- .../src/Services/Util/AzureRegionChecker.cs | 4 ++-- .../Util/Usage/CognitiveServicesUsageChecker.cs | 2 +- .../Services/Util/Usage/ComputeUsageChecker.cs | 2 +- .../Util/Usage/ContainerAppUsageChecker.cs | 2 +- .../Util/Usage/ContainerInstanceUsageChecker.cs | 2 +- .../Services/Util/Usage/HDInsightUsageChecker.cs | 2 +- .../Util/Usage/MachineLearningUsageChecker.cs | 2 +- .../Services/Util/Usage/NetworkUsageChecker.cs | 2 +- .../src/Services/Util/Usage/SearchUsageChecker.cs | 2 +- .../Services/Util/Usage/StorageUsageChecker.cs | 2 +- .../src/Services/RedisService.cs | 4 ++-- .../src/Services/SearchService.cs | 6 +++--- .../src/Services/SignalRService.cs | 4 ++-- .../src/Services/StorageSyncService.cs | 10 +++++----- 21 files changed, 50 insertions(+), 48 deletions(-) diff --git a/core/Azure.Mcp.Core/src/Services/Azure/Tenant/TenantService.cs b/core/Azure.Mcp.Core/src/Services/Azure/Tenant/TenantService.cs index 40d2c6e9ad..48c546163a 100644 --- a/core/Azure.Mcp.Core/src/Services/Azure/Tenant/TenantService.cs +++ b/core/Azure.Mcp.Core/src/Services/Azure/Tenant/TenantService.cs @@ -53,7 +53,7 @@ public async Task> GetTenants(CancellationToken cancellatio options.Environment = CloudConfiguration.ArmEnvironment; var client = new ArmClient(await GetCredential(cancellationToken), default, options); - await foreach (var tenant in client.GetTenants()) + await foreach (var tenant in client.GetTenants().WithCancellation(cancellationToken)) { results.Add(tenant); } diff --git a/tools/Azure.Mcp.Tools.Cosmos/src/Services/CosmosService.cs b/tools/Azure.Mcp.Tools.Cosmos/src/Services/CosmosService.cs index 9ee5ff5d1d..cbe5ed80c1 100644 --- a/tools/Azure.Mcp.Tools.Cosmos/src/Services/CosmosService.cs +++ b/tools/Azure.Mcp.Tools.Cosmos/src/Services/CosmosService.cs @@ -32,13 +32,14 @@ private async Task GetCosmosAccountAsync( string subscription, string accountName, string? tenant = null, - RetryPolicyOptions? retryPolicy = null) + RetryPolicyOptions? retryPolicy = null, + CancellationToken cancellationToken = default) { ValidateRequiredParameters((nameof(subscription), subscription), (nameof(accountName), accountName)); - var subscriptionResource = await _subscriptionService.GetSubscription(subscription, tenant, retryPolicy); + var subscriptionResource = await _subscriptionService.GetSubscription(subscription, tenant, retryPolicy, cancellationToken); - await foreach (var account in subscriptionResource.GetCosmosDBAccountsAsync()) + await foreach (var account in subscriptionResource.GetCosmosDBAccountsAsync(cancellationToken)) { if (account.Data.Name == accountName) { @@ -74,7 +75,7 @@ private async Task CreateCosmosClientWithAuth( switch (authMethod) { case AuthMethod.Key: - var cosmosAccount = await GetCosmosAccountAsync(subscription, accountName, tenant); + var cosmosAccount = await GetCosmosAccountAsync(subscription, accountName, tenant, cancellationToken: cancellationToken); var keys = await cosmosAccount.GetKeysAsync(cancellationToken); cosmosClient = new CosmosClient( string.Format(CosmosBaseUri, accountName), diff --git a/tools/Azure.Mcp.Tools.EventHubs/src/Services/EventHubsService.cs b/tools/Azure.Mcp.Tools.EventHubs/src/Services/EventHubsService.cs index 76577b04b1..40653cc933 100644 --- a/tools/Azure.Mcp.Tools.EventHubs/src/Services/EventHubsService.cs +++ b/tools/Azure.Mcp.Tools.EventHubs/src/Services/EventHubsService.cs @@ -43,7 +43,7 @@ public async Task> GetNamespacesAsync( throw new InvalidOperationException($"Resource group '{resourceGroup}' not found"); } - await foreach (var namespaceResource in resourceGroupResource.Value.GetEventHubsNamespaces()) + await foreach (var namespaceResource in resourceGroupResource.Value.GetEventHubsNamespaces().WithCancellation(cancellationToken)) { namespaces.Add(ConvertToNamespace(namespaceResource.Data, resourceGroup)); } @@ -51,9 +51,9 @@ public async Task> GetNamespacesAsync( else { // Get namespaces from all resource groups in subscription - await foreach (var rg in subscriptionResource.GetResourceGroups()) + await foreach (var rg in subscriptionResource.GetResourceGroups().WithCancellation(cancellationToken)) { - await foreach (var namespaceResource in rg.GetEventHubsNamespaces()) + await foreach (var namespaceResource in rg.GetEventHubsNamespaces().WithCancellation(cancellationToken)) { namespaces.Add(ConvertToNamespace(namespaceResource.Data, rg.Data.Name)); } @@ -313,7 +313,7 @@ public async Task> GetEventHubsAsync( var eventHubList = new List(); - await foreach (var eventHub in namespaceResource.Value.GetEventHubs()) + await foreach (var eventHub in namespaceResource.Value.GetEventHubs().WithCancellation(cancellationToken)) { eventHubList.Add(ConvertToEventHub(eventHub.Data, resourceGroup)); } @@ -692,7 +692,7 @@ public async Task> GetConsumerGroupsAsync( var consumerGroups = new List(); - await foreach (var consumerGroup in eventHubResource.Value.GetEventHubsConsumerGroups()) + await foreach (var consumerGroup in eventHubResource.Value.GetEventHubsConsumerGroups().WithCancellation(cancellationToken)) { consumerGroups.Add(ConvertToConsumerGroup(consumerGroup.Data, resourceGroup, namespaceName, eventHubName)); } diff --git a/tools/Azure.Mcp.Tools.FileShares/src/Services/FileSharesService.cs b/tools/Azure.Mcp.Tools.FileShares/src/Services/FileSharesService.cs index 8e06d78912..40a2dd78bd 100644 --- a/tools/Azure.Mcp.Tools.FileShares/src/Services/FileSharesService.cs +++ b/tools/Azure.Mcp.Tools.FileShares/src/Services/FileSharesService.cs @@ -59,7 +59,7 @@ public async Task> ListFileSharesAsync( } var collection = resourceGroupResource.GetFileShares(); - await foreach (var fileShareResource in collection) + await foreach (var fileShareResource in collection.WithCancellation(cancellationToken)) { fileShares.Add(FileShareInfo.FromResource(fileShareResource)); } @@ -488,7 +488,7 @@ public async Task GetSnapshotAsync( var fileShareResource = await resourceGroupResource.Value.GetFileShares().GetAsync(fileShareName, cancellationToken); var snapshotCollection = fileShareResource.Value.GetFileShareSnapshots(); - await foreach (var snapshotResource in snapshotCollection) + await foreach (var snapshotResource in snapshotCollection.WithCancellation(cancellationToken)) { if (snapshotResource.Data.Name.Equals(snapshotId, StringComparison.OrdinalIgnoreCase) || snapshotResource.Data.Id.ToString().Contains(snapshotId, StringComparison.OrdinalIgnoreCase)) @@ -539,7 +539,7 @@ public async Task> ListSnapshotsAsync( var snapshotCollection = fileShareResource.Value.GetFileShareSnapshots(); var snapshots = new List(); - await foreach (var snapshotResource in snapshotCollection) + await foreach (var snapshotResource in snapshotCollection.WithCancellation(cancellationToken)) { snapshots.Add(FileShareSnapshotInfo.FromResource(snapshotResource)); } diff --git a/tools/Azure.Mcp.Tools.Foundry/src/Services/FoundryService.cs b/tools/Azure.Mcp.Tools.Foundry/src/Services/FoundryService.cs index 53e7173520..105e839597 100644 --- a/tools/Azure.Mcp.Tools.Foundry/src/Services/FoundryService.cs +++ b/tools/Azure.Mcp.Tools.Foundry/src/Services/FoundryService.cs @@ -1105,7 +1105,7 @@ public async Task ListThreads( List threads = []; try { - await foreach (var thread in threadsIterator) + await foreach (var thread in threadsIterator.WithCancellation(cancellationToken)) { threads.Add(new() { @@ -1172,7 +1172,7 @@ public async Task GetMessages( { List messages = []; var messagesIterator = agentsClient.Messages.GetMessagesAsync(threadId, cancellationToken: cancellationToken); - await foreach (var message in messagesIterator) + await foreach (var message in messagesIterator.WithCancellation(cancellationToken)) { messages.Add(message); } @@ -1499,7 +1499,7 @@ public async Task> ListAiResourcesAsync( // List all AI resources in the subscription await foreach (var account in subscriptionResource.GetCognitiveServicesAccountsAsync(cancellationToken: cancellationToken)) { - var resourceInfo = await BuildResourceInformation(account, subscriptionResource.Data.DisplayName); + var resourceInfo = await BuildResourceInformation(account, subscriptionResource.Data.DisplayName, cancellationToken); resources.Add(resourceInfo); } } @@ -1510,7 +1510,7 @@ public async Task> ListAiResourcesAsync( { if (account.Data.Id.ResourceGroupName?.Equals(resourceGroup, StringComparison.OrdinalIgnoreCase) == true) { - var resourceInfo = await BuildResourceInformation(account, subscriptionResource.Data.DisplayName); + var resourceInfo = await BuildResourceInformation(account, subscriptionResource.Data.DisplayName, cancellationToken); resources.Add(resourceInfo); } } @@ -1554,7 +1554,7 @@ public async Task GetAiResourceAsync( throw new Exception($"AI resource '{resourceName}' not found in resource group '{resourceGroup}'"); } - return await BuildResourceInformation(account.Value, subscriptionResource.Data.DisplayName); + return await BuildResourceInformation(account.Value, subscriptionResource.Data.DisplayName, cancellationToken); } catch (Exception ex) { @@ -1596,7 +1596,8 @@ public AgentsGetSdkCodeSampleResult GetSdkCodeSample(string programmingLanguage) private async Task BuildResourceInformation( CognitiveServicesAccountResource account, - string subscriptionName) + string subscriptionName, + CancellationToken cancellationToken) { var resourceInfo = new AiResourceInformation { @@ -1613,7 +1614,7 @@ private async Task BuildResourceInformation( // Get deployments for this resource try { - await foreach (var deployment in account.GetCognitiveServicesAccountDeployments()) + await foreach (var deployment in account.GetCognitiveServicesAccountDeployments().WithCancellation(cancellationToken)) { var deploymentInfo = new DeploymentInformation { diff --git a/tools/Azure.Mcp.Tools.FunctionApp/src/Services/FunctionAppService.cs b/tools/Azure.Mcp.Tools.FunctionApp/src/Services/FunctionAppService.cs index a5e916fa7f..f83a596b75 100644 --- a/tools/Azure.Mcp.Tools.FunctionApp/src/Services/FunctionAppService.cs +++ b/tools/Azure.Mcp.Tools.FunctionApp/src/Services/FunctionAppService.cs @@ -48,7 +48,7 @@ public sealed class FunctionAppService( { if (string.IsNullOrEmpty(resourceGroup)) { - await RetrieveAndAddFunctionApp(subscriptionResource.GetWebSitesAsync(cancellationToken), functionApps); + await RetrieveAndAddFunctionApp(subscriptionResource.GetWebSitesAsync(cancellationToken), functionApps, cancellationToken); } else { @@ -58,7 +58,7 @@ public sealed class FunctionAppService( throw new Exception($"Resource group '{resourceGroup}' not found in subscription '{subscription}'"); } - await RetrieveAndAddFunctionApp(resourceGroupResource.Value.GetWebSites().GetAllAsync(cancellationToken: cancellationToken), functionApps); + await RetrieveAndAddFunctionApp(resourceGroupResource.Value.GetWebSites().GetAllAsync(cancellationToken: cancellationToken), functionApps, cancellationToken); } await _cacheService.SetAsync(CacheGroup, cacheKey, functionApps, s_cacheDuration, cancellationToken); @@ -105,9 +105,9 @@ public sealed class FunctionAppService( return functionApps; } - private static async Task RetrieveAndAddFunctionApp(AsyncPageable sites, List functionApps) + private static async Task RetrieveAndAddFunctionApp(AsyncPageable sites, List functionApps, CancellationToken cancellationToken) { - await foreach (var site in sites) + await foreach (var site in sites.WithCancellation(cancellationToken)) { TryAddFunctionApp(site, functionApps); } diff --git a/tools/Azure.Mcp.Tools.LoadTesting/src/Services/LoadTestingService.cs b/tools/Azure.Mcp.Tools.LoadTesting/src/Services/LoadTestingService.cs index 4e8f8e47d2..7b90b5de2e 100644 --- a/tools/Azure.Mcp.Tools.LoadTesting/src/Services/LoadTestingService.cs +++ b/tools/Azure.Mcp.Tools.LoadTesting/src/Services/LoadTestingService.cs @@ -188,7 +188,7 @@ public async Task> GetLoadTestRunsFromTestIdAsync( } var testRuns = new List(); - await foreach (var binaryData in loadTestRunResponse) + await foreach (var binaryData in loadTestRunResponse.WithCancellation(cancellationToken)) { var testRun = JsonSerializer.Deserialize(binaryData.ToString(), LoadTestJsonContext.Default.TestRun); if (testRun != null) diff --git a/tools/Azure.Mcp.Tools.Quota/src/Services/Util/AzureRegionChecker.cs b/tools/Azure.Mcp.Tools.Quota/src/Services/Util/AzureRegionChecker.cs index 9d1b9752d3..4639778b8d 100644 --- a/tools/Azure.Mcp.Tools.Quota/src/Services/Util/AzureRegionChecker.cs +++ b/tools/Azure.Mcp.Tools.Quota/src/Services/Util/AzureRegionChecker.cs @@ -103,7 +103,7 @@ public override async Task> GetAvailableRegionsAsync(string resourc { var quotas = subscription.GetModelsAsync(region); - await foreach (CognitiveServicesModel modelElement in quotas) + await foreach (CognitiveServicesModel modelElement in quotas.WithCancellation(cancellationToken)) { var nameMatch = string.IsNullOrEmpty(_modelName) || (modelElement.Model?.Name == _modelName); @@ -153,7 +153,7 @@ public override async Task> GetAvailableRegionsAsync(string resourc try { AsyncPageable result = subscription.ExecuteLocationBasedCapabilitiesAsync(region); - await foreach (var capability in result) + await foreach (var capability in result.WithCancellation(cancellationToken)) { if (capability.SupportedServerEditions?.Any() == true) { diff --git a/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/CognitiveServicesUsageChecker.cs b/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/CognitiveServicesUsageChecker.cs index 344f8f1f8d..429cf67c41 100644 --- a/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/CognitiveServicesUsageChecker.cs +++ b/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/CognitiveServicesUsageChecker.cs @@ -18,7 +18,7 @@ public override async Task> GetUsageForLocationAsync(string loca var usages = subscription.GetUsagesAsync(location, cancellationToken: cancellationToken); var result = new List(); - await foreach (ServiceAccountUsage item in usages) + await foreach (ServiceAccountUsage item in usages.WithCancellation(cancellationToken)) { result.Add(new UsageInfo( Name: item.Name?.LocalizedValue ?? item.Name?.Value ?? string.Empty, diff --git a/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/ComputeUsageChecker.cs b/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/ComputeUsageChecker.cs index 60a5178c5b..4117e7ff75 100644 --- a/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/ComputeUsageChecker.cs +++ b/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/ComputeUsageChecker.cs @@ -18,7 +18,7 @@ public override async Task> GetUsageForLocationAsync(string loca var usages = subscription.GetUsagesAsync(location, cancellationToken); var result = new List(); - await foreach (ComputeUsage item in usages) + await foreach (ComputeUsage item in usages.WithCancellation(cancellationToken)) { result.Add(new UsageInfo( Name: item.Name?.LocalizedValue ?? item.Name?.Value ?? string.Empty, diff --git a/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/ContainerAppUsageChecker.cs b/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/ContainerAppUsageChecker.cs index 4932d110fd..012e2ce5c6 100644 --- a/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/ContainerAppUsageChecker.cs +++ b/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/ContainerAppUsageChecker.cs @@ -17,7 +17,7 @@ public override async Task> GetUsageForLocationAsync(string loca var usages = subscription.GetUsagesAsync(location, cancellationToken); var result = new List(); - await foreach (var item in usages) + await foreach (var item in usages.WithCancellation(cancellationToken)) { result.Add(new UsageInfo( Name: item.Name?.Value ?? string.Empty, diff --git a/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/ContainerInstanceUsageChecker.cs b/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/ContainerInstanceUsageChecker.cs index 613ccbdc11..cb31d483bb 100644 --- a/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/ContainerInstanceUsageChecker.cs +++ b/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/ContainerInstanceUsageChecker.cs @@ -18,7 +18,7 @@ public override async Task> GetUsageForLocationAsync(string loca var usages = subscription.GetUsagesWithLocationAsync(location, cancellationToken); var result = new List(); - await foreach (ContainerInstanceUsage item in usages) + await foreach (ContainerInstanceUsage item in usages.WithCancellation(cancellationToken)) { result.Add(new UsageInfo( Name: item.Name?.LocalizedValue ?? item.Name?.Value ?? string.Empty, diff --git a/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/HDInsightUsageChecker.cs b/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/HDInsightUsageChecker.cs index 34374a1843..c98d4398df 100644 --- a/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/HDInsightUsageChecker.cs +++ b/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/HDInsightUsageChecker.cs @@ -18,7 +18,7 @@ public override async Task> GetUsageForLocationAsync(string loca var usages = subscription.GetHDInsightUsagesAsync(location, cancellationToken); var result = new List(); - await foreach (HDInsightUsage item in usages) + await foreach (HDInsightUsage item in usages.WithCancellation(cancellationToken)) { result.Add(new UsageInfo( Name: item.Name?.LocalizedValue ?? item.Name?.Value ?? string.Empty, diff --git a/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/MachineLearningUsageChecker.cs b/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/MachineLearningUsageChecker.cs index 14d9e9d835..e03d7a438a 100644 --- a/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/MachineLearningUsageChecker.cs +++ b/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/MachineLearningUsageChecker.cs @@ -17,7 +17,7 @@ public override async Task> GetUsageForLocationAsync(string loca var usages = subscription.GetMachineLearningUsagesAsync(location, cancellationToken); var result = new List(); - await foreach (var item in usages) + await foreach (var item in usages.WithCancellation(cancellationToken)) { result.Add(new UsageInfo( Name: item.Name?.Value ?? string.Empty, diff --git a/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/NetworkUsageChecker.cs b/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/NetworkUsageChecker.cs index 72c4400bae..ed2894090f 100644 --- a/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/NetworkUsageChecker.cs +++ b/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/NetworkUsageChecker.cs @@ -17,7 +17,7 @@ public override async Task> GetUsageForLocationAsync(string loca var usages = subscription.GetUsagesAsync(location, cancellationToken); var result = new List(); - await foreach (var item in usages) + await foreach (var item in usages.WithCancellation(cancellationToken)) { result.Add(new UsageInfo( Name: item.Name?.Value ?? string.Empty, diff --git a/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/SearchUsageChecker.cs b/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/SearchUsageChecker.cs index 7e074113e4..9f0aef545b 100644 --- a/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/SearchUsageChecker.cs +++ b/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/SearchUsageChecker.cs @@ -18,7 +18,7 @@ public override async Task> GetUsageForLocationAsync(string loca var usages = subscription.GetUsagesBySubscriptionAsync(location, cancellationToken: cancellationToken); var result = new List(); - await foreach (QuotaUsageResult item in usages) + await foreach (QuotaUsageResult item in usages.WithCancellation(cancellationToken)) { result.Add(new UsageInfo( Name: item.Name?.Value ?? string.Empty, diff --git a/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/StorageUsageChecker.cs b/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/StorageUsageChecker.cs index 7ee7abfa8f..86f44bad01 100644 --- a/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/StorageUsageChecker.cs +++ b/tools/Azure.Mcp.Tools.Quota/src/Services/Util/Usage/StorageUsageChecker.cs @@ -17,7 +17,7 @@ public override async Task> GetUsageForLocationAsync(string loca var usages = subscription.GetUsagesByLocationAsync(location, cancellationToken); var result = new List(); - await foreach (var item in usages) + await foreach (var item in usages.WithCancellation(cancellationToken)) { result.Add(new UsageInfo( Name: item.Name?.Value ?? string.Empty, diff --git a/tools/Azure.Mcp.Tools.Redis/src/Services/RedisService.cs b/tools/Azure.Mcp.Tools.Redis/src/Services/RedisService.cs index b8b9d22a80..73f6f06ba7 100644 --- a/tools/Azure.Mcp.Tools.Redis/src/Services/RedisService.cs +++ b/tools/Azure.Mcp.Tools.Redis/src/Services/RedisService.cs @@ -174,7 +174,7 @@ private async Task> ListAcrResourcesAsync(SubscriptionReso try { var accessPolicyAssignmentCollection = acrResource.GetRedisCacheAccessPolicyAssignments(); - await foreach (var accessPolicyAssignmentResource in accessPolicyAssignmentCollection) + await foreach (var accessPolicyAssignmentResource in accessPolicyAssignmentCollection.WithCancellation(cancellationToken)) { if (string.IsNullOrWhiteSpace(accessPolicyAssignmentResource?.Id.ToString()) || string.IsNullOrWhiteSpace(accessPolicyAssignmentResource.Data.Name)) @@ -278,7 +278,7 @@ private async Task> ListAmrResourcesAsync(SubscriptionReso try { var databaseCollection = amrResource.GetRedisEnterpriseDatabases(); - await foreach (var databaseResource in databaseCollection) + await foreach (var databaseResource in databaseCollection.WithCancellation(cancellationToken)) { if (string.IsNullOrWhiteSpace(databaseResource?.Id.ToString()) || string.IsNullOrWhiteSpace(databaseResource.Data.Name)) diff --git a/tools/Azure.Mcp.Tools.Search/src/Services/SearchService.cs b/tools/Azure.Mcp.Tools.Search/src/Services/SearchService.cs index 78fdf1944c..01ac165f6e 100644 --- a/tools/Azure.Mcp.Tools.Search/src/Services/SearchService.cs +++ b/tools/Azure.Mcp.Tools.Search/src/Services/SearchService.cs @@ -147,7 +147,7 @@ public async Task> QueryIndex( var searchResponse = await client.SearchAsync(searchText, SearchJsonContext.Default.JsonElement, options, cancellationToken: cancellationToken); - return await ProcessSearchResults(searchResponse); + return await ProcessSearchResults(searchResponse, cancellationToken); } catch (Exception ex) { @@ -393,10 +393,10 @@ private static void ConfigureSearchOptions(string q, SearchOptions options, Sear } } - private static async Task> ProcessSearchResults(Response> searchResponse) + private static async Task> ProcessSearchResults(Response> searchResponse, CancellationToken cancellationToken) { var results = new List(); - await foreach (var result in searchResponse.Value.GetResultsAsync()) + await foreach (var result in searchResponse.Value.GetResultsAsync().WithCancellation(cancellationToken)) { results.Add(result.Document); } diff --git a/tools/Azure.Mcp.Tools.SignalR/src/Services/SignalRService.cs b/tools/Azure.Mcp.Tools.SignalR/src/Services/SignalRService.cs index 38eaf80181..0bd54970e7 100644 --- a/tools/Azure.Mcp.Tools.SignalR/src/Services/SignalRService.cs +++ b/tools/Azure.Mcp.Tools.SignalR/src/Services/SignalRService.cs @@ -58,7 +58,7 @@ public async Task> GetRuntimeAsync( if (string.IsNullOrEmpty(resourceGroup)) { var signalRResources = subscriptionResource.GetSignalRsAsync(cancellationToken); - await foreach (var runtime in signalRResources) + await foreach (var runtime in signalRResources.WithCancellation(cancellationToken)) { runtimes.Add(ConvertToRuntimeModel(runtime)); } @@ -72,7 +72,7 @@ public async Task> GetRuntimeAsync( } var signalRResources = resourceGroupResource.Value.GetSignalRs().GetAllAsync(cancellationToken); - await foreach (var runtime in signalRResources) + await foreach (var runtime in signalRResources.WithCancellation(cancellationToken)) { runtimes.Add(ConvertToRuntimeModel(runtime)); } diff --git a/tools/Azure.Mcp.Tools.StorageSync/src/Services/StorageSyncService.cs b/tools/Azure.Mcp.Tools.StorageSync/src/Services/StorageSyncService.cs index fecf8f1fb0..adda03de5d 100644 --- a/tools/Azure.Mcp.Tools.StorageSync/src/Services/StorageSyncService.cs +++ b/tools/Azure.Mcp.Tools.StorageSync/src/Services/StorageSyncService.cs @@ -61,7 +61,7 @@ public async Task> ListStorageSyncServicesAsy } var collection = resourceGroupResource.GetStorageSyncServices(); - await foreach (var serviceResource in collection) + await foreach (var serviceResource in collection.WithCancellation(cancellationToken)) { services.Add(StorageSyncServiceDataSchema.FromResource(serviceResource)); } @@ -307,7 +307,7 @@ public async Task> ListSyncGroupsAsync( var serviceResource = await resourceGroupResource.Value.GetStorageSyncServices().GetAsync(storageSyncServiceName, cancellationToken); var syncGroups = new List(); - await foreach (var syncGroupResource in serviceResource.Value.GetStorageSyncGroups()) + await foreach (var syncGroupResource in serviceResource.Value.GetStorageSyncGroups().WithCancellation(cancellationToken)) { syncGroups.Add(SyncGroupDataSchema.FromResource(syncGroupResource)); } @@ -460,7 +460,7 @@ public async Task> ListCloudEndpointsAsync( var syncGroupResource = await serviceResource.Value.GetStorageSyncGroups().GetAsync(syncGroupName, cancellationToken); var endpoints = new List(); - await foreach (var endpointResource in syncGroupResource.Value.GetCloudEndpoints()) + await foreach (var endpointResource in syncGroupResource.Value.GetCloudEndpoints().WithCancellation(cancellationToken)) { endpoints.Add(CloudEndpointDataSchema.FromResource(endpointResource)); } @@ -710,7 +710,7 @@ public async Task> ListServerEndpointsAsync( var syncGroupResource = await serviceResource.Value.GetStorageSyncGroups().GetAsync(syncGroupName, cancellationToken); var endpoints = new List(); - await foreach (var endpointResource in syncGroupResource.Value.GetStorageSyncServerEndpoints()) + await foreach (var endpointResource in syncGroupResource.Value.GetStorageSyncServerEndpoints().WithCancellation(cancellationToken)) { endpoints.Add(ServerEndpointDataSchema.FromResource(endpointResource)); } @@ -962,7 +962,7 @@ public async Task> ListRegisteredServersAsync( var serviceResource = await resourceGroupResource.Value.GetStorageSyncServices().GetAsync(storageSyncServiceName, cancellationToken); var servers = new List(); - await foreach (var serverResource in serviceResource.Value.GetStorageSyncRegisteredServers()) + await foreach (var serverResource in serviceResource.Value.GetStorageSyncRegisteredServers().WithCancellation(cancellationToken)) { servers.Add(RegisteredServerDataSchema.FromResource(serverResource)); } From 9ee378873a31e03fe2df96a543e88853f6a36398 Mon Sep 17 00:00:00 2001 From: Marcos Salamanca Date: Wed, 4 Feb 2026 16:24:20 -0600 Subject: [PATCH 2/2] update new command instructions to use cancellation token for async enumerable scenario --- servers/Azure.Mcp.Server/docs/new-command.md | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/servers/Azure.Mcp.Server/docs/new-command.md b/servers/Azure.Mcp.Server/docs/new-command.md index f9f8a44f24..c36db25fb5 100644 --- a/servers/Azure.Mcp.Server/docs/new-command.md +++ b/servers/Azure.Mcp.Server/docs/new-command.md @@ -911,10 +911,28 @@ public interface IMyService **Service Implementation Requirements:** - Pass the `CancellationToken` parameter to all async method calls - Use `cancellationToken: cancellationToken` when calling Azure SDK methods +- Use `.WithCancellation(cancellationToken)` when iterating over async enumerables with `await foreach` - Always include `CancellationToken cancellationToken` as the final parameter (only use a default value if and only if other parameters have default values) - Force callers to explicitly provide a CancellationToken - Never pass `CancellationToken.None` or `default` as a value to a `CancellationToken` method parameter +**Example - Async Enumerable Pattern:** +```csharp +// ✅ Correct: Use .WithCancellation() for async enumerables +var subscription = _armClient.GetSubscriptionResource(new($"/subscriptions/{_subscriptionId}")); +await foreach (var resourceGroup in subscription.GetResourceGroups().WithCancellation(cancellationToken)) +{ + return resourceGroup.Data.Name; +} + +// ❌ Wrong: Missing .WithCancellation() +var subscription = _armClient.GetSubscriptionResource(new($"/subscriptions/{_subscriptionId}")); +await foreach (var resourceGroup in subscription.GetResourceGroups()) +{ + return resourceGroup.Data.Name; +} +``` + **Unit Testing Requirements:** - **Mock setup**: Use `Arg.Any()` for CancellationToken parameters in mock setups - **Product code invocation**: Use `TestContext.Current.CancellationToken` when invoking product code from unit tests