Add max query-timeout validation and refactor MCP tests

Copilot · JerryNixon · Copilot · commit 79f7e3643096 · 2026-03-02T21:45:34.000Z
Co-authored-by: JerryNixon &lt;1749983+JerryNixon@users.noreply.github.com&gt;
diff --git a/schemas/dab.draft.schema.json b/schemas/dab.draft.schema.json
@@ -277,9 +277,10 @@
             },
             "query-timeout": {
               "type": "integer",
-              "description": "Execution timeout in seconds for MCP tool operations. Applies to all MCP tools.",
+              "description": "Execution timeout in seconds for MCP tool operations. Applies to all MCP tools. Maximum: 600 seconds.",
               "default": 30,
-              "minimum": 1
+              "minimum": 1,
+              "maximum": 600
             },
             "dml-tools": {
               "oneOf": [
diff --git a/src/Azure.DataApiBuilder.Mcp/BuiltInTools/AggregateRecordsTool.cs b/src/Azure.DataApiBuilder.Mcp/BuiltInTools/AggregateRecordsTool.cs
@@ -36,6 +36,7 @@ public class AggregateRecordsTool : IMcpTool
         public ToolType ToolType { get; } = ToolType.BuiltIn;
 
         private static readonly HashSet<string> _validFunctions = new(StringComparer.OrdinalIgnoreCase) { "count", "avg", "sum", "min", "max" };
+        private static readonly HashSet<string> _validHavingOperators = new(StringComparer.OrdinalIgnoreCase) { "eq", "neq", "gt", "gte", "lt", "lte", "in" };
 
         public Tool GetToolMetadata()
         {
@@ -247,18 +248,44 @@ public async Task<CallToolResult> ExecuteAsync(
                     havingOps = new Dictionary<string, double>(StringComparer.OrdinalIgnoreCase);
                     foreach (JsonProperty prop in havingEl.EnumerateObject())
                     {
+                        if (!_validHavingOperators.Contains(prop.Name))
+                        {
+                            return McpResponseBuilder.BuildErrorResult(
+                                toolName,
+                                "InvalidArguments",
+                                $"Unsupported having operator '{prop.Name}'. Supported operators are: eq, neq, gt, gte, lt, lte, in.",
+                                logger);
+                        }
+
                         if (prop.Name.Equals("in", StringComparison.OrdinalIgnoreCase) && prop.Value.ValueKind == JsonValueKind.Array)
                         {
                             havingIn = new List<double>();
                             foreach (JsonElement item in prop.Value.EnumerateArray())
                             {
+                                if (item.ValueKind != JsonValueKind.Number)
+                                {
+                                    return McpResponseBuilder.BuildErrorResult(
+                                        toolName,
+                                        "InvalidArguments",
+                                        "Argument 'having.in' must contain only numeric values.",
+                                        logger);
+                                }
+
                                 havingIn.Add(item.GetDouble());
                             }
                         }
                         else if (prop.Value.ValueKind == JsonValueKind.Number)
                         {
                             havingOps[prop.Name] = prop.Value.GetDouble();
                         }
+                        else
+                        {
+                            return McpResponseBuilder.BuildErrorResult(
+                                toolName,
+                                "InvalidArguments",
+                                $"Argument 'having.{prop.Name}' must be a numeric value.",
+                                logger);
+                        }
                     }
                 }
 
@@ -374,7 +401,7 @@ public async Task<CallToolResult> ExecuteAsync(
                 // Apply pagination if first is specified with groupby
                 if (first.HasValue && groupby.Count > 0)
                 {
-                    PaginationResult paginatedResult = ApplyPagination(aggregatedResults, first.Value, after);
+                    PaginationResult paginatedResult = ApplyPagination(aggregatedResults, first.Value, after, logger);
                     return McpResponseBuilder.BuildSuccessResult(
                         new Dictionary<string, object?>
                         {
@@ -578,7 +605,8 @@ internal sealed class PaginationResult
         internal static PaginationResult ApplyPagination(
             List<Dictionary<string, object?>> allResults,
             int first,
-            string? after)
+            string? after,
+            ILogger? logger = null)
         {
             int startIndex = 0;
 
@@ -595,7 +623,7 @@ internal static PaginationResult ApplyPagination(
                 }
                 catch (FormatException)
                 {
-                    // Invalid cursor format; start from beginning
+                    logger?.LogWarning("Invalid pagination cursor provided to aggregate_records; starting from beginning.");
                 }
             }
 
diff --git a/src/Azure.DataApiBuilder.Mcp/Utils/McpTelemetryHelper.cs b/src/Azure.DataApiBuilder.Mcp/Utils/McpTelemetryHelper.cs
@@ -69,6 +69,13 @@ public static async Task<CallToolResult> ExecuteWithTelemetryAsync(
                     timeoutSeconds = config.Runtime?.Mcp?.EffectiveQueryTimeoutSeconds ?? McpRuntimeOptions.DEFAULT_QUERY_TIMEOUT_SECONDS;
                 }
 
+                if (timeoutSeconds < 1 || timeoutSeconds > McpRuntimeOptions.MAX_QUERY_TIMEOUT_SECONDS)
+                {
+                    throw new InvalidOperationException(
+                        $"MCP query-timeout must be within 1 and {McpRuntimeOptions.MAX_QUERY_TIMEOUT_SECONDS} seconds. " +
+                        $"Provided value: {timeoutSeconds}.");
+                }
+
                 // Wrap tool execution with the configured timeout using a linked CancellationTokenSource.
                 using CancellationTokenSource timeoutCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
                 timeoutCts.CancelAfter(TimeSpan.FromSeconds(timeoutSeconds));
diff --git a/src/Config/ObjectModel/McpRuntimeOptions.cs b/src/Config/ObjectModel/McpRuntimeOptions.cs
@@ -11,6 +11,7 @@ public record McpRuntimeOptions
 {
     public const string DEFAULT_PATH = "/mcp";
     public const int DEFAULT_QUERY_TIMEOUT_SECONDS = 30;
+    public const int MAX_QUERY_TIMEOUT_SECONDS = 600;
 
     /// <summary>
     /// Whether MCP endpoints are enabled
@@ -41,7 +42,7 @@ public record McpRuntimeOptions
     /// Execution timeout in seconds for MCP tool operations.
     /// This timeout wraps the entire tool execution including database queries.
     /// It applies to all MCP tools, not just aggregation.
-    /// Default: 30 seconds.
+    /// Default: 30 seconds. Maximum: 600 seconds.
     /// </summary>
     [JsonPropertyName("query-timeout")]
     public int? QueryTimeout { get; init; }
diff --git a/src/Core/Configurations/RuntimeConfigValidator.cs b/src/Core/Configurations/RuntimeConfigValidator.cs
@@ -916,11 +916,13 @@ public void ValidateMcpUri(RuntimeConfig runtimeConfig)
         }
 
         // Validate query-timeout if provided
-        if (runtimeConfig.Runtime.Mcp.QueryTimeout is not null && runtimeConfig.Runtime.Mcp.QueryTimeout < 1)
+        if (runtimeConfig.Runtime.Mcp.QueryTimeout is not null &&
+            (runtimeConfig.Runtime.Mcp.QueryTimeout < 1 ||
+             runtimeConfig.Runtime.Mcp.QueryTimeout > McpRuntimeOptions.MAX_QUERY_TIMEOUT_SECONDS))
         {
             HandleOrRecordException(new DataApiBuilderException(
-                message: "MCP query-timeout must be a positive integer (>= 1 second). " +
-                         $"Provided value: {runtimeConfig.Runtime.Mcp.QueryTimeout}.",
+                message: $"MCP query-timeout must be within 1 and {McpRuntimeOptions.MAX_QUERY_TIMEOUT_SECONDS} seconds. " +
+                          $"Provided value: {runtimeConfig.Runtime.Mcp.QueryTimeout}.",
                 statusCode: HttpStatusCode.ServiceUnavailable,
                 subStatusCode: DataApiBuilderException.SubStatusCodes.ConfigValidationError));
         }
diff --git a/src/Service.Tests/Mcp/AggregateRecordsToolTests.cs b/src/Service.Tests/Mcp/AggregateRecordsToolTests.cs
@@ -83,30 +83,15 @@ public void GetToolMetadata_HasInputSchema()
 
         #region Configuration Tests
 
-        [TestMethod]
-        public async Task AggregateRecords_DisabledAtRuntimeLevel_ReturnsToolDisabledError()
-        {
-            RuntimeConfig config = CreateConfig(aggregateRecordsEnabled: false);
-            IServiceProvider sp = CreateServiceProvider(config);
-            AggregateRecordsTool tool = new();
-
-            JsonDocument args = JsonDocument.Parse("{\"entity\": \"Book\", \"function\": \"count\", \"field\": \"*\"}");
-            CallToolResult result = await tool.ExecuteAsync(args, sp, CancellationToken.None);
-
-            Assert.IsTrue(result.IsError == true);
-            JsonElement content = ParseContent(result);
-            AssertToolDisabledError(content);
-        }
-
-        [TestMethod]
-        public async Task AggregateRecords_DisabledAtEntityLevel_ReturnsToolDisabledError()
+        [DataTestMethod]
+        [DataRow(false, false, DisplayName = "Disabled at runtime level returns ToolDisabled")]
+        [DataRow(true, true, DisplayName = "Disabled at entity level returns ToolDisabled")]
+        public async Task AggregateRecords_DisabledConfigurations_ReturnToolDisabledError(bool runtimeEnabled, bool disableAtEntityLevel)
         {
-            RuntimeConfig config = CreateConfigWithEntityDmlDisabled();
-            IServiceProvider sp = CreateServiceProvider(config);
-            AggregateRecordsTool tool = new();
-
-            JsonDocument args = JsonDocument.Parse("{\"entity\": \"Book\", \"function\": \"count\", \"field\": \"*\"}");
-            CallToolResult result = await tool.ExecuteAsync(args, sp, CancellationToken.None);
+            RuntimeConfig config = disableAtEntityLevel
+                ? CreateConfigWithEntityDmlDisabled()
+                : CreateConfig(aggregateRecordsEnabled: runtimeEnabled);
+            CallToolResult result = await ExecuteWithConfigAsync(config, "{\"entity\": \"Book\", \"function\": \"count\", \"field\": \"*\"}");
 
             Assert.IsTrue(result.IsError == true);
             JsonElement content = ParseContent(result);
@@ -117,75 +102,28 @@ public async Task AggregateRecords_DisabledAtEntityLevel_ReturnsToolDisabledErro
 
         #region Input Validation Tests
 
-        [TestMethod]
-        public async Task AggregateRecords_NullArguments_ReturnsInvalidArguments()
+        [DataTestMethod]
+        [DataRow(null, null, DisplayName = "Null arguments return InvalidArguments")]
+        [DataRow("{\"function\": \"count\", \"field\": \"*\"}", null, DisplayName = "Missing entity returns InvalidArguments")]
+        [DataRow("{\"entity\": \"Book\", \"field\": \"*\"}", null, DisplayName = "Missing function returns InvalidArguments")]
+        [DataRow("{\"entity\": \"Book\", \"function\": \"count\"}", null, DisplayName = "Missing field returns InvalidArguments")]
+        [DataRow("{\"entity\": \"Book\", \"function\": \"median\", \"field\": \"price\"}", "median", DisplayName = "Invalid function returns InvalidArguments")]
+        [DataRow("{\"entity\": \"Book\", \"function\": \"count\", \"field\": \"*\", \"having\": { \"between\": 10 }}", "Unsupported having operator", DisplayName = "Unsupported having operator returns InvalidArguments")]
+        [DataRow("{\"entity\": \"Book\", \"function\": \"count\", \"field\": \"*\", \"having\": { \"eq\": \"10\" }}", "must be a numeric value", DisplayName = "Non-numeric having value returns InvalidArguments")]
+        [DataRow("{\"entity\": \"Book\", \"function\": \"count\", \"field\": \"*\", \"having\": { \"in\": [1, \"x\"] }}", "must contain only numeric values", DisplayName = "Non-numeric having.in value returns InvalidArguments")]
+        public async Task AggregateRecords_InvalidArguments_ReturnInvalidArguments(string? argsJson, string? expectedMessageFragment)
         {
             RuntimeConfig config = CreateConfig();
-            IServiceProvider sp = CreateServiceProvider(config);
-            AggregateRecordsTool tool = new();
+            CallToolResult result = await ExecuteWithConfigAsync(config, argsJson);
 
-            CallToolResult result = await tool.ExecuteAsync(null, sp, CancellationToken.None);
-            Assert.IsTrue(result.IsError == true);
-            JsonElement content = ParseContent(result);
-            Assert.IsTrue(content.TryGetProperty("error", out JsonElement error));
-            Assert.AreEqual("InvalidArguments", error.GetProperty("type").GetString());
-        }
-
-        [TestMethod]
-        public async Task AggregateRecords_MissingEntity_ReturnsInvalidArguments()
-        {
-            RuntimeConfig config = CreateConfig();
-            IServiceProvider sp = CreateServiceProvider(config);
-            AggregateRecordsTool tool = new();
-
-            JsonDocument args = JsonDocument.Parse("{\"function\": \"count\", \"field\": \"*\"}");
-            CallToolResult result = await tool.ExecuteAsync(args, sp, CancellationToken.None);
             Assert.IsTrue(result.IsError == true);
             JsonElement content = ParseContent(result);
             Assert.AreEqual("InvalidArguments", content.GetProperty("error").GetProperty("type").GetString());
-        }
 
-        [TestMethod]
-        public async Task AggregateRecords_MissingFunction_ReturnsInvalidArguments()
-        {
-            RuntimeConfig config = CreateConfig();
-            IServiceProvider sp = CreateServiceProvider(config);
-            AggregateRecordsTool tool = new();
-
-            JsonDocument args = JsonDocument.Parse("{\"entity\": \"Book\", \"field\": \"*\"}");
-            CallToolResult result = await tool.ExecuteAsync(args, sp, CancellationToken.None);
-            Assert.IsTrue(result.IsError == true);
-            JsonElement content = ParseContent(result);
-            Assert.AreEqual("InvalidArguments", content.GetProperty("error").GetProperty("type").GetString());
-        }
-
-        [TestMethod]
-        public async Task AggregateRecords_MissingField_ReturnsInvalidArguments()
-        {
-            RuntimeConfig config = CreateConfig();
-            IServiceProvider sp = CreateServiceProvider(config);
-            AggregateRecordsTool tool = new();
-
-            JsonDocument args = JsonDocument.Parse("{\"entity\": \"Book\", \"function\": \"count\"}");
-            CallToolResult result = await tool.ExecuteAsync(args, sp, CancellationToken.None);
-            Assert.IsTrue(result.IsError == true);
-            JsonElement content = ParseContent(result);
-            Assert.AreEqual("InvalidArguments", content.GetProperty("error").GetProperty("type").GetString());
-        }
-
-        [TestMethod]
-        public async Task AggregateRecords_InvalidFunction_ReturnsInvalidArguments()
-        {
-            RuntimeConfig config = CreateConfig();
-            IServiceProvider sp = CreateServiceProvider(config);
-            AggregateRecordsTool tool = new();
-
-            JsonDocument args = JsonDocument.Parse("{\"entity\": \"Book\", \"function\": \"median\", \"field\": \"price\"}");
-            CallToolResult result = await tool.ExecuteAsync(args, sp, CancellationToken.None);
-            Assert.IsTrue(result.IsError == true);
-            JsonElement content = ParseContent(result);
-            Assert.AreEqual("InvalidArguments", content.GetProperty("error").GetProperty("type").GetString());
-            Assert.IsTrue(content.GetProperty("error").GetProperty("message").GetString()!.Contains("median"));
+            if (!string.IsNullOrWhiteSpace(expectedMessageFragment))
+            {
+                Assert.IsTrue(content.GetProperty("error").GetProperty("message").GetString()!.Contains(expectedMessageFragment));
+            }
         }
 
         #endregion
@@ -1194,6 +1132,14 @@ private static JsonElement ParseContent(CallToolResult result)
             return JsonDocument.Parse(firstContent.Text).RootElement;
         }
 
+        private static async Task<CallToolResult> ExecuteWithConfigAsync(RuntimeConfig config, string? argsJson, CancellationToken cancellationToken = default)
+        {
+            IServiceProvider serviceProvider = CreateServiceProvider(config);
+            AggregateRecordsTool tool = new();
+            JsonDocument? args = argsJson is null ? null : JsonDocument.Parse(argsJson);
+            return await tool.ExecuteAsync(args, serviceProvider, cancellationToken);
+        }
+
         private static void AssertToolDisabledError(JsonElement content)
         {
             Assert.IsTrue(content.TryGetProperty("error", out JsonElement error));
diff --git a/src/Service.Tests/Mcp/McpQueryTimeoutTests.cs b/src/Service.Tests/Mcp/McpQueryTimeoutTests.cs

Original file line number	Diff line number	Diff line change
`@@ -36,6 +36,7 @@ public class AggregateRecordsTool : IMcpTool`
`36`	`36`	`public ToolType ToolType { get; } = ToolType.BuiltIn;`
`37`	`37`
`38`	`38`	`private static readonly HashSet<string> _validFunctions = new(StringComparer.OrdinalIgnoreCase) { "count", "avg", "sum", "min", "max" };`
	`39`	`+ private static readonly HashSet<string> _validHavingOperators = new(StringComparer.OrdinalIgnoreCase) { "eq", "neq", "gt", "gte", "lt", "lte", "in" };`
`39`	`40`
`40`	`41`	`public Tool GetToolMetadata()`
`41`	`42`	`{`
`@@ -247,18 +248,44 @@ public async Task<CallToolResult> ExecuteAsync(`
`247`	`248`	`havingOps = new Dictionary<string, double>(StringComparer.OrdinalIgnoreCase);`
`248`	`249`	`foreach (JsonProperty prop in havingEl.EnumerateObject())`
`249`	`250`	`{`
	`251`	`+ if (!_validHavingOperators.Contains(prop.Name))`
	`252`	`+ {`
	`253`	`+ return McpResponseBuilder.BuildErrorResult(`
	`254`	`+ toolName,`
	`255`	`+ "InvalidArguments",`
	`256`	`+ $"Unsupported having operator '{prop.Name}'. Supported operators are: eq, neq, gt, gte, lt, lte, in.",`
	`257`	`+ logger);`
	`258`	`+ }`
	`259`	`+`
`250`	`260`	`if (prop.Name.Equals("in", StringComparison.OrdinalIgnoreCase) && prop.Value.ValueKind == JsonValueKind.Array)`
`251`	`261`	`{`
`252`	`262`	`havingIn = new List<double>();`
`253`	`263`	`foreach (JsonElement item in prop.Value.EnumerateArray())`
`254`	`264`	`{`
	`265`	`+ if (item.ValueKind != JsonValueKind.Number)`
	`266`	`+ {`
	`267`	`+ return McpResponseBuilder.BuildErrorResult(`
	`268`	`+ toolName,`
	`269`	`+ "InvalidArguments",`
	`270`	`+ "Argument 'having.in' must contain only numeric values.",`
	`271`	`+ logger);`
	`272`	`+ }`
	`273`	`+`
`255`	`274`	`havingIn.Add(item.GetDouble());`
`256`	`275`	`}`
`257`	`276`	`}`
`258`	`277`	`else if (prop.Value.ValueKind == JsonValueKind.Number)`
`259`	`278`	`{`
`260`	`279`	`havingOps[prop.Name] = prop.Value.GetDouble();`
`261`	`280`	`}`
	`281`	`+ else`
	`282`	`+ {`
	`283`	`+ return McpResponseBuilder.BuildErrorResult(`
	`284`	`+ toolName,`
	`285`	`+ "InvalidArguments",`
	`286`	`+ $"Argument 'having.{prop.Name}' must be a numeric value.",`
	`287`	`+ logger);`
	`288`	`+ }`
`262`	`289`	`}`
`263`	`290`	`}`
`264`	`291`
`@@ -374,7 +401,7 @@ public async Task<CallToolResult> ExecuteAsync(`
`374`	`401`	`// Apply pagination if first is specified with groupby`
`375`	`402`	`if (first.HasValue && groupby.Count > 0)`
`376`	`403`	`{`
`377`		`- PaginationResult paginatedResult = ApplyPagination(aggregatedResults, first.Value, after);`
	`404`	`+ PaginationResult paginatedResult = ApplyPagination(aggregatedResults, first.Value, after, logger);`
`378`	`405`	`return McpResponseBuilder.BuildSuccessResult(`
`379`	`406`	`new Dictionary<string, object?>`
`380`	`407`	`{`
`@@ -578,7 +605,8 @@ internal sealed class PaginationResult`
`578`	`605`	`internal static PaginationResult ApplyPagination(`
`579`	`606`	`List<Dictionary<string, object?>> allResults,`
`580`	`607`	`int first,`
`581`		`- string? after)`
	`608`	`+ string? after,`
	`609`	`+ ILogger? logger = null)`
`582`	`610`	`{`
`583`	`611`	`int startIndex = 0;`
`584`	`612`
`@@ -595,7 +623,7 @@ internal static PaginationResult ApplyPagination(`
`595`	`623`	`}`
`596`	`624`	`catch (FormatException)`
`597`	`625`	`{`
`598`		`- // Invalid cursor format; start from beginning`
	`626`	`+ logger?.LogWarning("Invalid pagination cursor provided to aggregate_records; starting from beginning.");`
`599`	`627`	`}`
`600`	`628`	`}`
`601`	`629`
Original file line number	Diff line number	Diff line change
`@@ -916,11 +916,13 @@ public void ValidateMcpUri(RuntimeConfig runtimeConfig)`
`916`	`916`	`}`
`917`	`917`
`918`	`918`	`// Validate query-timeout if provided`
`919`		`- if (runtimeConfig.Runtime.Mcp.QueryTimeout is not null && runtimeConfig.Runtime.Mcp.QueryTimeout < 1)`
	`919`	`+ if (runtimeConfig.Runtime.Mcp.QueryTimeout is not null &&`
	`920`	`+ (runtimeConfig.Runtime.Mcp.QueryTimeout < 1 \|\|`
	`921`	`+ runtimeConfig.Runtime.Mcp.QueryTimeout > McpRuntimeOptions.MAX_QUERY_TIMEOUT_SECONDS))`
`920`	`922`	`{`
`921`	`923`	`HandleOrRecordException(new DataApiBuilderException(`
`922`		`- message: "MCP query-timeout must be a positive integer (>= 1 second). " +`
`923`		`- $"Provided value: {runtimeConfig.Runtime.Mcp.QueryTimeout}.",`
	`924`	`+ message: $"MCP query-timeout must be within 1 and {McpRuntimeOptions.MAX_QUERY_TIMEOUT_SECONDS} seconds. " +`
	`925`	`+ $"Provided value: {runtimeConfig.Runtime.Mcp.QueryTimeout}.",`
`924`	`926`	`statusCode: HttpStatusCode.ServiceUnavailable,`
`925`	`927`	`subStatusCode: DataApiBuilderException.SubStatusCodes.ConfigValidationError));`
`926`	`928`	`}`