Skip to content

Commit 883f4a1

Browse files
souvikghosh04CopilotJerryNixonanushakolanJerry Nixon
authored
[MCP] aggregate_records DML tool (#3199)
## Why make this change? - Closes #3178 - Adds `aggregate_records` as a new DML tool to the MCP server, enabling models to answer common aggregation questions like "How many products are there?" and "What is our most expensive product?" - Continuation of work from #3179 (stale PR with resolved review comments). ## What is this change? - New `aggregate_records` DML tool (`AggregateRecordsTool.cs`) that generates SQL-level aggregation queries (`COUNT`, `AVG`, `SUM`, `MIN`, `MAX`) with support for `DISTINCT`, OData `$filter` (WHERE), `GROUP BY`, `HAVING` operators (`eq`, `neq`, `gt`, `gte`, `lt`, `lte`, `in`), `ORDER BY` (asc/desc), and cursor-based pagination (`first`/`after`) — all per the spec in #3178. query-timeout configuration for MCP runtime options — allows setting a per-query timeout (1–600 seconds, default 30s) via `McpRuntimeOptions.QueryTimeout`, validated at startup. - Config & CLI plumbing: `DmlToolsConfig` gains `AggregateRecords`/`UserProvidedAggregateRecords`; `McpRuntimeOptionsConverterFactory` handles query-timeout serialization; `ConfigureOptions.cs` updated for CLI configure support; JSON schema updated. - Telemetry: Aggregation operations emit OpenTelemetry traces with structured error codes. - Updated 37 CLI + 4 Service.Tests snapshot files to reflect new default properties. ## How was this tested? - [x] Integration Tests — `AggregateRecordsToolTests.cs` (all 13 spec examples + edge cases), `McpQueryTimeoutTests.cs`, `EntityLevelDmlToolConfigurationTests.cs`, `McpToolRegistryTests.cs` - [x] Unit Tests — `AggregateRecordsToolTests.cs` (unit), `McpTelemetryTests.cs`, `RequestParserUnitTests.cs`, `SqlQueryExecutorUnitTests.cs` --------- Co-authored-by: copilot-swe-agent[bot] <[email protected]> Co-authored-by: JerryNixon <[email protected]> Co-authored-by: JerryNixon <[email protected]> Co-authored-by: anushakolan <[email protected]> Co-authored-by: Jerry Nixon <[email protected]> Co-authored-by: Copilot <[email protected]> Co-authored-by: souvikghosh04 <[email protected]> Co-authored-by: Aniruddh Munde <[email protected]>
1 parent c4c5943 commit 883f4a1

58 files changed

Lines changed: 2618 additions & 309 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎schemas/dab.draft.schema.json‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -315,6 +315,34 @@
315315
"type": "boolean",
316316
"description": "Enable/disable the execute-entity tool.",
317317
"default": false
318+
},
319+
"aggregate-records": {
320+
"oneOf": [
321+
{
322+
"type": "boolean",
323+
"description": "Enable/disable the aggregate-records tool."
324+
},
325+
{
326+
"type": "object",
327+
"description": "Aggregate records tool configuration",
328+
"additionalProperties": false,
329+
"properties": {
330+
"enabled": {
331+
"type": "boolean",
332+
"description": "Enable/disable the aggregate-records tool.",
333+
"default": true
334+
},
335+
"query-timeout": {
336+
"type": "integer",
337+
"description": "Execution timeout in seconds for aggregate queries. Range: 1-600.",
338+
"default": 30,
339+
"minimum": 1,
340+
"maximum": 600
341+
}
342+
}
343+
}
344+
],
345+
"default": false
318346
}
319347
}
320348
}

‎src/Azure.DataApiBuilder.Mcp/BuiltInTools/AggregateRecordsTool.cs‎

Lines changed: 1087 additions & 0 deletions
Large diffs are not rendered by default.

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,5 +24,16 @@ public static CallToolResult ToolDisabled(string toolName, ILogger? logger, stri
2424
string message = customMessage ?? $"The {toolName} tool is disabled in the configuration.";
2525
return McpResponseBuilder.BuildErrorResult(toolName, Model.McpErrorCode.ToolDisabled.ToString(), message, logger);
2626
}
27+
28+
/// <summary>
29+
/// Returns a model-friendly error when a field name is not found for an entity.
30+
/// Guides the model to call describe_entities to discover valid field names.
31+
/// </summary>
32+
public static CallToolResult FieldNotFound(string toolName, string entityName, string fieldName, string parameterName, ILogger? logger)
33+
{
34+
string message = $"Field '{fieldName}' in '{parameterName}' was not found for entity '{entityName}'. "
35+
+ $"Call describe_entities to get valid field names for '{entityName}'.";
36+
return McpResponseBuilder.BuildErrorResult(toolName, "FieldNotFound", message, logger);
37+
}
2738
}
2839
}

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,5 +37,10 @@ internal static class McpTelemetryErrorCodes
3737
/// Operation cancelled error code.
3838
/// </summary>
3939
public const string OPERATION_CANCELLED = "OperationCancelled";
40+
41+
/// <summary>
42+
/// Operation timed out error code.
43+
/// </summary>
44+
public const string OPERATION_TIMEOUT = "OperationTimeout";
4045
}
4146
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,6 @@ public static async Task<CallToolResult> ExecuteWithTelemetryAsync(
6060
operation: operation,
6161
dbProcedure: dbProcedure);
6262

63-
// Execute the tool
6463
CallToolResult result = await tool.ExecuteAsync(arguments, serviceProvider, cancellationToken);
6564

6665
// Check if the tool returned an error result (tools catch exceptions internally
@@ -124,6 +123,7 @@ public static string InferOperationFromTool(IMcpTool tool, string toolName)
124123
"delete_record" => "delete",
125124
"describe_entities" => "describe",
126125
"execute_entity" => "execute",
126+
"aggregate_records" => "aggregate",
127127
_ => "execute" // Fallback for any unknown built-in tools
128128
};
129129
}
@@ -188,6 +188,7 @@ public static string MapExceptionToErrorCode(Exception ex)
188188
return ex switch
189189
{
190190
OperationCanceledException => McpTelemetryErrorCodes.OPERATION_CANCELLED,
191+
TimeoutException => McpTelemetryErrorCodes.OPERATION_TIMEOUT,
191192
DataApiBuilderException dabEx when dabEx.SubStatusCode == DataApiBuilderException.SubStatusCodes.AuthenticationChallenge
192193
=> McpTelemetryErrorCodes.AUTHENTICATION_FAILED,
193194
DataApiBuilderException dabEx when dabEx.SubStatusCode == DataApiBuilderException.SubStatusCodes.AuthorizationCheckFailed

‎src/Cli.Tests/InitTests.cs‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -493,6 +493,74 @@ public Task VerifyCorrectConfigGenerationWithMultipleMutationOptions(DatabaseTyp
493493
return ExecuteVerifyTest(options, verifySettings);
494494
}
495495

496+
/// <summary>
497+
/// Test that init with/without --mcp.aggregate-records.query-timeout produces a config
498+
/// with the correct aggregate-records query-timeout in the DmlTools section.
499+
/// When null (not specified), defaults to 30 seconds. When provided, the config reflects the value.
500+
/// </summary>
501+
[DataTestMethod]
502+
[DataRow(null, false, DmlToolsConfig.DEFAULT_QUERY_TIMEOUT_SECONDS, DisplayName = "Init without query-timeout uses default 30s")]
503+
[DataRow(1, true, 1, DisplayName = "Init with query-timeout 1s (minimum)")]
504+
[DataRow(120, true, 120, DisplayName = "Init with query-timeout 120s")]
505+
[DataRow(600, true, 600, DisplayName = "Init with query-timeout 600s (maximum)")]
506+
public void InitWithAggregateRecordsQueryTimeout_SetsOrDefaultsTimeout(int? inputTimeout, bool expectedUserProvided, int expectedEffectiveTimeout)
507+
{
508+
InitOptions options = new(
509+
databaseType: DatabaseType.MSSQL,
510+
connectionString: "testconnectionstring",
511+
cosmosNoSqlDatabase: null,
512+
cosmosNoSqlContainer: null,
513+
graphQLSchemaPath: null,
514+
setSessionContext: false,
515+
hostMode: HostMode.Development,
516+
corsOrigin: null,
517+
authenticationProvider: EasyAuthType.AppService.ToString(),
518+
mcpAggregateRecordsQueryTimeout: inputTimeout,
519+
config: TEST_RUNTIME_CONFIG_FILE);
520+
521+
Assert.IsTrue(TryCreateRuntimeConfig(options, _runtimeConfigLoader!, _fileSystem!, out RuntimeConfig? runtimeConfig));
522+
Assert.IsNotNull(runtimeConfig?.Runtime?.Mcp?.DmlTools);
523+
Assert.AreEqual(inputTimeout, runtimeConfig.Runtime.Mcp.DmlTools.AggregateRecordsQueryTimeout);
524+
Assert.AreEqual(expectedUserProvided, runtimeConfig.Runtime.Mcp.DmlTools.UserProvidedAggregateRecordsQueryTimeout);
525+
Assert.AreEqual(expectedEffectiveTimeout, runtimeConfig.Runtime.Mcp.DmlTools.EffectiveAggregateRecordsQueryTimeoutSeconds);
526+
}
527+
528+
/// <summary>
529+
/// Test that init with --mcp.aggregate-records.query-timeout produces valid JSON
530+
/// that round-trips correctly through serialization/deserialization.
531+
/// </summary>
532+
[TestMethod]
533+
public void InitWithAggregateRecordsQueryTimeout_RoundTripsCorrectly()
534+
{
535+
InitOptions options = new(
536+
databaseType: DatabaseType.MSSQL,
537+
connectionString: "testconnectionstring",
538+
cosmosNoSqlDatabase: null,
539+
cosmosNoSqlContainer: null,
540+
graphQLSchemaPath: null,
541+
setSessionContext: false,
542+
hostMode: HostMode.Development,
543+
corsOrigin: null,
544+
authenticationProvider: EasyAuthType.AppService.ToString(),
545+
mcpAggregateRecordsQueryTimeout: 90,
546+
config: TEST_RUNTIME_CONFIG_FILE);
547+
548+
Assert.IsTrue(TryCreateRuntimeConfig(options, _runtimeConfigLoader!, _fileSystem!, out RuntimeConfig? runtimeConfig));
549+
550+
// Serialize to JSON and deserialize back
551+
JsonSerializerOptions serializerOptions = RuntimeConfigLoader.GetSerializationOptions();
552+
string json = JsonSerializer.Serialize(runtimeConfig, serializerOptions);
553+
RuntimeConfig? deserialized = JsonSerializer.Deserialize<RuntimeConfig>(json, serializerOptions);
554+
555+
Assert.IsNotNull(deserialized?.Runtime?.Mcp?.DmlTools);
556+
Assert.AreEqual(90, deserialized.Runtime.Mcp.DmlTools.AggregateRecordsQueryTimeout);
557+
Assert.AreEqual(90, deserialized.Runtime.Mcp.DmlTools.EffectiveAggregateRecordsQueryTimeoutSeconds);
558+
559+
// Verify the JSON contains the object format for aggregate-records
560+
Assert.IsTrue(json.Contains("\"query-timeout\""), $"Expected 'query-timeout' in serialized JSON. Got: {json}");
561+
Assert.IsTrue(json.Contains("90"), $"Expected timeout value 90 in serialized JSON. Got: {json}");
562+
}
563+
496564
private Task ExecuteVerifyTest(InitOptions options, VerifySettings? settings = null)
497565
{
498566
Assert.IsTrue(TryCreateRuntimeConfig(options, _runtimeConfigLoader!, _fileSystem!, out RuntimeConfig? runtimeConfig));

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,13 +27,17 @@
2727
UpdateRecord: true,
2828
DeleteRecord: true,
2929
ExecuteEntity: true,
30+
AggregateRecords: true,
3031
UserProvidedAllTools: false,
3132
UserProvidedDescribeEntities: false,
3233
UserProvidedCreateRecord: false,
3334
UserProvidedReadRecords: false,
3435
UserProvidedUpdateRecord: false,
3536
UserProvidedDeleteRecord: false,
36-
UserProvidedExecuteEntity: false
37+
UserProvidedExecuteEntity: false,
38+
UserProvidedAggregateRecords: false,
39+
UserProvidedAggregateRecordsQueryTimeout: false,
40+
EffectiveAggregateRecordsQueryTimeoutSeconds: 30
3741
}
3842
},
3943
Host: {

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,13 +27,17 @@
2727
UpdateRecord: true,
2828
DeleteRecord: true,
2929
ExecuteEntity: true,
30+
AggregateRecords: true,
3031
UserProvidedAllTools: false,
3132
UserProvidedDescribeEntities: false,
3233
UserProvidedCreateRecord: false,
3334
UserProvidedReadRecords: false,
3435
UserProvidedUpdateRecord: false,
3536
UserProvidedDeleteRecord: false,
36-
UserProvidedExecuteEntity: false
37+
UserProvidedExecuteEntity: false,
38+
UserProvidedAggregateRecords: false,
39+
UserProvidedAggregateRecordsQueryTimeout: false,
40+
EffectiveAggregateRecordsQueryTimeoutSeconds: 30
3741
}
3842
},
3943
Host: {

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,13 +27,17 @@
2727
UpdateRecord: true,
2828
DeleteRecord: true,
2929
ExecuteEntity: true,
30+
AggregateRecords: true,
3031
UserProvidedAllTools: false,
3132
UserProvidedDescribeEntities: false,
3233
UserProvidedCreateRecord: false,
3334
UserProvidedReadRecords: false,
3435
UserProvidedUpdateRecord: false,
3536
UserProvidedDeleteRecord: false,
36-
UserProvidedExecuteEntity: false
37+
UserProvidedExecuteEntity: false,
38+
UserProvidedAggregateRecords: false,
39+
UserProvidedAggregateRecordsQueryTimeout: false,
40+
EffectiveAggregateRecordsQueryTimeoutSeconds: 30
3741
}
3842
},
3943
Host: {

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,13 +27,17 @@
2727
UpdateRecord: true,
2828
DeleteRecord: true,
2929
ExecuteEntity: true,
30+
AggregateRecords: true,
3031
UserProvidedAllTools: false,
3132
UserProvidedDescribeEntities: false,
3233
UserProvidedCreateRecord: false,
3334
UserProvidedReadRecords: false,
3435
UserProvidedUpdateRecord: false,
3536
UserProvidedDeleteRecord: false,
36-
UserProvidedExecuteEntity: false
37+
UserProvidedExecuteEntity: false,
38+
UserProvidedAggregateRecords: false,
39+
UserProvidedAggregateRecordsQueryTimeout: false,
40+
EffectiveAggregateRecordsQueryTimeoutSeconds: 30
3741
}
3842
},
3943
Host: {

0 commit comments

Comments
 (0)