diff --git a/src/Azure.DataApiBuilder.Mcp/BuiltInTools/ExecuteEntityTool.cs b/src/Azure.DataApiBuilder.Mcp/BuiltInTools/ExecuteEntityTool.cs index 90b51d2837..5291108fc8 100644 --- a/src/Azure.DataApiBuilder.Mcp/BuiltInTools/ExecuteEntityTool.cs +++ b/src/Azure.DataApiBuilder.Mcp/BuiltInTools/ExecuteEntityTool.cs @@ -381,22 +381,7 @@ private static CallToolResult BuildExecuteSuccessResponse( // Handle different result types if (queryResult is OkObjectResult okResult && okResult.Value != null) { - // Extract the actual data from the action result - if (okResult.Value is JsonDocument jsonDoc) - { - JsonElement root = jsonDoc.RootElement; - responseData["value"] = root.ValueKind == JsonValueKind.Array ? root : JsonSerializer.SerializeToElement(new[] { root }); - } - else if (okResult.Value is JsonElement jsonElement) - { - responseData["value"] = jsonElement.ValueKind == JsonValueKind.Array ? jsonElement : JsonSerializer.SerializeToElement(new[] { jsonElement }); - } - else - { - // Serialize the value directly - JsonElement serialized = JsonSerializer.SerializeToElement(okResult.Value); - responseData["value"] = serialized; - } + responseData["value"] = McpResponseBuilder.NormalizeStoredProcedureResult(okResult.Value); } else if (queryResult is BadRequestObjectResult badRequest) { diff --git a/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs b/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs index 1d78e81680..868a47995d 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs @@ -580,22 +580,7 @@ private static CallToolResult BuildExecuteSuccessResponse( // Handle different result types if (queryResult is OkObjectResult okResult && okResult.Value != null) { - // Extract the actual data from the action result - if (okResult.Value is JsonDocument jsonDoc) - { - JsonElement root = jsonDoc.RootElement; - responseData["value"] = root.ValueKind == JsonValueKind.Array ? root : JsonSerializer.SerializeToElement(new[] { root }); - } - else if (okResult.Value is JsonElement jsonElement) - { - responseData["value"] = jsonElement.ValueKind == JsonValueKind.Array ? jsonElement : JsonSerializer.SerializeToElement(new[] { jsonElement }); - } - else - { - // Serialize the value directly - JsonElement serialized = JsonSerializer.SerializeToElement(okResult.Value); - responseData["value"] = serialized; - } + responseData["value"] = McpResponseBuilder.NormalizeStoredProcedureResult(okResult.Value); } else if (queryResult is BadRequestObjectResult badRequest) { diff --git a/src/Azure.DataApiBuilder.Mcp/Utils/McpResponseBuilder.cs b/src/Azure.DataApiBuilder.Mcp/Utils/McpResponseBuilder.cs index 401f270f42..34de1956b2 100644 --- a/src/Azure.DataApiBuilder.Mcp/Utils/McpResponseBuilder.cs +++ b/src/Azure.DataApiBuilder.Mcp/Utils/McpResponseBuilder.cs @@ -102,6 +102,30 @@ public static string ExtractResultJson(IActionResult? result) } } + /// + /// Normalizes a stored procedure result to an array, unwrapping the DAB response envelope when present. + /// + public static JsonElement NormalizeStoredProcedureResult(object value) + { + JsonElement result = value switch + { + JsonDocument document => document.RootElement, + JsonElement element => element, + _ => JsonSerializer.SerializeToElement(value) + }; + + if (result.ValueKind == JsonValueKind.Object && + result.TryGetProperty("value", out JsonElement valueArray) && + valueArray.ValueKind == JsonValueKind.Array) + { + return valueArray.Clone(); + } + + return result.ValueKind == JsonValueKind.Array + ? result.Clone() + : JsonSerializer.SerializeToElement(new[] { result }); + } + /// /// Extracts value from a JsonElement. /// diff --git a/src/Service.Tests/Mcp/DynamicCustomToolMsSqlIntegrationTests.cs b/src/Service.Tests/Mcp/DynamicCustomToolMsSqlIntegrationTests.cs index 6ea4bb7f45..edc387636e 100644 --- a/src/Service.Tests/Mcp/DynamicCustomToolMsSqlIntegrationTests.cs +++ b/src/Service.Tests/Mcp/DynamicCustomToolMsSqlIntegrationTests.cs @@ -82,13 +82,9 @@ public async Task DynamicCustomTool_GetBookById_ReturnsMatchingRecord() Assert.IsTrue(root.TryGetProperty("value", out JsonElement valueWrapper), "Response should contain 'value' property."); - JsonElement records = valueWrapper.ValueKind == JsonValueKind.Object - ? valueWrapper.GetProperty("value") - : valueWrapper; - - Assert.AreEqual(JsonValueKind.Array, records.ValueKind); - Assert.IsTrue(records.GetArrayLength() > 0, "Expected at least one book record."); - Assert.AreEqual(1, records[0].GetProperty("id").GetInt32()); + Assert.AreEqual(JsonValueKind.Array, valueWrapper.ValueKind); + Assert.IsTrue(valueWrapper.GetArrayLength() > 0, "Expected at least one book record."); + Assert.AreEqual(1, valueWrapper[0].GetProperty("id").GetInt32()); } /// diff --git a/src/Service.Tests/Mcp/DynamicCustomToolTests.cs b/src/Service.Tests/Mcp/DynamicCustomToolTests.cs index aad83ed358..10e4e95df5 100644 --- a/src/Service.Tests/Mcp/DynamicCustomToolTests.cs +++ b/src/Service.Tests/Mcp/DynamicCustomToolTests.cs @@ -554,6 +554,25 @@ public async Task ExecuteAsync_JsonDocumentObjectResult_WrapsValueInArray() Assert.AreEqual(JsonValueKind.Array, response.RootElement.GetProperty("value").ValueKind); } + [DataTestMethod] + [DataRow("[]", DisplayName = "Empty result")] + [DataRow("[{\"id\":1},{\"id\":2}]", DisplayName = "Multiple rows")] + [DataRow("[{\"id\":1,\"value\":[\"first\",{\"nested\":\"second\"}]}]", DisplayName = "Array-valued value column")] + public async Task ExecuteAsync_EngineResponse_RemovesOnlyOuterEnvelope(string rowsJson) + { + using JsonDocument rows = JsonDocument.Parse(rowsJson); + CallToolResult result = await ExecuteCustomToolAsync( + dbParameters: new(), + userParameters: null, + queryOutcome: SqlResponseHelpers.OkResponse(rows.RootElement.Clone())); + + AssertSuccess(result, "The DAB response envelope should be removed from the stored procedure result."); + using JsonDocument response = JsonDocument.Parse(GetFirstText(result)); + JsonElement value = response.RootElement.GetProperty("value"); + Assert.AreEqual(JsonValueKind.Array, value.ValueKind); + Assert.AreEqual(JsonSerializer.Serialize(rows.RootElement), JsonSerializer.Serialize(value), "Rows should be unchanged after removing only the engine envelope."); + } + [TestMethod] public async Task ExecuteAsync_PocoResult_IsSerialized() { diff --git a/src/Service.Tests/Mcp/ExecuteEntityToolMsSqlIntegrationTests.cs b/src/Service.Tests/Mcp/ExecuteEntityToolMsSqlIntegrationTests.cs index f4089d8e19..75ce3d5730 100644 --- a/src/Service.Tests/Mcp/ExecuteEntityToolMsSqlIntegrationTests.cs +++ b/src/Service.Tests/Mcp/ExecuteEntityToolMsSqlIntegrationTests.cs @@ -76,13 +76,9 @@ public async Task ExecuteEntity_GetBookById_ReturnsMatchingRecord() Assert.IsTrue(root.TryGetProperty("value", out JsonElement valueWrapper), "Response should contain 'value' property."); - JsonElement records = valueWrapper.ValueKind == JsonValueKind.Object - ? valueWrapper.GetProperty("value") - : valueWrapper; - - Assert.AreEqual(JsonValueKind.Array, records.ValueKind); - Assert.IsTrue(records.GetArrayLength() > 0, "Expected at least one book record."); - Assert.AreEqual(1, records[0].GetProperty("id").GetInt32()); + Assert.AreEqual(JsonValueKind.Array, valueWrapper.ValueKind); + Assert.IsTrue(valueWrapper.GetArrayLength() > 0, "Expected at least one book record."); + Assert.AreEqual(1, valueWrapper[0].GetProperty("id").GetInt32()); } /// diff --git a/src/Service.Tests/Mcp/ExecuteEntityToolTests.cs b/src/Service.Tests/Mcp/ExecuteEntityToolTests.cs index 843cbb4cd3..e27b667b75 100644 --- a/src/Service.Tests/Mcp/ExecuteEntityToolTests.cs +++ b/src/Service.Tests/Mcp/ExecuteEntityToolTests.cs @@ -538,6 +538,26 @@ public async Task ExecuteEntity_ObjectJsonResult_IsWrappedInArray() Assert.AreEqual(1, value[0].GetProperty("id").GetInt32()); } + [DataTestMethod] + [DataRow("[]", DisplayName = "Empty result")] + [DataRow("[{\"id\":1},{\"id\":2}]", DisplayName = "Multiple rows")] + [DataRow("[{\"id\":1,\"value\":[\"first\",{\"nested\":\"second\"}]}]", DisplayName = "Array-valued value column")] + public async Task ExecuteEntity_EngineResponse_RemovesOnlyOuterEnvelope(string rowsJson) + { + using JsonDocument rows = JsonDocument.Parse(rowsJson); + CallToolResult result = await ExecuteWithMockedEngineAsync( + TEST_ENTITY, + new(), + null, + queryOutcome: SqlResponseHelpers.OkResponse(rows.RootElement.Clone())); + + AssertSuccess(result, "The DAB response envelope should be removed from the stored procedure result."); + using JsonDocument response = JsonDocument.Parse(GetFirstText(result)); + JsonElement value = response.RootElement.GetProperty("value"); + Assert.AreEqual(JsonValueKind.Array, value.ValueKind); + Assert.AreEqual(JsonSerializer.Serialize(rows.RootElement), JsonSerializer.Serialize(value), "Rows should be unchanged after removing only the engine envelope."); + } + [TestMethod] public async Task ExecuteEntity_UnknownResult_ReturnsEmptyValue() {