Add aggregate_records MCP tool with query-timeout support

schemas/dab.draft.schema.json

@@ -275,6 +275,12 @@
               "description": "Allow enabling/disabling MCP requests for all entities.",
               "default": true
             },
+            "query-timeout": {
+              "type": "integer",
+              "description": "Execution timeout in seconds for MCP tool operations. Applies to all MCP tools.",
+              "default": 30,
+              "minimum": 1
+            },
             "dml-tools": {
               "oneOf": [
                 {
@@ -315,6 +321,11 @@
                       "type": "boolean",
                       "description": "Enable/disable the execute-entity tool.",
                       "default": false
+                    },
+                    "aggregate-records": {
+                      "type": "boolean",
+                      "description": "Enable/disable the aggregate-records tool.",
+                      "default": false
                     }
                   }
                 }

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 TIMEOUT = "Timeout";
     }
 }

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

@@ -60,8 +60,33 @@ public static async Task<CallToolResult> ExecuteWithTelemetryAsync(
                     operation: operation,
                     dbProcedure: dbProcedure);
 
-                // Execute the tool
-                CallToolResult result = await tool.ExecuteAsync(arguments, serviceProvider, cancellationToken);
+                // Read query-timeout from current config per invocation (hot-reload safe).
+                int timeoutSeconds = McpRuntimeOptions.DEFAULT_QUERY_TIMEOUT_SECONDS;
+                RuntimeConfigProvider? runtimeConfigProvider = serviceProvider.GetService<RuntimeConfigProvider>();
+                if (runtimeConfigProvider is not null)
+                {
+                    RuntimeConfig config = runtimeConfigProvider.GetConfig();
+                    timeoutSeconds = config.Runtime?.Mcp?.EffectiveQueryTimeoutSeconds ?? McpRuntimeOptions.DEFAULT_QUERY_TIMEOUT_SECONDS;
+                }
+
+                // Wrap tool execution with the configured timeout using a linked CancellationTokenSource.
+                using CancellationTokenSource timeoutCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+                timeoutCts.CancelAfter(TimeSpan.FromSeconds(timeoutSeconds));
+
+                CallToolResult result;
+                try
+                {
+                    result = await tool.ExecuteAsync(arguments, serviceProvider, timeoutCts.Token);
+                }
+                catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested)
+                {
+                    // The timeout CTS fired, not the caller's token. Surface as TimeoutException
+                    // so downstream telemetry and tool handlers see TIMEOUT, not cancellation.
+                    throw new TimeoutException(
+                        $"The MCP tool '{toolName}' did not complete within {timeoutSeconds} {(timeoutSeconds == 1 ? "second" : "seconds")}. "
+                        + "This is NOT a tool error. The operation exceeded the configured query-timeout. "
+                        + "Try narrowing results with a filter, reducing groupby fields, or using pagination.");
+                }
 
                 // Check if the tool returned an error result (tools catch exceptions internally
                 // and return CallToolResult with IsError=true instead of throwing)
@@ -124,6 +149,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 +214,7 @@ public static string MapExceptionToErrorCode(Exception ex)
             return ex switch
             {
                 OperationCanceledException => McpTelemetryErrorCodes.OPERATION_CANCELLED,
+                TimeoutException => McpTelemetryErrorCodes.TIMEOUT,
                 DataApiBuilderException dabEx when dabEx.SubStatusCode == DataApiBuilderException.SubStatusCodes.AuthenticationChallenge
                     => McpTelemetryErrorCodes.AUTHENTICATION_FAILED,
                 DataApiBuilderException dabEx when dabEx.SubStatusCode == DataApiBuilderException.SubStatusCodes.AuthorizationCheckFailed

src/Cli/Commands/ConfigureOptions.cs

@@ -42,6 +42,7 @@ public ConfigureOptions(
             bool? runtimeMcpEnabled = null,
             string? runtimeMcpPath = null,
             string? runtimeMcpDescription = null,
+            int? runtimeMcpQueryTimeout = null,
             bool? runtimeMcpDmlToolsEnabled = null,
             bool? runtimeMcpDmlToolsDescribeEntitiesEnabled = null,
             bool? runtimeMcpDmlToolsCreateRecordEnabled = null,
@@ -102,6 +103,7 @@ public ConfigureOptions(
             RuntimeMcpEnabled = runtimeMcpEnabled;
             RuntimeMcpPath = runtimeMcpPath;
             RuntimeMcpDescription = runtimeMcpDescription;
+            RuntimeMcpQueryTimeout = runtimeMcpQueryTimeout;
             RuntimeMcpDmlToolsEnabled = runtimeMcpDmlToolsEnabled;
             RuntimeMcpDmlToolsDescribeEntitiesEnabled = runtimeMcpDmlToolsDescribeEntitiesEnabled;
             RuntimeMcpDmlToolsCreateRecordEnabled = runtimeMcpDmlToolsCreateRecordEnabled;
@@ -203,6 +205,9 @@ public ConfigureOptions(
         [Option("runtime.mcp.description", Required = false, HelpText = "Set the MCP server description to be exposed in the initialize response.")]
         public string? RuntimeMcpDescription { get; }
 
+        [Option("runtime.mcp.query-timeout", Required = false, HelpText = "Set the execution timeout in seconds for MCP tool operations. Applies to all MCP tools. Default: 30 (integer). Must be >= 1.")]
+        public int? RuntimeMcpQueryTimeout { get; }
+
         [Option("runtime.mcp.dml-tools.enabled", Required = false, HelpText = "Enable DAB's MCP DML tools endpoint. Default: true (boolean).")]
         public bool? RuntimeMcpDmlToolsEnabled { get; }
 
@@ -224,6 +229,9 @@ public ConfigureOptions(
         [Option("runtime.mcp.dml-tools.execute-entity.enabled", Required = false, HelpText = "Enable DAB's MCP execute entity tool. Default: true (boolean).")]
         public bool? RuntimeMcpDmlToolsExecuteEntityEnabled { get; }
 
+        [Option("runtime.mcp.dml-tools.aggregate-records.enabled", Required = false, HelpText = "Enable DAB's MCP aggregate records tool. Default: true (boolean).")]
+        public bool? RuntimeMcpDmlToolsAggregateRecordsEnabled { get; }
+
         [Option("runtime.cache.enabled", Required = false, HelpText = "Enable DAB's cache globally. (You must also enable each entity's cache separately.). Default: false (boolean).")]
         public bool? RuntimeCacheEnabled { get; }
 

src/Cli/ConfigGenerator.cs

@@ -876,13 +876,15 @@ private static bool TryUpdateConfiguredRuntimeOptions(
             if (options.RuntimeMcpEnabled != null ||
                 options.RuntimeMcpPath != null ||
                 options.RuntimeMcpDescription != null ||
+                options.RuntimeMcpQueryTimeout != null ||
                 options.RuntimeMcpDmlToolsEnabled != null ||
                 options.RuntimeMcpDmlToolsDescribeEntitiesEnabled != null ||
                 options.RuntimeMcpDmlToolsCreateRecordEnabled != null ||
                 options.RuntimeMcpDmlToolsReadRecordsEnabled != null ||
                 options.RuntimeMcpDmlToolsUpdateRecordEnabled != null ||
                 options.RuntimeMcpDmlToolsDeleteRecordEnabled != null ||
-                options.RuntimeMcpDmlToolsExecuteEntityEnabled != null)
+                options.RuntimeMcpDmlToolsExecuteEntityEnabled != null ||
+                options.RuntimeMcpDmlToolsAggregateRecordsEnabled != null)
             {
                 McpRuntimeOptions updatedMcpOptions = runtimeConfig?.Runtime?.Mcp ?? new();
                 bool status = TryUpdateConfiguredMcpValues(options, ref updatedMcpOptions);
@@ -1161,6 +1163,14 @@ private static bool TryUpdateConfiguredMcpValues(
                     _logger.LogInformation("Updated RuntimeConfig with Runtime.Mcp.Description as '{updatedValue}'", updatedValue);
                 }
 
+                // Runtime.Mcp.QueryTimeout
+                updatedValue = options?.RuntimeMcpQueryTimeout;
+                if (updatedValue != null)
+                {
+                    updatedMcpOptions = updatedMcpOptions! with { QueryTimeout = (int)updatedValue };
+                    _logger.LogInformation("Updated RuntimeConfig with Runtime.Mcp.QueryTimeout as '{updatedValue}'", updatedValue);
+                }
+
                 // Handle DML tools configuration
                 bool hasToolUpdates = false;
                 DmlToolsConfig? currentDmlTools = updatedMcpOptions?.DmlTools;
@@ -1181,6 +1191,7 @@ private static bool TryUpdateConfiguredMcpValues(
                 bool? updateRecord = currentDmlTools?.UpdateRecord;
                 bool? deleteRecord = currentDmlTools?.DeleteRecord;
                 bool? executeEntity = currentDmlTools?.ExecuteEntity;
+                bool? aggregateRecords = currentDmlTools?.AggregateRecords;
 
                 updatedValue = options?.RuntimeMcpDmlToolsDescribeEntitiesEnabled;
                 if (updatedValue != null)
@@ -1230,6 +1241,14 @@ private static bool TryUpdateConfiguredMcpValues(
                     _logger.LogInformation("Updated RuntimeConfig with runtime.mcp.dml-tools.execute-entity as '{updatedValue}'", updatedValue);
                 }
 
+                updatedValue = options?.RuntimeMcpDmlToolsAggregateRecordsEnabled;
+                if (updatedValue != null)
+                {
+                    aggregateRecords = (bool)updatedValue;
+                    hasToolUpdates = true;
+                    _logger.LogInformation("Updated RuntimeConfig with runtime.mcp.dml-tools.aggregate-records as '{updatedValue}'", updatedValue);
+                }
+
                 if (hasToolUpdates)
                 {
                     updatedMcpOptions = updatedMcpOptions! with
@@ -1242,7 +1261,8 @@ private static bool TryUpdateConfiguredMcpValues(
                             ReadRecords = readRecord,
                             UpdateRecord = updateRecord,
                             DeleteRecord = deleteRecord,
-                            ExecuteEntity = executeEntity
+                            ExecuteEntity = executeEntity,
+                            AggregateRecords = aggregateRecords
                         }
                     };
                 }

src/Config/Converters/DmlToolsConfigConverter.cs

@@ -44,6 +44,7 @@ internal class DmlToolsConfigConverter : JsonConverter<DmlToolsConfig>
             bool? updateRecord = null;
             bool? deleteRecord = null;
             bool? executeEntity = null;
+            bool? aggregateRecords = null;
 
             while (reader.Read())
             {
@@ -82,6 +83,9 @@ internal class DmlToolsConfigConverter : JsonConverter<DmlToolsConfig>
                             case "execute-entity":
                                 executeEntity = value;
                                 break;
+                            case "aggregate-records":
+                                aggregateRecords = value;
+                                break;
                             default:
                                 // Skip unknown properties
                                 break;
@@ -91,7 +95,8 @@ internal class DmlToolsConfigConverter : JsonConverter<DmlToolsConfig>
                     {
                         // Error on non-boolean values for known properties
                         if (property?.ToLowerInvariant() is "describe-entities" or "create-record"
-                            or "read-records" or "update-record" or "delete-record" or "execute-entity")
+                            or "read-records" or "update-record" or "delete-record" or "execute-entity"
+                            or "aggregate-records")
                         {
                             throw new JsonException($"Property '{property}' must be a boolean value.");
                         }
@@ -110,7 +115,8 @@ internal class DmlToolsConfigConverter : JsonConverter<DmlToolsConfig>
                 readRecords: readRecords,
                 updateRecord: updateRecord,
                 deleteRecord: deleteRecord,
-                executeEntity: executeEntity);
+                executeEntity: executeEntity,
+                aggregateRecords: aggregateRecords);
         }
 
         // For any other unexpected token type, return default (all enabled)
@@ -135,7 +141,8 @@ public override void Write(Utf8JsonWriter writer, DmlToolsConfig? value, JsonSer
                                     value.UserProvidedReadRecords ||
                                     value.UserProvidedUpdateRecord ||
                                     value.UserProvidedDeleteRecord ||
-                                    value.UserProvidedExecuteEntity;
+                                    value.UserProvidedExecuteEntity ||
+                                    value.UserProvidedAggregateRecords;
 
         // Only write the boolean value if it's provided by user
         // This prevents writing "dml-tools": true when it's the default
@@ -181,6 +188,11 @@ public override void Write(Utf8JsonWriter writer, DmlToolsConfig? value, JsonSer
                 writer.WriteBoolean("execute-entity", value.ExecuteEntity.Value);
             }
 
+            if (value.UserProvidedAggregateRecords && value.AggregateRecords.HasValue)
+            {
+                writer.WriteBoolean("aggregate-records", value.AggregateRecords.Value);
+            }
+
             writer.WriteEndObject();
         }
     }

src/Config/Converters/McpRuntimeOptionsConverterFactory.cs

@@ -66,12 +66,13 @@ internal McpRuntimeOptionsConverter(DeserializationVariableReplacementSettings?
                 string? path = null;
                 DmlToolsConfig? dmlTools = null;
                 string? description = null;
+                int? queryTimeout = null;
 
                 while (reader.Read())
                 {
                     if (reader.TokenType == JsonTokenType.EndObject)
                     {
-                        return new McpRuntimeOptions(enabled, path, dmlTools, description);
+                        return new McpRuntimeOptions(enabled, path, dmlTools, description, queryTimeout);
                     }
 
                     string? propertyName = reader.GetString();
@@ -107,6 +108,14 @@ internal McpRuntimeOptionsConverter(DeserializationVariableReplacementSettings?
 
                             break;
 
+                        case "query-timeout":
+                            if (reader.TokenType is not JsonTokenType.Null)
+                            {
+                                queryTimeout = reader.GetInt32();
+                            }
+
+                            break;
+
                         default:
                             throw new JsonException($"Unexpected property {propertyName}");
                     }
@@ -150,6 +159,12 @@ public override void Write(Utf8JsonWriter writer, McpRuntimeOptions value, JsonS
                 JsonSerializer.Serialize(writer, value.Description, options);
             }
 
+            // Write query-timeout if it's user provided
+            if (value?.UserProvidedQueryTimeout is true && value.QueryTimeout.HasValue)
+            {
+                writer.WriteNumber("query-timeout", value.QueryTimeout.Value);
+            }
+
             writer.WriteEndObject();
         }
     }

src/Config/ObjectModel/DmlToolsConfig.cs

@@ -51,6 +51,11 @@ public record DmlToolsConfig
     /// </summary>
     public bool? ExecuteEntity { get; init; }
 
+    /// <summary>
+    /// Whether aggregate-records tool is enabled
+    /// </summary>
+    public bool? AggregateRecords { get; init; }
+
     [JsonConstructor]
     public DmlToolsConfig(
         bool? allToolsEnabled = null,
@@ -59,7 +64,8 @@ public DmlToolsConfig(
         bool? readRecords = null,
         bool? updateRecord = null,
         bool? deleteRecord = null,
-        bool? executeEntity = null)
+        bool? executeEntity = null,
+        bool? aggregateRecords = null)
     {
         if (allToolsEnabled is not null)
         {
@@ -75,6 +81,7 @@ public DmlToolsConfig(
             UpdateRecord = updateRecord ?? toolDefault;
             DeleteRecord = deleteRecord ?? toolDefault;
             ExecuteEntity = executeEntity ?? toolDefault;
+            AggregateRecords = aggregateRecords ?? toolDefault;
         }
         else
         {
@@ -87,6 +94,7 @@ public DmlToolsConfig(
             UpdateRecord = updateRecord ?? DEFAULT_ENABLED;
             DeleteRecord = deleteRecord ?? DEFAULT_ENABLED;
             ExecuteEntity = executeEntity ?? DEFAULT_ENABLED;
+            AggregateRecords = aggregateRecords ?? DEFAULT_ENABLED;
         }
 
         // Track user-provided status - only true if the parameter was not null
@@ -96,6 +104,7 @@ public DmlToolsConfig(
         UserProvidedUpdateRecord = updateRecord is not null;
         UserProvidedDeleteRecord = deleteRecord is not null;
         UserProvidedExecuteEntity = executeEntity is not null;
+        UserProvidedAggregateRecords = aggregateRecords is not null;
     }
 
     /// <summary>
@@ -112,7 +121,8 @@ public static DmlToolsConfig FromBoolean(bool enabled)
             readRecords: null,
             updateRecord: null,
             deleteRecord: null,
-            executeEntity: null
+            executeEntity: null,
+            aggregateRecords: null
         );
     }
 
@@ -127,7 +137,8 @@ public static DmlToolsConfig FromBoolean(bool enabled)
         readRecords: null,
         updateRecord: null,
         deleteRecord: null,
-        executeEntity: null
+        executeEntity: null,
+        aggregateRecords: null
     );
 
     /// <summary>
@@ -185,4 +196,12 @@ public static DmlToolsConfig FromBoolean(bool enabled)
     [JsonIgnore(Condition = JsonIgnoreCondition.Always)]
     [MemberNotNullWhen(true, nameof(ExecuteEntity))]
     public bool UserProvidedExecuteEntity { get; init; } = false;
+
+    /// <summary>
+    /// Flag which informs CLI and JSON serializer whether to write aggregate-records
+    /// property/value to the runtime config file.
+    /// </summary>
+    [JsonIgnore(Condition = JsonIgnoreCondition.Always)]
+    [MemberNotNullWhen(true, nameof(AggregateRecords))]
+    public bool UserProvidedAggregateRecords { get; init; } = false;
 }