Add aggregate_records DML tool and query-timeout to MCP server

schemas/dab.draft.schema.json

@@ -275,6 +275,13 @@
               "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. Range: 1-600.",
+              "default": 30,
+              "minimum": 1,
+              "maximum": 600
+            },
             "dml-tools": {
               "oneOf": [
                 {
@@ -315,6 +322,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/BuiltInTools/AggregateRecordsTool.md

@@ -0,0 +1,51 @@
+# AggregateRecordsTool
+
+MCP tool that computes SQL-level aggregations (COUNT, AVG, SUM, MIN, MAX) on DAB entities. All aggregation is pushed to the database engine — no in-memory computation.
+
+## Class Structure
+
+| Member | Kind | Purpose |
+|---|---|---|
+| `ToolType` | Property | Returns `ToolType.BuiltIn` for reflection-based discovery. |
+| `_validFunctions` | Static field | Allowlist of aggregation functions: count, avg, sum, min, max. |
+| `GetToolMetadata()` | Method | Returns the MCP `Tool` descriptor (name, description, JSON input schema). |
+| `ExecuteAsync()` | Method | Main entry point — validates input, resolves metadata, authorizes, builds the SQL query via the engine's `IQueryBuilder.Build(SqlQueryStructure)`, executes it, and formats the response. |
+| `ComputeAlias()` | Static method | Produces the result column alias: `"count"` for count(\*), otherwise `"{function}_{field}"`. |
+| `DecodeCursorOffset()` | Static method | Decodes a base64 opaque cursor string to an integer offset for OFFSET/FETCH pagination. Returns 0 on any invalid input. |
+| `BuildPaginatedResponse()` | Private method | Formats a grouped result set into `{ items, endCursor, hasNextPage }` when `first` is provided. |
+| `BuildSimpleResponse()` | Private method | Formats a scalar or grouped result set without pagination. |
+
+## ExecuteAsync Sequence
+
+```mermaid
+sequenceDiagram
+    participant Client as MCP Client
+    participant Tool as AggregateRecordsTool
+    participant Engine as DAB Engine
+    participant DB as Database
+
+    Client->>Tool: ExecuteAsync(arguments)
+    Tool->>Tool: Validate inputs & check tool enabled
+    Tool->>Engine: Resolve entity metadata & validate fields
+    Tool->>Engine: Authorize (column-level permissions)
+    Tool->>Engine: Build SQL via queryBuilder.Build(SqlQueryStructure)
+    Tool->>Tool: Post-process SQL (ORDER BY, pagination)
+    Tool->>DB: ExecuteQueryAsync → JSON result
+    alt Paginated (first provided)
+        Tool-->>Client: { items, endCursor, hasNextPage }
+    else Simple
+        Tool-->>Client: { entity, result: [{alias: value}] }
+    end
+
+    Note over Tool,Client: On error: TimeoutError, OperationCanceled, or DatabaseOperationFailed
+```
+
+## Key Design Decisions
+
+- **No in-memory aggregation.** The engine's `GroupByMetadata` / `AggregationColumn` types drive SQL generation via `queryBuilder.Build(structure)`. All aggregation is performed by the database.
+- **COUNT(\*) workaround.** The engine's `Build(AggregationColumn)` doesn't support `*` as a column name (it produces invalid SQL like `count([].[*])`), so the primary key column is used instead. `COUNT(pk)` ≡ `COUNT(*)` since PK is NOT NULL.
+- **ORDER BY post-processing.** Neither the GraphQL nor REST code paths support ORDER BY on an aggregate expression, so this tool inserts `ORDER BY {func}({col}) ASC|DESC` into the generated SQL before `FOR JSON PATH`.
+- **TOP vs OFFSET/FETCH.** SQL Server forbids both in the same query. When pagination (`first`) is used, `TOP N` is stripped via regex before appending `OFFSET/FETCH NEXT`.
+- **Early field validation.** All user-supplied field names (aggregation field, groupby fields) are validated against the entity's metadata before authorization or query building, so typos surface immediately with actionable guidance.
+- **Timeout vs cancellation.** `TimeoutException` (from `query-timeout` config) and `OperationCanceledException` (from client disconnect) are handled separately with distinct model-facing messages. Timeouts guide the model to narrow filters or paginate; cancellations suggest retry.
+- **Database support.** Only MsSql / DWSQL — matches the engine's GraphQL aggregation support. PostgreSQL, MySQL, and CosmosDB return an `UnsupportedDatabase` error.

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

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

@@ -60,8 +60,39 @@ 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;
+                }
+
+                // Defensive runtime guard: clamp timeout to valid range [1, MAX_QUERY_TIMEOUT_SECONDS].
+                if (timeoutSeconds < 1 || timeoutSeconds > McpRuntimeOptions.MAX_QUERY_TIMEOUT_SECONDS)
+                {
+                    timeoutSeconds = 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 +155,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 +220,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.Tests/Snapshots/EndToEndTests.TestAddingStoredProcedureWithRestMethodsAndGraphQLOperations.verified.txt

@@ -1,4 +1,4 @@
-{
+{
   DataSource: {
     DatabaseType: MSSQL,
     Options: {
@@ -27,14 +27,18 @@
         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
+      },
+      UserProvidedQueryTimeout: false,
+      EffectiveQueryTimeoutSeconds: 30
     },
     Host: {
       Cors: {

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

@@ -1,4 +1,4 @@
-{
+{
   DataSource: {
     DatabaseType: MSSQL,
     Options: {
@@ -27,14 +27,18 @@
         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
+      },
+      UserProvidedQueryTimeout: false,
+      EffectiveQueryTimeoutSeconds: 30
     },
     Host: {
       Cors: {

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

@@ -1,4 +1,4 @@
-{
+{
   DataSource: {
     DatabaseType: MSSQL,
     Options: {
@@ -27,14 +27,18 @@
         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
+      },
+      UserProvidedQueryTimeout: false,
+      EffectiveQueryTimeoutSeconds: 30
     },
     Host: {
       Cors: {

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

@@ -1,4 +1,4 @@
-{
+{
   DataSource: {
     DatabaseType: MSSQL,
     Options: {
@@ -27,14 +27,18 @@
         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
+      },
+      UserProvidedQueryTimeout: false,
+      EffectiveQueryTimeoutSeconds: 30
     },
     Host: {
       Cors: {

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

@@ -1,4 +1,4 @@
-{
+{
   DataSource: {
     Options: {
       container: planet,
@@ -28,14 +28,18 @@
         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
+      },
+      UserProvidedQueryTimeout: false,
+      EffectiveQueryTimeoutSeconds: 30
     },
     Host: {
       Cors: {

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

@@ -27,14 +27,18 @@
         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
+      },
+      UserProvidedQueryTimeout: false,
+      EffectiveQueryTimeoutSeconds: 30
     },
     Host: {
       Cors: {

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

@@ -1,4 +1,4 @@
-{
+{
   DataSource: {
     DatabaseType: MSSQL,
     Options: {
@@ -27,14 +27,18 @@
         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
+      },
+      UserProvidedQueryTimeout: false,
+      EffectiveQueryTimeoutSeconds: 30
     },
     Host: {
       Cors: {

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

@@ -1,4 +1,4 @@
-{
+{
   DataSource: {
     Options: {
       container: testcontainer,
@@ -28,14 +28,18 @@
         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
+      },
+      UserProvidedQueryTimeout: false,
+      EffectiveQueryTimeoutSeconds: 30
     },
     Host: {
       Cors: {