[MCP] aggregate_records DML tool

schemas/dab.draft.schema.json

@@ -315,6 +315,34 @@
                       "type": "boolean",
                       "description": "Enable/disable the execute-entity tool.",
                       "default": false
+                    },
+                    "aggregate-records": {
+                      "oneOf": [
+                        {
+                          "type": "boolean",
+                          "description": "Enable/disable the aggregate-records tool."
+                        },
+                        {
+                          "type": "object",
+                          "description": "Aggregate records tool configuration",
+                          "additionalProperties": false,
+                          "properties": {
+                            "enabled": {
+                              "type": "boolean",
+                              "description": "Enable/disable the aggregate-records tool.",
+                              "default": true
+                            },
+                            "query-timeout": {
+                              "type": "integer",
+                              "description": "Execution timeout in seconds for aggregate queries. Range: 1-600.",
+                              "default": 30,
+                              "minimum": 1,
+                              "maximum": 600
+                            }
+                          }
+                        }
+                      ],
+                      "default": false
                     }
                   }
                 }

src/Azure.DataApiBuilder.Mcp/Utils/McpErrorHelpers.cs

@@ -24,5 +24,16 @@ public static CallToolResult ToolDisabled(string toolName, ILogger? logger, stri
             string message = customMessage ?? $"The {toolName} tool is disabled in the configuration.";
             return McpResponseBuilder.BuildErrorResult(toolName, Model.McpErrorCode.ToolDisabled.ToString(), message, logger);
         }
+
+        /// <summary>
+        /// Returns a model-friendly error when a field name is not found for an entity.
+        /// Guides the model to call describe_entities to discover valid field names.
+        /// </summary>
+        public static CallToolResult FieldNotFound(string toolName, string entityName, string fieldName, string parameterName, ILogger? logger)
+        {
+            string message = $"Field '{fieldName}' in '{parameterName}' was not found for entity '{entityName}'. "
+                + $"Call describe_entities to get valid field names for '{entityName}'.";
+            return McpResponseBuilder.BuildErrorResult(toolName, "FieldNotFound", message, logger);
+        }
     }
 }

src/Azure.DataApiBuilder.Mcp/Utils/McpTelemetryErrorCodes.cs

@@ -37,5 +37,10 @@ internal static class McpTelemetryErrorCodes
         /// Operation cancelled error code.
         /// </summary>
         public const string OPERATION_CANCELLED = "OperationCancelled";
+
+        /// <summary>
+        /// Operation timed out error code.
+        /// </summary>
+        public const string OPERATION_TIMEOUT = "OperationTimeout";
     }
 }

src/Azure.DataApiBuilder.Mcp/Utils/McpTelemetryHelper.cs

@@ -60,7 +60,6 @@ public static async Task<CallToolResult> ExecuteWithTelemetryAsync(
                     operation: operation,
                     dbProcedure: dbProcedure);
 
-                // Execute the tool
                 CallToolResult result = await tool.ExecuteAsync(arguments, serviceProvider, cancellationToken);
 
                 // Check if the tool returned an error result (tools catch exceptions internally
@@ -124,6 +123,7 @@ public static string InferOperationFromTool(IMcpTool tool, string toolName)
                 "delete_record" => "delete",
                 "describe_entities" => "describe",
                 "execute_entity" => "execute",
+                "aggregate_records" => "aggregate",
                 _ => "execute" // Fallback for any unknown built-in tools
             };
         }
@@ -188,6 +188,7 @@ public static string MapExceptionToErrorCode(Exception ex)
             return ex switch
             {
                 OperationCanceledException => McpTelemetryErrorCodes.OPERATION_CANCELLED,
+                TimeoutException => McpTelemetryErrorCodes.OPERATION_TIMEOUT,
                 DataApiBuilderException dabEx when dabEx.SubStatusCode == DataApiBuilderException.SubStatusCodes.AuthenticationChallenge
                     => McpTelemetryErrorCodes.AUTHENTICATION_FAILED,
                 DataApiBuilderException dabEx when dabEx.SubStatusCode == DataApiBuilderException.SubStatusCodes.AuthorizationCheckFailed

src/Cli.Tests/InitTests.cs

@@ -493,6 +493,74 @@ public Task VerifyCorrectConfigGenerationWithMultipleMutationOptions(DatabaseTyp
             return ExecuteVerifyTest(options, verifySettings);
         }
 
+        /// <summary>
+        /// Test that init with/without --mcp.aggregate-records.query-timeout produces a config
+        /// with the correct aggregate-records query-timeout in the DmlTools section.
+        /// When null (not specified), defaults to 30 seconds. When provided, the config reflects the value.
+        /// </summary>
+        [DataTestMethod]
+        [DataRow(null, false, DmlToolsConfig.DEFAULT_QUERY_TIMEOUT_SECONDS, DisplayName = "Init without query-timeout uses default 30s")]
+        [DataRow(1, true, 1, DisplayName = "Init with query-timeout 1s (minimum)")]
+        [DataRow(120, true, 120, DisplayName = "Init with query-timeout 120s")]
+        [DataRow(600, true, 600, DisplayName = "Init with query-timeout 600s (maximum)")]
+        public void InitWithAggregateRecordsQueryTimeout_SetsOrDefaultsTimeout(int? inputTimeout, bool expectedUserProvided, int expectedEffectiveTimeout)
+        {
+            InitOptions options = new(
+                databaseType: DatabaseType.MSSQL,
+                connectionString: "testconnectionstring",
+                cosmosNoSqlDatabase: null,
+                cosmosNoSqlContainer: null,
+                graphQLSchemaPath: null,
+                setSessionContext: false,
+                hostMode: HostMode.Development,
+                corsOrigin: null,
+                authenticationProvider: EasyAuthType.AppService.ToString(),
+                mcpAggregateRecordsQueryTimeout: inputTimeout,
+                config: TEST_RUNTIME_CONFIG_FILE);
+
+            Assert.IsTrue(TryCreateRuntimeConfig(options, _runtimeConfigLoader!, _fileSystem!, out RuntimeConfig? runtimeConfig));
+            Assert.IsNotNull(runtimeConfig?.Runtime?.Mcp?.DmlTools);
+            Assert.AreEqual(inputTimeout, runtimeConfig.Runtime.Mcp.DmlTools.AggregateRecordsQueryTimeout);
+            Assert.AreEqual(expectedUserProvided, runtimeConfig.Runtime.Mcp.DmlTools.UserProvidedAggregateRecordsQueryTimeout);
+            Assert.AreEqual(expectedEffectiveTimeout, runtimeConfig.Runtime.Mcp.DmlTools.EffectiveAggregateRecordsQueryTimeoutSeconds);
+        }
+
+        /// <summary>
+        /// Test that init with --mcp.aggregate-records.query-timeout produces valid JSON
+        /// that round-trips correctly through serialization/deserialization.
+        /// </summary>
+        [TestMethod]
+        public void InitWithAggregateRecordsQueryTimeout_RoundTripsCorrectly()
+        {
+            InitOptions options = new(
+                databaseType: DatabaseType.MSSQL,
+                connectionString: "testconnectionstring",
+                cosmosNoSqlDatabase: null,
+                cosmosNoSqlContainer: null,
+                graphQLSchemaPath: null,
+                setSessionContext: false,
+                hostMode: HostMode.Development,
+                corsOrigin: null,
+                authenticationProvider: EasyAuthType.AppService.ToString(),
+                mcpAggregateRecordsQueryTimeout: 90,
+                config: TEST_RUNTIME_CONFIG_FILE);
+
+            Assert.IsTrue(TryCreateRuntimeConfig(options, _runtimeConfigLoader!, _fileSystem!, out RuntimeConfig? runtimeConfig));
+
+            // Serialize to JSON and deserialize back
+            JsonSerializerOptions serializerOptions = RuntimeConfigLoader.GetSerializationOptions();
+            string json = JsonSerializer.Serialize(runtimeConfig, serializerOptions);
+            RuntimeConfig? deserialized = JsonSerializer.Deserialize<RuntimeConfig>(json, serializerOptions);
+
+            Assert.IsNotNull(deserialized?.Runtime?.Mcp?.DmlTools);
+            Assert.AreEqual(90, deserialized.Runtime.Mcp.DmlTools.AggregateRecordsQueryTimeout);
+            Assert.AreEqual(90, deserialized.Runtime.Mcp.DmlTools.EffectiveAggregateRecordsQueryTimeoutSeconds);
+
+            // Verify the JSON contains the object format for aggregate-records
+            Assert.IsTrue(json.Contains("\"query-timeout\""), $"Expected 'query-timeout' in serialized JSON. Got: {json}");
+            Assert.IsTrue(json.Contains("90"), $"Expected timeout value 90 in serialized JSON. Got: {json}");
+        }
+
         private Task ExecuteVerifyTest(InitOptions options, VerifySettings? settings = null)
         {
             Assert.IsTrue(TryCreateRuntimeConfig(options, _runtimeConfigLoader!, _fileSystem!, out RuntimeConfig? runtimeConfig));

src/Cli.Tests/Snapshots/EndToEndTests.TestAddingStoredProcedureWithRestMethodsAndGraphQLOperations.verified.txt

@@ -27,13 +27,17 @@
         UpdateRecord: true,
         DeleteRecord: true,
         ExecuteEntity: true,
+        AggregateRecords: true,
         UserProvidedAllTools: false,
         UserProvidedDescribeEntities: false,
         UserProvidedCreateRecord: false,
         UserProvidedReadRecords: false,
         UserProvidedUpdateRecord: false,
         UserProvidedDeleteRecord: false,
-        UserProvidedExecuteEntity: false
+        UserProvidedExecuteEntity: false,
+        UserProvidedAggregateRecords: false,
+        UserProvidedAggregateRecordsQueryTimeout: false,
+        EffectiveAggregateRecordsQueryTimeoutSeconds: 30
       }
     },
     Host: {

src/Cli.Tests/Snapshots/EndToEndTests.TestConfigGeneratedAfterAddingEntityWithSourceAsStoredProcedure.verified.txt

@@ -27,13 +27,17 @@
         UpdateRecord: true,
         DeleteRecord: true,
         ExecuteEntity: true,
+        AggregateRecords: true,
         UserProvidedAllTools: false,
         UserProvidedDescribeEntities: false,
         UserProvidedCreateRecord: false,
         UserProvidedReadRecords: false,
         UserProvidedUpdateRecord: false,
         UserProvidedDeleteRecord: false,
-        UserProvidedExecuteEntity: false
+        UserProvidedExecuteEntity: false,
+        UserProvidedAggregateRecords: false,
+        UserProvidedAggregateRecordsQueryTimeout: false,
+        EffectiveAggregateRecordsQueryTimeoutSeconds: 30
       }
     },
     Host: {

src/Cli.Tests/Snapshots/EndToEndTests.TestConfigGeneratedAfterAddingEntityWithSourceWithDefaultType.verified.txt

@@ -27,13 +27,17 @@
         UpdateRecord: true,
         DeleteRecord: true,
         ExecuteEntity: true,
+        AggregateRecords: true,
         UserProvidedAllTools: false,
         UserProvidedDescribeEntities: false,
         UserProvidedCreateRecord: false,
         UserProvidedReadRecords: false,
         UserProvidedUpdateRecord: false,
         UserProvidedDeleteRecord: false,
-        UserProvidedExecuteEntity: false
+        UserProvidedExecuteEntity: false,
+        UserProvidedAggregateRecords: false,
+        UserProvidedAggregateRecordsQueryTimeout: false,
+        EffectiveAggregateRecordsQueryTimeoutSeconds: 30
       }
     },
     Host: {

src/Cli.Tests/Snapshots/EndToEndTests.TestConfigGeneratedAfterAddingEntityWithoutIEnumerables.verified.txt

@@ -27,13 +27,17 @@
         UpdateRecord: true,
         DeleteRecord: true,
         ExecuteEntity: true,
+        AggregateRecords: true,
         UserProvidedAllTools: false,
         UserProvidedDescribeEntities: false,
         UserProvidedCreateRecord: false,
         UserProvidedReadRecords: false,
         UserProvidedUpdateRecord: false,
         UserProvidedDeleteRecord: false,
-        UserProvidedExecuteEntity: false
+        UserProvidedExecuteEntity: false,
+        UserProvidedAggregateRecords: false,
+        UserProvidedAggregateRecordsQueryTimeout: false,
+        EffectiveAggregateRecordsQueryTimeoutSeconds: 30
       }
     },
     Host: {

src/Cli.Tests/Snapshots/EndToEndTests.TestInitForCosmosDBNoSql.verified.txt

@@ -28,13 +28,17 @@
         UpdateRecord: true,
         DeleteRecord: true,
         ExecuteEntity: true,
+        AggregateRecords: true,
         UserProvidedAllTools: false,
         UserProvidedDescribeEntities: false,
         UserProvidedCreateRecord: false,
         UserProvidedReadRecords: false,
         UserProvidedUpdateRecord: false,
         UserProvidedDeleteRecord: false,
-        UserProvidedExecuteEntity: false
+        UserProvidedExecuteEntity: false,
+        UserProvidedAggregateRecords: false,
+        UserProvidedAggregateRecordsQueryTimeout: false,
+        EffectiveAggregateRecordsQueryTimeoutSeconds: 30
       }
     },
     Host: {

src/Cli.Tests/Snapshots/EndToEndTests.TestUpdatingStoredProcedureWithRestMethods.verified.txt

@@ -1,4 +1,4 @@
-{
+{
   DataSource: {
     DatabaseType: MSSQL,
     Options: {
@@ -27,13 +27,17 @@
         UpdateRecord: true,
         DeleteRecord: true,
         ExecuteEntity: true,
+        AggregateRecords: true,
         UserProvidedAllTools: false,
         UserProvidedDescribeEntities: false,
         UserProvidedCreateRecord: false,
         UserProvidedReadRecords: false,
         UserProvidedUpdateRecord: false,
         UserProvidedDeleteRecord: false,
-        UserProvidedExecuteEntity: false
+        UserProvidedExecuteEntity: false,
+        UserProvidedAggregateRecords: false,
+        UserProvidedAggregateRecordsQueryTimeout: false,
+        EffectiveAggregateRecordsQueryTimeoutSeconds: 30
       }
     },
     Host: {

src/Cli.Tests/Snapshots/EndToEndTests.TestUpdatingStoredProcedureWithRestMethodsAndGraphQLOperations.verified.txt

@@ -27,13 +27,17 @@
         UpdateRecord: true,
         DeleteRecord: true,
         ExecuteEntity: true,
+        AggregateRecords: true,
         UserProvidedAllTools: false,
         UserProvidedDescribeEntities: false,
         UserProvidedCreateRecord: false,
         UserProvidedReadRecords: false,
         UserProvidedUpdateRecord: false,
         UserProvidedDeleteRecord: false,
-        UserProvidedExecuteEntity: false
+        UserProvidedExecuteEntity: false,
+        UserProvidedAggregateRecords: false,
+        UserProvidedAggregateRecordsQueryTimeout: false,
+        EffectiveAggregateRecordsQueryTimeoutSeconds: 30
       }
     },
     Host: {

src/Cli.Tests/Snapshots/InitTests.CosmosDbNoSqlDatabase.verified.txt

@@ -28,13 +28,17 @@
         UpdateRecord: true,
         DeleteRecord: true,
         ExecuteEntity: true,
+        AggregateRecords: true,
         UserProvidedAllTools: false,
         UserProvidedDescribeEntities: false,
         UserProvidedCreateRecord: false,
         UserProvidedReadRecords: false,
         UserProvidedUpdateRecord: false,
         UserProvidedDeleteRecord: false,
-        UserProvidedExecuteEntity: false
+        UserProvidedExecuteEntity: false,
+        UserProvidedAggregateRecords: false,
+        UserProvidedAggregateRecordsQueryTimeout: false,
+        EffectiveAggregateRecordsQueryTimeoutSeconds: 30
       }
     },
     Host: {

src/Cli.Tests/Snapshots/InitTests.CosmosDbPostgreSqlDatabase.verified.txt

@@ -24,13 +24,17 @@
         UpdateRecord: true,
         DeleteRecord: true,
         ExecuteEntity: true,
+        AggregateRecords: true,
         UserProvidedAllTools: false,
         UserProvidedDescribeEntities: false,
         UserProvidedCreateRecord: false,
         UserProvidedReadRecords: false,
         UserProvidedUpdateRecord: false,
         UserProvidedDeleteRecord: false,
-        UserProvidedExecuteEntity: false
+        UserProvidedExecuteEntity: false,
+        UserProvidedAggregateRecords: false,
+        UserProvidedAggregateRecordsQueryTimeout: false,
+        EffectiveAggregateRecordsQueryTimeoutSeconds: 30
       }
     },
     Host: {

src/Cli.Tests/Snapshots/InitTests.EnsureCorrectConfigGenerationWithDifferentAuthenticationProviders_171ea8114ff71814.verified.txt

@@ -27,13 +27,17 @@
         UpdateRecord: true,
         DeleteRecord: true,
         ExecuteEntity: true,
+        AggregateRecords: true,
         UserProvidedAllTools: false,
         UserProvidedDescribeEntities: false,
         UserProvidedCreateRecord: false,
         UserProvidedReadRecords: false,
         UserProvidedUpdateRecord: false,
         UserProvidedDeleteRecord: false,
-        UserProvidedExecuteEntity: false
+        UserProvidedExecuteEntity: false,
+        UserProvidedAggregateRecords: false,
+        UserProvidedAggregateRecordsQueryTimeout: false,
+        EffectiveAggregateRecordsQueryTimeoutSeconds: 30
       }
     },
     Host: {

src/Cli.Tests/Snapshots/InitTests.EnsureCorrectConfigGenerationWithDifferentAuthenticationProviders_2df7a1794712f154.verified.txt

@@ -27,13 +27,17 @@
         UpdateRecord: true,
         DeleteRecord: true,
         ExecuteEntity: true,
+        AggregateRecords: true,
         UserProvidedAllTools: false,
         UserProvidedDescribeEntities: false,
         UserProvidedCreateRecord: false,
         UserProvidedReadRecords: false,
         UserProvidedUpdateRecord: false,
         UserProvidedDeleteRecord: false,
-        UserProvidedExecuteEntity: false
+        UserProvidedExecuteEntity: false,
+        UserProvidedAggregateRecords: false,
+        UserProvidedAggregateRecordsQueryTimeout: false,
+        EffectiveAggregateRecordsQueryTimeoutSeconds: 30
       }
     },
     Host: {

src/Cli.Tests/Snapshots/InitTests.EnsureCorrectConfigGenerationWithDifferentAuthenticationProviders_59fe1a10aa78899d.verified.txt

@@ -27,13 +27,17 @@
         UpdateRecord: true,
         DeleteRecord: true,
         ExecuteEntity: true,
+        AggregateRecords: true,
         UserProvidedAllTools: false,
         UserProvidedDescribeEntities: false,
         UserProvidedCreateRecord: false,
         UserProvidedReadRecords: false,
         UserProvidedUpdateRecord: false,
         UserProvidedDeleteRecord: false,
-        UserProvidedExecuteEntity: false
+        UserProvidedExecuteEntity: false,
+        UserProvidedAggregateRecords: false,
+        UserProvidedAggregateRecordsQueryTimeout: false,
+        EffectiveAggregateRecordsQueryTimeoutSeconds: 30
       }
     },
     Host: {