General cleanups to improve reliability (#2189)

* General cleanups to improve reliability

* Few small tweaks

* Change log level, remove default param value, update parameter description

* Use ConcurrentDictionary inside group keys instead of managing two dictionaries

* Fix logic

* Give some breadcrumps to display error

* Add ConcurrentHashSet for convenience

Dockerfile

@@ -34,4 +34,9 @@ RUN if [ ! -f $EXECUTABLE_NAME ]; then \
 COPY ${PUBLISH_DIR}/${EXECUTABLE_NAME} server-binary
 RUN chmod +x server-binary && test -x server-binary
 
+# Run as non-root user for security hardening
+RUN adduser -D -H -s /sbin/nologin mcp && \
+    chown -R mcp:mcp /mcp-server
+USER mcp
+
 ENTRYPOINT ["./server-binary", "server", "start"]

core/Azure.Mcp.Core/src/Services/Azure/Subscription/SubscriptionService.cs

@@ -6,13 +6,19 @@ using Azure.Mcp.Core.Options;
 using Azure.Mcp.Core.Services.Azure.Tenant;
 using Azure.Mcp.Core.Services.Caching;
 using Azure.ResourceManager.Resources;
+using Microsoft.Extensions.Logging;
 
 namespace Azure.Mcp.Core.Services.Azure.Subscription;
 
-public class SubscriptionService(ICacheService cacheService, ITenantService tenantService)
+public class SubscriptionService(
+    ICacheService cacheService,
+    ITenantService tenantService,
+    ILogger<SubscriptionService> logger)
     : BaseAzureService(tenantService), ISubscriptionService
 {
+    private const int MaxSubscriptions = 10_000;
     private readonly ICacheService _cacheService = cacheService ?? throw new ArgumentNullException(nameof(cacheService));
+    private readonly ILogger<SubscriptionService> _logger = logger ?? throw new ArgumentNullException(nameof(logger));
     private const string CacheGroup = "subscription";
     private const string CacheKey = "subscriptions";
     private const string SubscriptionCacheKey = "subscription";
@@ -36,6 +42,11 @@ public class SubscriptionService(ICacheService cacheService, ITenantService tena
         await foreach (var subscription in subscriptions.WithCancellation(cancellationToken))
         {
             results.Add(subscription.Data);
+            if (results.Count >= MaxSubscriptions)
+            {
+                _logger.LogWarning("Reached maximum subscription limit of {MaxSubscriptions}. Some subscriptions may not be included in the results.", MaxSubscriptions);
+                break;
+            }
         }
 
         // Cache the results
@@ -74,6 +85,7 @@ public class SubscriptionService(ICacheService cacheService, ITenantService tena
         return response.Value;
     }
 
+    // TODO (alzimmer): Why does this method take tenant?
     public bool IsSubscriptionId(string subscription, string? tenant = null)
     {
         return Guid.TryParse(subscription, out _);

core/Azure.Mcp.Core/src/Services/Azure/Tenant/TenantService.cs

@@ -7,14 +7,17 @@ using Azure.Mcp.Core.Services.Azure.Authentication;
 using Azure.Mcp.Core.Services.Caching;
 using Azure.ResourceManager;
 using Azure.ResourceManager.Resources;
+using Microsoft.Extensions.Logging;
 
 namespace Azure.Mcp.Core.Services.Azure.Tenant;
 
 public class TenantService : BaseAzureService, ITenantService
 {
+    private const int MaxTenants = 10_000;
     private readonly IAzureTokenCredentialProvider _credentialProvider;
     private readonly ICacheService _cacheService;
     private readonly IHttpClientFactory _httpClientFactory;
+    private readonly ILogger<TenantService> _logger;
     private const string CacheGroup = "tenant";
     private const string CacheKey = "tenants";
     private static readonly TimeSpan s_cacheDuration = CacheDurations.Tenant;
@@ -23,12 +26,14 @@ public class TenantService : BaseAzureService, ITenantService
         IAzureTokenCredentialProvider credentialProvider,
         ICacheService cacheService,
         IHttpClientFactory clientFactory,
-        IAzureCloudConfiguration cloudConfiguration)
+        IAzureCloudConfiguration cloudConfiguration,
+        ILogger<TenantService> logger)
     {
         _credentialProvider = credentialProvider;
         _cacheService = cacheService ?? throw new ArgumentNullException(nameof(cacheService));
         _httpClientFactory = clientFactory ?? throw new ArgumentNullException(nameof(clientFactory));
         CloudConfiguration = cloudConfiguration ?? throw new ArgumentNullException(nameof(cloudConfiguration));
+        _logger = logger ?? throw new ArgumentNullException(nameof(logger));
         TenantService = this;
     }
 
@@ -56,6 +61,11 @@ public class TenantService : BaseAzureService, ITenantService
         await foreach (var tenant in client.GetTenants().WithCancellation(cancellationToken))
         {
             results.Add(tenant);
+            if (results.Count >= MaxTenants)
+            {
+                _logger.LogWarning("Reached maximum tenant limit of {MaxTenants}. Some tenants may not be included in the results.", MaxTenants);
+                break;
+            }
         }
 
         // Cache the results

core/Azure.Mcp.Core/tests/Azure.Mcp.Core.LiveTests/Services/Azure/Authentication/AuthenticationIntegrationTests.cs

@@ -12,6 +12,7 @@ using Azure.Mcp.Core.Services.Azure.Tenant;
 using Azure.Mcp.Core.Services.Caching;
 using Azure.ResourceManager.Resources;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Logging;
 using Microsoft.Mcp.Tests;
 using NSubstitute;
 using Xunit;
@@ -33,6 +34,7 @@ public class AuthenticationIntegrationTests : IAsyncLifetime
         services.AddLogging();
         services.AddSingleton(Substitute.For<ICacheService>());
         services.AddSingleton(Substitute.For<ITenantService>());
+        services.AddSingleton(Substitute.For<ILogger<SubscriptionService>>());
         services.AddSingleton<ISubscriptionService, SubscriptionService>();
 
         _serviceProvider = services.BuildServiceProvider();

core/Microsoft.Mcp.Core/src/Areas/Server/Commands/Runtime/McpRuntime.cs

@@ -130,7 +130,7 @@ public sealed class McpRuntime : IMcpRuntime
             {
                 Content = [new TextContentBlock
                 {
-                    Text = ex.Message,
+                    Text = !string.IsNullOrWhiteSpace(ex.Message) ? ex.Message : "An unknown error occurred while trying to call the tool.",
                 }],
                 IsError = true,
             };

core/Microsoft.Mcp.Core/src/Helpers/EmbeddedResourceHelper.cs

@@ -3,6 +3,7 @@
 
 using System.Reflection;
 using System.Text.RegularExpressions;
+using Microsoft.Mcp.Core.Helpers;
 
 namespace Azure.Mcp.Core.Helpers
 {
@@ -36,7 +37,7 @@ namespace Azure.Mcp.Core.Helpers
             Regex regex;
             try
             {
-                regex = new Regex(resourcePattern);
+                regex = RegexHelper.CreateRegex(resourcePattern, RegexOptions.None);
             }
             catch (ArgumentException ex)
             {

core/Microsoft.Mcp.Core/src/Helpers/RegexHelper.cs

@@ -0,0 +1,30 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System.Text.RegularExpressions;
+
+namespace Microsoft.Mcp.Core.Helpers;
+
+public static class RegexHelper
+{
+    public static readonly TimeSpan DefaultRegexTimeout = TimeSpan.FromSeconds(3);
+
+    /// <summary>
+    /// Creates a Regex object with a specified pattern and options, and a default timeout to prevent excessive backtracking.
+    /// </summary>
+    /// <param name="pattern">The regex pattern.</param>
+    /// <param name="options">Regex options.</param>
+    /// <returns>The regex.</returns>
+    /// <exception cref="ArgumentException">If the regex pattern is invalid.</exception>
+    public static Regex CreateRegex(string pattern, RegexOptions options)
+    {
+        try
+        {
+            return new Regex(pattern, options, DefaultRegexTimeout);
+        }
+        catch (ArgumentException ex)
+        {
+            throw new ArgumentException($"Invalid regex pattern: '{pattern}'", nameof(pattern), ex);
+        }
+    }
+}

core/Microsoft.Mcp.Core/src/Models/ConcurrentHashSet.cs

@@ -0,0 +1,60 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System.Collections;
+using System.Collections.Concurrent;
+
+namespace Microsoft.Mcp.Core.Models;
+
+/// <summary>
+/// Simple concurrent hash set implementation using ConcurrentDictionary.
+/// This is not a full implementation of ISet and only supports a subset of operations needed for our use case.
+/// </summary>
+/// <typeparam name="T">The key type.</typeparam>
+public sealed class ConcurrentHashSet<T> : IEnumerable<T> where T : notnull
+{
+    private readonly ConcurrentDictionary<T, byte> _dictionary = new();
+
+    /// <summary>
+    /// Gets the number of items contained in the set.
+    /// </summary>
+    public int Count => _dictionary.Count;
+
+    /// <summary>
+    /// Adds an item to the current set and returns a value to indicate if the item was successfully added.
+    /// </summary>
+    /// <param name="item">The item to add to the set.</param>
+    /// <returns>true if the item is added to the set; false if the item is already in the set.</returns>
+    public bool Add(T item) => _dictionary.TryAdd(item, 0);
+
+    /// <summary>
+    /// Determines whether the set contains a specific item.
+    /// </summary>
+    /// <param name="item">The item to locate in the set.</param>
+    /// <returns>true if the item is found in the set; otherwise, false.</returns>
+    public bool Contains(T item) => _dictionary.ContainsKey(item);
+
+    /// <summary>
+    /// Removes all items from the set.
+    /// </summary>
+    public void Clear() => _dictionary.Clear();
+
+    /// <summary>
+    /// Removes the specified item from the set.
+    /// </summary>
+    /// <param name="item">The item to remove from the set.</param>
+    /// <returns>true if the item was successfully removed from the set; otherwise, false. false is also returned if the item is not in the set.</returns>
+    public bool Remove(T item) => _dictionary.TryRemove(item, out _);
+
+    /// <summary>
+    /// Returns an enumerator that iterates through the set.
+    /// </summary>
+    /// <returns>An enumerator that can be used to iterate through the set.</returns>
+    public IEnumerator<T> GetEnumerator() => _dictionary.Keys.GetEnumerator();
+
+    /// <summary>
+    /// Returns an enumerator that iterates through the set.
+    /// </summary>
+    /// <returns>An enumerator that can be used to iterate through the set.</returns>
+    IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
+}

core/Microsoft.Mcp.Core/src/Services/Azure/Authentication/WindowHandleProvider.cs

@@ -31,10 +31,12 @@ public static partial class WindowHandleProvider
                     : "X display is available.");
                 return display;
             }
-            catch (Exception ex)
+            catch (Exception)
             {
+                // X11 display detection failed; running in headless mode.
+                // Avoid logging exception details to console to prevent information disclosure.
                 Console.ForegroundColor = ConsoleColor.Red;
-                Console.WriteLine(ex.ToString());
+                Console.WriteLine($"Failed to detect X display. Running in headless mode. Exception type: {typeof(Exception).FullName}");
                 Console.ResetColor();
             }
         }

core/Microsoft.Mcp.Core/src/Services/Caching/SingleUserCliCacheService.cs

@@ -3,6 +3,7 @@
 
 using System.Collections.Concurrent;
 using Microsoft.Extensions.Caching.Memory;
+using Microsoft.Mcp.Core.Models;
 
 namespace Azure.Mcp.Core.Services.Caching;
 
@@ -20,8 +21,10 @@ namespace Azure.Mcp.Core.Services.Caching;
 /// </remarks>
 public class SingleUserCliCacheService(IMemoryCache memoryCache) : ICacheService
 {
+    private static readonly object s_sentinal = new();
     private readonly IMemoryCache _memoryCache = memoryCache;
-    private static readonly ConcurrentDictionary<string, HashSet<string>> s_groupKeys = new();
+
+    private static readonly ConcurrentDictionary<string, ConcurrentHashSet<string>> s_groupKeys = new();
 
     public ValueTask<T?> GetAsync<T>(string group, string key, TimeSpan? expiration = null, CancellationToken cancellationToken = default)
     {
@@ -43,10 +46,9 @@ public class SingleUserCliCacheService(IMemoryCache memoryCache) : ICacheService
 
         _memoryCache.Set(cacheKey, data, options);
 
-        // Track the key in the group
         s_groupKeys.AddOrUpdate(
             group,
-            new HashSet<string> { key },
+            [key],
             (_, keys) =>
             {
                 keys.Add(key);
@@ -74,7 +76,8 @@ public class SingleUserCliCacheService(IMemoryCache memoryCache) : ICacheService
     {
         if (s_groupKeys.TryGetValue(group, out var keys))
         {
-            return new ValueTask<IEnumerable<string>>(keys.AsEnumerable());
+            // Return a snapshot to avoid concurrent modification during enumeration.
+            return new ValueTask<IEnumerable<string>>([.. keys]);
         }
 
         return new ValueTask<IEnumerable<string>>([]);
@@ -102,8 +105,10 @@ public class SingleUserCliCacheService(IMemoryCache memoryCache) : ICacheService
             return default;
         }
 
+        string[] keysSnapshot = [.. keys];
+
         // Remove each key in the group from the cache
-        foreach (var key in keys)
+        foreach (var key in keysSnapshot)
         {
             string cacheKey = GetGroupKey(group, key);
             _memoryCache.Remove(cacheKey);

tools/Azure.Mcp.Tools.AppConfig/tests/Azure.Mcp.Tools.AppConfig.LiveTests/AppConfigCommandTests.cs

@@ -44,8 +44,8 @@ public class AppConfigCommandTests : RecordedCommandTestsBase
         _httpClientProvider = TestHttpClientFactoryProvider.Create(fixture);
         var httpClientFactory = _httpClientProvider.GetRequiredService<IHttpClientFactory>();
         var cloudConfiguration = new AzureCloudConfiguration(new ConfigurationBuilder().Build());
-        var tenantService = new TenantService(tokenProvider, cacheService, httpClientFactory, cloudConfiguration);
-        var subscriptionService = new SubscriptionService(cacheService, tenantService);
+        var tenantService = new TenantService(tokenProvider, cacheService, httpClientFactory, cloudConfiguration, NullLogger<TenantService>.Instance);
+        var subscriptionService = new SubscriptionService(cacheService, tenantService, NullLogger<SubscriptionService>.Instance);
         _appConfigService = new AppConfigService(subscriptionService, tenantService, _logger, httpClientFactory);
     }
 

tools/Azure.Mcp.Tools.AzureMigrate/src/Helpers/AzureMigrateProjectHelper.cs

@@ -51,7 +51,7 @@ public sealed class AzureMigrateProjectHelper(
                 retryPolicy,
                 cancellationToken);
 
-            var subscriptionResource = await _subscriptionService.GetSubscription(subscription, cancellationToken: cancellationToken);
+            var subscriptionResource = await _subscriptionService.GetSubscription(subscription, tenant: tenant, cancellationToken: cancellationToken);
             ResourceIdentifier projectId = new ResourceIdentifier(
                 $"/subscriptions/{subscriptionResource.Data.SubscriptionId}/resourceGroups/{resourceGroup}/providers/{MigrateProjectResourceType}/{projectName}");
 

tools/Azure.Mcp.Tools.Cosmos/src/Validation/CosmosQueryValidator.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 using System.Text.RegularExpressions;
+using Microsoft.Mcp.Core.Helpers;
 
 namespace Azure.Mcp.Tools.Cosmos.Validation;
 
@@ -21,33 +22,28 @@ internal static class CosmosQueryValidator
 {
     private const int MaxQueryLength = 5000; // Safety cap similar to Postgres validator.
     private const string StringLiteralPlaceholder = "'str'";
-    private static readonly TimeSpan RegexTimeout = TimeSpan.FromSeconds(3); // Prevent ReDoS attacks
 
     // Regex to strip string literals, replacing them with a placeholder for safe token analysis.
-    private static readonly Regex StringLiteralPattern = new(
+    private static readonly Regex StringLiteralPattern = RegexHelper.CreateRegex(
         "'([^']|'')*'",
-        RegexOptions.Compiled,
-        RegexTimeout);
+        RegexOptions.Compiled);
 
     // Matches: OR <word> = <same_word> with optional whitespace around =
     // Catches: OR 1=1, OR 2=2, OR a=a, OR 1 = 1, etc.
-    private static readonly Regex TautologyIdentifierPattern = new(
+    private static readonly Regex TautologyIdentifierPattern = RegexHelper.CreateRegex(
         @"\bor\s+(\w+)\s*=\s*\1\b",
-        RegexOptions.IgnoreCase | RegexOptions.Compiled,
-        RegexTimeout);
+        RegexOptions.IgnoreCase | RegexOptions.Compiled);
 
     // Matches: OR '<literal>' = '<same_literal>' using a backreference on the original query.
     // Only flags when both sides are identical (e.g., OR 'x'='x'), not different values (OR 'a'='b').
-    private static readonly Regex TautologyStringLiteralPattern = new(
+    private static readonly Regex TautologyStringLiteralPattern = RegexHelper.CreateRegex(
         @"\bor\s+'((?:[^']|'')*)'\s*=\s*'\1'",
-        RegexOptions.IgnoreCase | RegexOptions.Compiled,
-        RegexTimeout);
+        RegexOptions.IgnoreCase | RegexOptions.Compiled);
 
     // Matches: OR true / OR TRUE (boolean tautology)
-    private static readonly Regex TautologyBooleanPattern = new(
+    private static readonly Regex TautologyBooleanPattern = RegexHelper.CreateRegex(
         @"\bor\s+true\b",
-        RegexOptions.IgnoreCase | RegexOptions.Compiled,
-        RegexTimeout);
+        RegexOptions.IgnoreCase | RegexOptions.Compiled);
 
     // Keywords/patterns that might indicate attempts to execute stored procedures or triggers
     private static readonly HashSet<string> BlockedPatterns = new(StringComparer.OrdinalIgnoreCase)
@@ -110,7 +106,7 @@ internal static class CosmosQueryValidator
 
         // Check for attempts to execute stored procedures or triggers
         // While these cannot be executed via SQL queries, this is defense-in-depth
-        var matches = Regex.Matches(withoutStrings, "[A-Za-z_]+", RegexOptions.Compiled, RegexTimeout);
+        var matches = Regex.Matches(withoutStrings, "[A-Za-z_]+", RegexOptions.Compiled, RegexHelper.DefaultRegexTimeout);
 
         foreach (Match m in matches)
         {

tools/Azure.Mcp.Tools.EventGrid/src/Commands/Subscription/SubscriptionListCommand.cs

@@ -99,7 +99,7 @@ public sealed class SubscriptionListCommand(ILogger<SubscriptionListCommand> log
             if (crossSubscriptionSearch)
             {
                 // Iterate all subscriptions and aggregate
-                var allSubs = await _subscriptionService.GetSubscriptions(null, options.RetryPolicy, cancellationToken);
+                var allSubs = await _subscriptionService.GetSubscriptions(options.Tenant, options.RetryPolicy, cancellationToken);
                 var aggregate = new List<EventGridSubscriptionInfo>();
                 foreach (var sub in allSubs)
                 {

tools/Azure.Mcp.Tools.FunctionApp/src/Services/FunctionAppService.cs

@@ -17,6 +17,7 @@ public sealed class FunctionAppService(
     ICacheService cacheService,
     ILogger<FunctionAppService> logger) : BaseAzureService(tenantService), IFunctionAppService
 {
+    private const int MaxFunctionApps = 10_000;
     private readonly ISubscriptionService _subscriptionService = subscriptionService ?? throw new ArgumentNullException(nameof(subscriptionService));
     private readonly ICacheService _cacheService = cacheService ?? throw new ArgumentNullException(nameof(cacheService));
     private readonly ILogger<FunctionAppService> _logger = logger ?? throw new ArgumentNullException(nameof(logger));
@@ -49,7 +50,7 @@ public sealed class FunctionAppService(
 
             if (string.IsNullOrEmpty(resourceGroup))
             {
-                await RetrieveAndAddFunctionApp(subscriptionResource.GetWebSitesAsync(cancellationToken), functionApps, cancellationToken);
+                await RetrieveAndAddFunctionApp(subscriptionResource.GetWebSitesAsync(cancellationToken), functionApps, _logger, cancellationToken);
             }
             else
             {
@@ -59,7 +60,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, cancellationToken);
+                await RetrieveAndAddFunctionApp(resourceGroupResource.Value.GetWebSites().GetAllAsync(cancellationToken: cancellationToken), functionApps, _logger, cancellationToken);
             }
 
             await _cacheService.SetAsync(CacheGroup, cacheKey, functionApps, s_cacheDuration, cancellationToken);
@@ -94,11 +95,20 @@ public sealed class FunctionAppService(
         return functionApps;
     }
 
-    private static async Task RetrieveAndAddFunctionApp(AsyncPageable<WebSiteResource> sites, List<FunctionAppInfo> functionApps, CancellationToken cancellationToken)
+    private static async Task RetrieveAndAddFunctionApp(
+        AsyncPageable<WebSiteResource> sites,
+        List<FunctionAppInfo> functionApps,
+        ILogger<FunctionAppService> logger,
+        CancellationToken cancellationToken)
     {
         await foreach (var site in sites.WithCancellation(cancellationToken))
         {
             TryAddFunctionApp(site, functionApps);
+            if (functionApps.Count >= MaxFunctionApps)
+            {
+                logger.LogWarning("Warning: Reached maximum function app limit of {MaxFunctionApps}. Some function apps may not be included in the results.", MaxFunctionApps);
+                break;
+            }
         }
     }
 

tools/Azure.Mcp.Tools.Kusto/src/Commands/SampleCommand.cs

@@ -60,7 +60,10 @@ public sealed class SampleCommand(ILogger<SampleCommand> logger, IKustoService k
         try
         {
             List<JsonElement> results;
-            var query = $"{KustoService.EscapeKqlIdentifier(options.Table!)} | sample {options.Limit}";
+            // Validate limit is within safe bounds to prevent resource abuse
+            var safeLimit = Math.Clamp(options.Limit ?? 10, 1, 10000);
+
+            var query = $"{KustoService.EscapeKqlIdentifier(options.Table!)} | sample {safeLimit}";
 
             if (UseClusterUri(options))
             {

tools/Azure.Mcp.Tools.Kusto/src/Options/KustoOptionDefinitions.cs

@@ -49,7 +49,7 @@ public static class KustoOptionDefinitions
         $"--{LimitName}"
     )
     {
-        Description = "The maximum number of results to return.",
+        Description = "The maximum number of results to return. Must be a positive integer between 1 and 100000. Default is 10.",
         DefaultValueFactory = _ => 10,
         Required = true
     };

tools/Azure.Mcp.Tools.Kusto/src/Services/KustoService.cs

@@ -8,6 +8,7 @@ using Azure.Mcp.Core.Services.Azure.Subscription;
 using Azure.Mcp.Core.Services.Azure.Tenant;
 using Azure.Mcp.Core.Services.Caching;
 using Azure.Mcp.Tools.Kusto.Models;
+using Azure.Mcp.Tools.Kusto.Validation;
 using Microsoft.Extensions.Logging;
 
 namespace Azure.Mcp.Tools.Kusto.Services;
@@ -203,6 +204,10 @@ public sealed class KustoService(
             (nameof(databaseName), databaseName),
             (nameof(tableName), tableName));
 
+        // Validate table name to prevent KQL injection — while the query endpoint is read-only
+        // (no data modification), injection could still enable information disclosure or resource abuse
+        KustoIdentifierValidator.ValidateIdentifier(tableName, nameof(tableName));
+
         var kustoClient = await GetOrCreateKustoClientAsync(clusterUri, tenant, cancellationToken);
         var kustoResult = await kustoClient.ExecuteQueryCommandAsync(
             databaseName,

tools/Azure.Mcp.Tools.Kusto/src/Validation/KustoIdentifierValidator.cs

@@ -0,0 +1,59 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System.Text.RegularExpressions;
+using Microsoft.Mcp.Core.Commands;
+using Microsoft.Mcp.Core.Helpers;
+
+namespace Azure.Mcp.Tools.Kusto.Validation;
+
+/// <summary>
+/// Validates Kusto identifiers (table names, database names) to prevent KQL injection.
+/// Kusto identifiers follow specific naming rules:
+/// - Letters, digits, underscores, spaces, hyphens, and periods are allowed
+/// - Must start with a letter or underscore
+/// - Maximum length of 1024 characters
+/// This prevents injection attacks where a user-supplied table name could contain
+/// KQL operators like "|" to chain additional commands.
+/// </summary>
+internal static class KustoIdentifierValidator
+{
+    private const int MaxIdentifierLength = 1024;
+
+    // Kusto identifiers: letters, digits, underscores, spaces, hyphens, periods.
+    // Must start with a letter or underscore.
+    private static readonly Regex ValidIdentifierPattern = RegexHelper.CreateRegex(
+        @"^[A-Za-z_][A-Za-z0-9_ .\-]*$",
+        RegexOptions.Compiled);
+
+    /// <summary>
+    /// Validates that the given identifier is a safe Kusto table/database name.
+    /// Throws <see cref="CommandValidationException"/> for invalid input.
+    /// </summary>
+    public static void ValidateIdentifier(string? identifier, string parameterName)
+    {
+        if (string.IsNullOrWhiteSpace(identifier))
+        {
+            throw new CommandValidationException($"{parameterName} cannot be empty.");
+        }
+
+        if (identifier.Length > MaxIdentifierLength)
+        {
+            throw new CommandValidationException(
+                $"{parameterName} length exceeds maximum of {MaxIdentifierLength} characters.");
+        }
+
+        // Reject any KQL operators or injection characters
+        if (identifier.IndexOfAny(['|', ';', '(', ')', '{', '}', '[', ']', '<', '>', '\'', '"', '`', '/', '\\', '\n', '\r']) != -1)
+        {
+            throw new CommandValidationException(
+                $"{parameterName} contains invalid characters. Only letters, digits, underscores, spaces, hyphens, and periods are allowed.");
+        }
+
+        if (!ValidIdentifierPattern.IsMatch(identifier))
+        {
+            throw new CommandValidationException(
+                $"{parameterName} is not a valid Kusto identifier. Must start with a letter or underscore and contain only letters, digits, underscores, spaces, hyphens, and periods.");
+        }
+    }
+}

tools/Azure.Mcp.Tools.Monitor/src/Services/MonitorWebTestService.cs

@@ -2,7 +2,6 @@
 // Licensed under the MIT License.
 
 using System.Text.RegularExpressions;
-using Azure.Core;
 using Azure.Mcp.Core.Options;
 using Azure.Mcp.Core.Services.Azure;
 using Azure.Mcp.Core.Services.Azure.ResourceGroup;
@@ -12,6 +11,7 @@ using Azure.Mcp.Tools.Monitor.Models.WebTests;
 using Azure.ResourceManager.ApplicationInsights;
 using Azure.ResourceManager.ApplicationInsights.Models;
 using Microsoft.Extensions.Logging;
+using Microsoft.Mcp.Core.Helpers;
 
 namespace Azure.Mcp.Tools.Monitor.Services;
 
@@ -497,5 +497,5 @@ public class MonitorWebTestService(
         return hiddenLinkMatch?.Groups[1].Value;
     }
 
-    private readonly Regex AppInsightsComponentHiddenLinkTagRegex = new("^hidden-link:(\\/subscriptions\\/[^\\/]+\\/resourceGroups\\/[^\\/]+\\/providers\\/microsoft\\.insights\\/components\\/[^\\/]+)$", RegexOptions.Compiled | RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
+    private readonly Regex AppInsightsComponentHiddenLinkTagRegex = RegexHelper.CreateRegex("^hidden-link:(\\/subscriptions\\/[^\\/]+\\/resourceGroups\\/[^\\/]+\\/providers\\/microsoft\\.insights\\/components\\/[^\\/]+)$", RegexOptions.Compiled | RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
 }

tools/Azure.Mcp.Tools.Monitor/tests/Azure.Mcp.Tools.Monitor.LiveTests/MonitorCommandTests.cs

@@ -45,8 +45,8 @@ public sealed class MonitorCommandTests : RecordedCommandTestsBase
         _httpClientFactory = _httpClientProvider.GetRequiredService<IHttpClientFactory>();
         var tokenProvider = new PlaybackAwareTokenCredentialProvider(() => TestMode, NullLoggerFactory.Instance);
         var cloudConfiguration = new AzureCloudConfiguration(new ConfigurationBuilder().Build());
-        _tenantService = new TenantService(tokenProvider, cacheService, _httpClientFactory, cloudConfiguration);
-        var subscriptionService = new SubscriptionService(cacheService, _tenantService);
+        _tenantService = new TenantService(tokenProvider, cacheService, _httpClientFactory, cloudConfiguration, NullLogger<TenantService>.Instance);
+        var subscriptionService = new SubscriptionService(cacheService, _tenantService, NullLogger<SubscriptionService>.Instance);
         var resourceGroupService = new ResourceGroupService(cacheService, subscriptionService, _tenantService);
         var resourceResolverService = new ResourceResolverService(subscriptionService, _tenantService);
         _logger = NullLogger<MonitorService>.Instance;

tools/Azure.Mcp.Tools.MySql/src/Services/MySqlService.cs

@@ -10,6 +10,7 @@ using Azure.Mcp.Core.Services.Azure.Tenant;
 using Azure.Mcp.Tools.MySql.Commands;
 using Azure.ResourceManager.MySql.FlexibleServers;
 using Microsoft.Extensions.Logging;
+using Microsoft.Mcp.Core.Helpers;
 using MySqlConnector;
 
 namespace Azure.Mcp.Tools.MySql.Services;
@@ -69,15 +70,13 @@ public class MySqlService(IResourceGroupService resourceGroupService, ITenantSer
     ];
 
     // Pre-compiled regex patterns for word-boundary keyword matching
-    private static readonly Regex DangerousKeywordsPattern = new(
+    private static readonly Regex DangerousKeywordsPattern = RegexHelper.CreateRegex(
         @"\b(" + string.Join("|", DangerousKeywords.Select(Regex.Escape)) + @")\b",
-        RegexOptions.IgnoreCase | RegexOptions.Compiled,
-        TimeSpan.FromSeconds(3));
+        RegexOptions.IgnoreCase | RegexOptions.Compiled);
 
-    private static readonly Regex ObfuscationFunctionsPattern = new(
+    private static readonly Regex ObfuscationFunctionsPattern = RegexHelper.CreateRegex(
         @"\b(" + string.Join("|", ObfuscationFunctions.Select(Regex.Escape)) + @")\s*\(",
-        RegexOptions.IgnoreCase | RegexOptions.Compiled,
-        TimeSpan.FromSeconds(3));
+        RegexOptions.IgnoreCase | RegexOptions.Compiled);
 
     private async Task<string> GetEntraIdAccessTokenAsync(CancellationToken cancellationToken)
     {
@@ -174,7 +173,7 @@ public class MySqlService(IResourceGroupService resourceGroupService, ITenantSer
         // false positives (e.g., 'C#Developer' or 'foo--bar' are not comments).
         // The pattern handles both SQL-standard doubled quotes ('') and
         // MySQL's default backslash escaping (\') inside string literals.
-        var queryWithoutStrings = Regex.Replace(query, "'([^'\\\\]|\\\\.|'')*'", "'str'", RegexOptions.None, TimeSpan.FromSeconds(3));
+        var queryWithoutStrings = Regex.Replace(query, "'([^'\\\\]|\\\\.|'')*'", "'str'", RegexOptions.None, RegexHelper.DefaultRegexTimeout);
 
         // Reject queries containing SQL comments to prevent bypass attacks
         // (e.g., MySQL version-specific comments /*!50000 ... */ that are executed as code)
@@ -193,7 +192,7 @@ public class MySqlService(IResourceGroupService resourceGroupService, ITenantSer
         }
 
         // Regex pattern to detect multiple SQL statements (semicolon not at end)
-        var multipleStatementsPattern = new Regex(
+        var multipleStatementsPattern = RegexHelper.CreateRegex(
             @";\s*\w",
             RegexOptions.IgnoreCase | RegexOptions.Compiled
         );

tools/Azure.Mcp.Tools.Postgres/src/Services/PostgresService.cs

@@ -13,7 +13,6 @@ using Azure.Mcp.Core.Services.Azure.Tenant;
 using Azure.Mcp.Tools.Postgres.Auth;
 using Azure.Mcp.Tools.Postgres.Options;
 using Azure.Mcp.Tools.Postgres.Providers;
-using Azure.ResourceManager;
 using Azure.ResourceManager.PostgreSql.FlexibleServers;
 using Azure.ResourceManager.Resources;
 using Microsoft.Mcp.Core.Commands;

tools/Azure.Mcp.Tools.Postgres/src/Validation/SqlQueryValidator.cs

@@ -4,6 +4,7 @@
 using System.Net;
 using System.Text.RegularExpressions;
 using Microsoft.Mcp.Core.Commands;
+using Microsoft.Mcp.Core.Helpers;
 
 namespace Azure.Mcp.Tools.Postgres.Validation;
 
@@ -18,7 +19,6 @@ namespace Azure.Mcp.Tools.Postgres.Validation;
 internal static class SqlQueryValidator
 {
     private const int MaxQueryLength = 5000; // Arbitrary safety cap to avoid extremely large inputs.
-    private static readonly TimeSpan RegexTimeout = TimeSpan.FromSeconds(3); // 3 second timeout for regex operations
 
     /// <summary>
     /// SQL set-operation keywords that must be blocked to prevent data-exfiltration attacks
@@ -142,7 +142,7 @@ internal static class SqlQueryValidator
         // E-prefixed strings (E'...') additionally support backslash escapes (e.g., \').
         // The E-string pattern must appear first so the alternation matches it before
         // the standard pattern consumes the opening quote.
-        var withoutStrings = Regex.Replace(core, "[eE]'([^'\\\\]|\\\\.|'')*'|'([^']|'')*'", "'str'", RegexOptions.Compiled, RegexTimeout);
+        var withoutStrings = Regex.Replace(core, "[eE]'([^'\\\\]|\\\\.|'')*'|'([^']|'')*'", "'str'", RegexOptions.Compiled, RegexHelper.DefaultRegexTimeout);
 
         // Reject inline / block comments which can hide stacked statements or alter logic.
         if (withoutStrings.Contains("--", StringComparison.Ordinal) || withoutStrings.Contains("/*", StringComparison.Ordinal))
@@ -165,7 +165,7 @@ internal static class SqlQueryValidator
         }
 
         // Tokenize: capture word tokens (letters / underscore). Numerics & punctuation ignored.
-        var matches = Regex.Matches(withoutStrings, "[A-Za-z_]+", RegexOptions.Compiled, RegexTimeout);
+        var matches = Regex.Matches(withoutStrings, "[A-Za-z_]+", RegexOptions.Compiled, RegexHelper.DefaultRegexTimeout);
         if (matches.Count == 0)
         {
             throw new CommandValidationException("Query must contain a SELECT statement.", HttpStatusCode.BadRequest);

tools/Azure.Mcp.Tools.Redis/src/Services/RedisService.cs

@@ -16,10 +16,14 @@ using Azure.ResourceManager.Redis.Models;
 using Azure.ResourceManager.RedisEnterprise;
 using Azure.ResourceManager.Resources;
 using Azure.ResourceManager.Resources.Models;
+using Microsoft.Extensions.Logging;
 
 namespace Azure.Mcp.Tools.Redis.Services;
 
-public class RedisService(ISubscriptionService _subscriptionService, ITenantService _tenantService)
+public class RedisService(
+    ISubscriptionService _subscriptionService,
+    ITenantService _tenantService,
+    ILogger<RedisService> _logger)
     : BaseAzureService(_tenantService), IRedisService
 {
     public async Task<IEnumerable<Resource>> ListResourcesAsync(
@@ -38,12 +42,15 @@ public class RedisService(ISubscriptionService _subscriptionService, ITenantServ
 
         try
         {
-            resourcesTasks.Add(ListAcrResourcesAsync(subscriptionResource, cancellationToken));
-            resourcesTasks.Add(ListAmrResourcesAsync(subscriptionResource, cancellationToken));
+            resourcesTasks.Add(ListAcrResourcesAsync(subscriptionResource, _logger, cancellationToken));
+            resourcesTasks.Add(ListAmrResourcesAsync(subscriptionResource, _logger, cancellationToken));
             await Task.WhenAll(resourcesTasks);
         }
-        catch (Exception)
-        { }
+        catch (Exception ex) when (ex is not OperationCanceledException)
+        {
+            // Log individual task failures; partial results are still collected in the finally block.
+            _logger.LogWarning(ex, "One or more Redis resource listing tasks failed.");
+        }
         finally
         {
             foreach (var resourceTask in resourcesTasks)
@@ -148,7 +155,7 @@ public class RedisService(ISubscriptionService _subscriptionService, ITenantServ
         };
     }
 
-    private static async Task<IEnumerable<Resource>> ListAcrResourcesAsync(SubscriptionResource subscriptionResource, CancellationToken cancellationToken)
+    private static async Task<IEnumerable<Resource>> ListAcrResourcesAsync(SubscriptionResource subscriptionResource, ILogger<RedisService> logger, CancellationToken cancellationToken)
     {
         var resources = new List<Resource>();
 
@@ -182,7 +189,11 @@ public class RedisService(ISubscriptionService _subscriptionService, ITenantServ
                     });
                 }
             }
-            catch (Exception) { }
+            catch (Exception ex) when (ex is not OperationCanceledException)
+            {
+                // Access policy assignments may not be available for all cache types; continue with partial data.
+                logger.LogWarning(ex, "Failed to retrieve access policy assignments for {ResourceName}.", acrResource.Data.Name);
+            }
 
             resources.Add(new()
             {
@@ -252,7 +263,7 @@ public class RedisService(ISubscriptionService _subscriptionService, ITenantServ
         return resources;
     }
 
-    private static async Task<IEnumerable<Resource>> ListAmrResourcesAsync(SubscriptionResource subscriptionResource, CancellationToken cancellationToken)
+    private static async Task<IEnumerable<Resource>> ListAmrResourcesAsync(SubscriptionResource subscriptionResource, ILogger<RedisService> logger, CancellationToken cancellationToken)
     {
         var resources = new List<Resource>();
 
@@ -301,7 +312,11 @@ public class RedisService(ISubscriptionService _subscriptionService, ITenantServ
                     });
                 }
             }
-            catch (Exception) { }
+            catch (Exception ex) when (ex is not OperationCanceledException)
+            {
+                // Database listing may fail for some enterprise clusters; continue with partial data.
+                logger.LogWarning(ex, "Failed to retrieve databases for {ResourceName}.", resource.Name);
+            }
 
             resources.Add(new()
             {