diff --git a/src/Azure.DataApiBuilder.Mcp/Azure.DataApiBuilder.Mcp.csproj b/src/Azure.DataApiBuilder.Mcp/Azure.DataApiBuilder.Mcp.csproj index f675f8d8d1..c6a8b7bf21 100644 --- a/src/Azure.DataApiBuilder.Mcp/Azure.DataApiBuilder.Mcp.csproj +++ b/src/Azure.DataApiBuilder.Mcp/Azure.DataApiBuilder.Mcp.csproj @@ -6,6 +6,10 @@ enable + + + + diff --git a/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs b/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs index ea2fa0cfea..f724d0d1ba 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/DynamicCustomTool.cs @@ -38,7 +38,6 @@ namespace Azure.DataApiBuilder.Mcp.Core /// public class DynamicCustomTool : IMcpTool { - private readonly string _entityName; private readonly Entity _entity; /// @@ -48,7 +47,7 @@ public class DynamicCustomTool : IMcpTool /// The entity configuration object. public DynamicCustomTool(string entityName, Entity entity) { - _entityName = entityName ?? throw new ArgumentNullException(nameof(entityName)); + EntityName = entityName ?? throw new ArgumentNullException(nameof(entityName)); _entity = entity ?? throw new ArgumentNullException(nameof(entity)); // Validate that this is a stored procedure @@ -65,12 +64,17 @@ public DynamicCustomTool(string entityName, Entity entity) /// public ToolType ToolType { get; } = ToolType.Custom; + /// + /// Gets the entity name associated with this custom tool. + /// + public string EntityName { get; } + /// /// Gets the metadata for this custom tool, including name, description, and input schema. /// public Tool GetToolMetadata() { - string toolName = ConvertToToolName(_entityName); + string toolName = ConvertToToolName(EntityName); string description = _entity.Description ?? $"Executes the {toolName} stored procedure"; // Build input schema based on parameters @@ -114,25 +118,25 @@ public async Task ExecuteAsync( } // 3) Validate entity still exists in configuration - if (!config.Entities.TryGetValue(_entityName, out Entity? entityConfig)) + if (!config.Entities.TryGetValue(EntityName, out Entity? entityConfig)) { - return McpResponseBuilder.BuildErrorResult(toolName, "EntityNotFound", $"Entity '{_entityName}' not found in configuration.", logger); + return McpResponseBuilder.BuildErrorResult(toolName, "EntityNotFound", $"Entity '{EntityName}' not found in configuration.", logger); } if (entityConfig.Source.Type != EntitySourceType.StoredProcedure) { - return McpResponseBuilder.BuildErrorResult(toolName, "InvalidEntity", $"Entity {_entityName} is not a stored procedure.", logger); + return McpResponseBuilder.BuildErrorResult(toolName, "InvalidEntity", $"Entity {EntityName} is not a stored procedure.", logger); } // Check if custom tool is still enabled for this entity if (entityConfig.Mcp?.CustomToolEnabled != true) { - return McpErrorHelpers.ToolDisabled(toolName, logger, $"Custom tool is disabled for entity '{_entityName}'."); + return McpErrorHelpers.ToolDisabled(toolName, logger, $"Custom tool is disabled for entity '{EntityName}'."); } // 4) Resolve metadata if (!McpMetadataHelper.TryResolveMetadata( - _entityName, + EntityName, config, serviceProvider, out ISqlMetadataProvider sqlMetadataProvider, @@ -150,18 +154,18 @@ public async Task ExecuteAsync( if (!McpAuthorizationHelper.ValidateRoleContext(httpContext, authResolver, out string roleError)) { - return McpErrorHelpers.PermissionDenied(toolName, _entityName, "execute", roleError, logger); + return McpErrorHelpers.PermissionDenied(toolName, EntityName, "execute", roleError, logger); } if (!McpAuthorizationHelper.TryResolveAuthorizedRole( httpContext!, authResolver, - _entityName, + EntityName, EntityActionOperation.Execute, out string? effectiveRole, out string authError)) { - return McpErrorHelpers.PermissionDenied(toolName, _entityName, "execute", authError, logger); + return McpErrorHelpers.PermissionDenied(toolName, EntityName, "execute", authError, logger); } // 6) Build request payload @@ -175,7 +179,7 @@ public async Task ExecuteAsync( // 7) Build stored procedure execution context StoredProcedureRequestContext context = new( - entityName: _entityName, + entityName: EntityName, dbo: dbObject, requestPayloadRoot: requestPayloadRoot, operationType: EntityActionOperation.Execute); @@ -218,7 +222,7 @@ public async Task ExecuteAsync( } catch (DataApiBuilderException dabEx) { - logger?.LogError(dabEx, "Error executing custom tool {ToolName} for entity {Entity}", toolName, _entityName); + logger?.LogError(dabEx, "Error executing custom tool {ToolName} for entity {Entity}", toolName, EntityName); return McpResponseBuilder.BuildErrorResult(toolName, "ExecutionError", dabEx.Message, logger); } catch (SqlException sqlEx) @@ -238,7 +242,7 @@ public async Task ExecuteAsync( } // 9) Build success response - return BuildExecuteSuccessResponse(toolName, _entityName, parameters, queryResult, logger); + return BuildExecuteSuccessResponse(toolName, EntityName, parameters, queryResult, logger); } catch (OperationCanceledException) { @@ -246,7 +250,7 @@ public async Task ExecuteAsync( } catch (Exception ex) { - logger?.LogError(ex, "Unexpected error in DynamicCustomTool for {EntityName}", _entityName); + logger?.LogError(ex, "Unexpected error in DynamicCustomTool for {EntityName}", EntityName); return McpResponseBuilder.BuildErrorResult(toolName, "UnexpectedError", "An unexpected error occurred.", logger); } } diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpServerConfiguration.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpServerConfiguration.cs index d76af816bd..bcba1a50e4 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpServerConfiguration.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpServerConfiguration.cs @@ -3,6 +3,7 @@ using System.Text.Json; using Azure.DataApiBuilder.Mcp.Model; +using Azure.DataApiBuilder.Mcp.Utils; using Microsoft.Extensions.DependencyInjection; using ModelContextProtocol; using ModelContextProtocol.Protocol; @@ -61,22 +62,23 @@ internal static IServiceCollection ConfigureMcpServer(this IServiceCollection se } JsonDocument? arguments = null; - if (request.Params?.Arguments != null) + try { - // Convert IReadOnlyDictionary to JsonDocument - Dictionary jsonObject = new(); - foreach (KeyValuePair kvp in request.Params.Arguments) + if (request.Params?.Arguments != null) { - jsonObject[kvp.Key] = kvp.Value; - } + // Convert IReadOnlyDictionary to JsonDocument + Dictionary jsonObject = new(); + foreach (KeyValuePair kvp in request.Params.Arguments) + { + jsonObject[kvp.Key] = kvp.Value; + } - string json = JsonSerializer.Serialize(jsonObject); - arguments = JsonDocument.Parse(json); - } + string json = JsonSerializer.Serialize(jsonObject); + arguments = JsonDocument.Parse(json); + } - try - { - return await tool!.ExecuteAsync(arguments, request.Services!, ct); + return await McpTelemetryHelper.ExecuteWithTelemetryAsync( + tool!, toolName, arguments, request.Services!, ct); } finally { diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs index 51d8295068..1ab1c73d05 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs @@ -7,6 +7,7 @@ using Azure.DataApiBuilder.Core.AuthenticationHelpers.AuthenticationSimulator; using Azure.DataApiBuilder.Core.Configurations; using Azure.DataApiBuilder.Mcp.Model; +using Azure.DataApiBuilder.Mcp.Utils; using Microsoft.AspNetCore.Http; using Microsoft.Extensions.Configuration; using Microsoft.Extensions.DependencyInjection; @@ -284,7 +285,7 @@ private async Task HandleCallToolAsync(JsonElement? id, JsonElement root, Cancel Console.Error.WriteLine($"[MCP DEBUG] callTool → tool: {toolName}, args: "); } - // Execute the tool. + // Execute the tool with telemetry. // If a MCP stdio role override is set in the environment, create // a request HttpContext with the X-MS-API-ROLE header so tools and authorization // helpers that read IHttpContextAccessor will see the role. We also ensure the @@ -319,7 +320,8 @@ private async Task HandleCallToolAsync(JsonElement? id, JsonElement root, Cancel try { // Execute the tool with the scoped service provider so any scoped services resolve correctly. - callResult = await tool.ExecuteAsync(argsDoc, scopedProvider, ct); + callResult = await McpTelemetryHelper.ExecuteWithTelemetryAsync( + tool, toolName!, argsDoc, scopedProvider, ct); } finally { @@ -332,7 +334,8 @@ private async Task HandleCallToolAsync(JsonElement? id, JsonElement root, Cancel } else { - callResult = await tool.ExecuteAsync(argsDoc, _serviceProvider, ct); + callResult = await McpTelemetryHelper.ExecuteWithTelemetryAsync( + tool, toolName!, argsDoc, _serviceProvider, ct); } // Normalize to MCP content blocks (array). We try to pass through if a 'Content' property exists, diff --git a/src/Azure.DataApiBuilder.Mcp/Utils/McpTelemetryErrorCodes.cs b/src/Azure.DataApiBuilder.Mcp/Utils/McpTelemetryErrorCodes.cs new file mode 100644 index 0000000000..f69a26fa5d --- /dev/null +++ b/src/Azure.DataApiBuilder.Mcp/Utils/McpTelemetryErrorCodes.cs @@ -0,0 +1,41 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +namespace Azure.DataApiBuilder.Mcp.Utils +{ + /// + /// Constants for MCP telemetry error codes. + /// + internal static class McpTelemetryErrorCodes + { + /// + /// Generic execution failure error code. + /// + public const string EXECUTION_FAILED = "ExecutionFailed"; + + /// + /// Authentication failure error code. + /// + public const string AUTHENTICATION_FAILED = "AuthenticationFailed"; + + /// + /// Authorization failure error code. + /// + public const string AUTHORIZATION_FAILED = "AuthorizationFailed"; + + /// + /// Database operation failure error code. + /// + public const string DATABASE_ERROR = "DatabaseError"; + + /// + /// Invalid request or arguments error code. + /// + public const string INVALID_REQUEST = "InvalidRequest"; + + /// + /// Operation cancelled error code. + /// + public const string OPERATION_CANCELLED = "OperationCancelled"; + } +} diff --git a/src/Azure.DataApiBuilder.Mcp/Utils/McpTelemetryHelper.cs b/src/Azure.DataApiBuilder.Mcp/Utils/McpTelemetryHelper.cs new file mode 100644 index 0000000000..1e1851c8f1 --- /dev/null +++ b/src/Azure.DataApiBuilder.Mcp/Utils/McpTelemetryHelper.cs @@ -0,0 +1,261 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.Diagnostics; +using System.Text.Json; +using Azure.DataApiBuilder.Config.ObjectModel; +using Azure.DataApiBuilder.Core.Configurations; +using Azure.DataApiBuilder.Core.Telemetry; +using Azure.DataApiBuilder.Mcp.Core; +using Azure.DataApiBuilder.Mcp.Model; +using Azure.DataApiBuilder.Service.Exceptions; +using Microsoft.Extensions.DependencyInjection; +using ModelContextProtocol.Protocol; +using static Azure.DataApiBuilder.Mcp.Model.McpEnums; + +namespace Azure.DataApiBuilder.Mcp.Utils +{ + /// + /// Utility class for MCP telemetry operations. + /// + internal static class McpTelemetryHelper + { + /// + /// Executes an MCP tool wrapped in an OpenTelemetry activity span. + /// Handles telemetry attribute extraction, success/failure tracking, + /// and exception recording with typed error codes. + /// + /// The MCP tool to execute. + /// The name of the tool being invoked. + /// The parsed JSON arguments for the tool (may be null). + /// The service provider for resolving dependencies. + /// Cancellation token. + /// The result of the tool execution. + public static async Task ExecuteWithTelemetryAsync( + IMcpTool tool, + string toolName, + JsonDocument? arguments, + IServiceProvider serviceProvider, + CancellationToken cancellationToken) + { + using Activity? activity = TelemetryTracesHelper.DABActivitySource.StartActivity("mcp.tool.execute"); + + try + { + // Extract telemetry metadata + string? entityName = ExtractEntityNameFromArguments(arguments); + string? operation = InferOperationFromTool(tool, toolName); + string? dbProcedure = null; + + // For custom tools (DynamicCustomTool), extract stored procedure information + if (tool is DynamicCustomTool customTool) + { + (entityName, dbProcedure) = ExtractCustomToolMetadata(customTool, serviceProvider); + } + + // Track the start of MCP tool execution with telemetry + activity?.TrackMcpToolExecutionStarted( + toolName: toolName, + entityName: entityName, + 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 + // and return CallToolResult with IsError=true instead of throwing) + if (result.IsError == true) + { + // Extract error code and message from the result content + (string? errorCode, string? errorMessage) = ExtractErrorFromCallToolResult(result); + + activity?.SetStatus(ActivityStatusCode.Error, errorMessage ?? "Tool returned an error result"); + activity?.SetTag("mcp.tool.error", true); + + if (!string.IsNullOrEmpty(errorCode)) + { + activity?.SetTag("error.code", errorCode); + } + + if (!string.IsNullOrEmpty(errorMessage)) + { + activity?.SetTag("error.message", errorMessage); + } + } + else + { + // Track successful completion + activity?.TrackMcpToolExecutionFinished(); + } + + return result; + } + catch (Exception ex) + { + // Track exception in telemetry with specific error code based on exception type + string errorCode = MapExceptionToErrorCode(ex); + activity?.TrackMcpToolExecutionFinishedWithException(ex, errorCode: errorCode); + throw; + } + } + + /// + /// Infers the operation type from the tool instance and name. + /// For built-in tools, maps tool name directly to operation. + /// For custom tools (stored procedures), always returns "execute". + /// + /// The tool instance. + /// The name of the tool. + /// The inferred operation type. + public static string InferOperationFromTool(IMcpTool tool, string toolName) + { + // Custom tools (stored procedures) are always "execute" + if (tool.ToolType == ToolType.Custom) + { + return "execute"; + } + + // Built-in tools: map tool name to operation + return toolName.ToLowerInvariant() switch + { + "read_records" => "read", + "create_record" => "create", + "update_record" => "update", + "delete_record" => "delete", + "describe_entities" => "describe", + "execute_entity" => "execute", + _ => "execute" // Fallback for any unknown built-in tools + }; + } + + /// + /// Extracts error code and message from a CallToolResult's content. + /// MCP tools may return errors as JSON with "code" and "message" properties. + /// + /// The tool result to extract error info from. + /// A tuple of (errorCode, errorMessage). + private static (string? errorCode, string? errorMessage) ExtractErrorFromCallToolResult(CallToolResult result) + { + string? errorCode = null; + string? errorMessage = null; + + if (result.Content != null) + { + foreach (ContentBlock block in result.Content) + { + // Check if this is a text block with JSON error information + if (block is TextContentBlock textBlock && !string.IsNullOrEmpty(textBlock.Text)) + { + try + { + using JsonDocument doc = JsonDocument.Parse(textBlock.Text); + JsonElement root = doc.RootElement; + + if (root.TryGetProperty("code", out JsonElement codeEl)) + { + errorCode = codeEl.GetString(); + } + + if (root.TryGetProperty("message", out JsonElement msgEl)) + { + errorMessage = msgEl.GetString(); + } + + // If we found error info, we can break + if (errorCode != null || errorMessage != null) + { + break; + } + } + catch + { + // Not JSON or doesn't have expected structure, skip + } + } + } + } + + return (errorCode, errorMessage); + } + + /// + /// Maps an exception to a telemetry error code. + /// + /// The exception to map. + /// The corresponding error code string. + public static string MapExceptionToErrorCode(Exception ex) + { + return ex switch + { + OperationCanceledException => McpTelemetryErrorCodes.OPERATION_CANCELLED, + DataApiBuilderException dabEx when dabEx.SubStatusCode == DataApiBuilderException.SubStatusCodes.AuthenticationChallenge + => McpTelemetryErrorCodes.AUTHENTICATION_FAILED, + DataApiBuilderException dabEx when dabEx.SubStatusCode == DataApiBuilderException.SubStatusCodes.AuthorizationCheckFailed + => McpTelemetryErrorCodes.AUTHORIZATION_FAILED, + UnauthorizedAccessException => McpTelemetryErrorCodes.AUTHORIZATION_FAILED, + System.Data.Common.DbException => McpTelemetryErrorCodes.DATABASE_ERROR, + ArgumentException => McpTelemetryErrorCodes.INVALID_REQUEST, + _ => McpTelemetryErrorCodes.EXECUTION_FAILED + }; + } + + /// + /// Extracts the entity name from parsed tool arguments, if present. + /// + /// The parsed JSON arguments. + /// The entity name, or null if not present. + private static string? ExtractEntityNameFromArguments(JsonDocument? arguments) + { + if (arguments != null && + arguments.RootElement.TryGetProperty("entity", out JsonElement entityEl) && + entityEl.ValueKind == JsonValueKind.String) + { + return entityEl.GetString(); + } + + return null; + } + + /// + /// Extracts metadata from a custom tool for telemetry purposes. + /// Returns best-effort metadata; failures in configuration access must not prevent tool execution. + /// + /// The custom tool instance. + /// The service provider. + /// A tuple containing the entity name and database procedure name. + public static (string? entityName, string? dbProcedure) ExtractCustomToolMetadata(DynamicCustomTool customTool, IServiceProvider serviceProvider) + { + // Access public properties instead of reflection + string? entityName = customTool.EntityName; + + if (entityName == null) + { + return (null, null); + } + + try + { + // Try to get the stored procedure name from the runtime configuration + RuntimeConfigProvider? runtimeConfigProvider = serviceProvider.GetService(); + if (runtimeConfigProvider != null) + { + RuntimeConfig config = runtimeConfigProvider.GetConfig(); + if (config.Entities.TryGetValue(entityName, out Entity? entityConfig)) + { + string? dbProcedure = entityConfig.Source.Object; + return (entityName, dbProcedure); + } + } + } + catch (Exception) + { + // If configuration access fails for any reason (including DataApiBuilderException + // when runtime config isn't set up), fall back to returning only the entity name. + // Telemetry metadata extraction is best-effort and must not prevent tool execution. + } + + return (entityName, null); + } + } +} diff --git a/src/Core/Telemetry/TelemetryTracesHelper.cs b/src/Core/Telemetry/TelemetryTracesHelper.cs index 01c5acbf51..a6b0ef2b0d 100644 --- a/src/Core/Telemetry/TelemetryTracesHelper.cs +++ b/src/Core/Telemetry/TelemetryTracesHelper.cs @@ -111,5 +111,78 @@ public static void TrackMainControllerActivityFinishedWithException( activity.SetTag("status.code", statusCode); } } + + /// + /// Tracks the start of an MCP tool execution activity. + /// + /// The activity instance. + /// The name of the MCP tool being executed. + /// The entity name associated with the tool (optional). + /// The operation being performed (e.g., execute, read, create). + /// The database procedure being executed (optional, schema-qualified if available). + public static void TrackMcpToolExecutionStarted( + this Activity activity, + string toolName, + string? entityName = null, + string? operation = null, + string? dbProcedure = null) + { + if (activity.IsAllDataRequested) + { + activity.SetTag("mcp.tool.name", toolName); + + if (!string.IsNullOrEmpty(entityName)) + { + activity.SetTag("dab.entity", entityName); + } + + if (!string.IsNullOrEmpty(operation)) + { + activity.SetTag("dab.operation", operation); + } + + if (!string.IsNullOrEmpty(dbProcedure)) + { + activity.SetTag("db.procedure", dbProcedure); + } + } + } + + /// + /// Tracks the successful completion of an MCP tool execution. + /// + /// The activity instance. + public static void TrackMcpToolExecutionFinished(this Activity activity) + { + if (activity.IsAllDataRequested) + { + activity.SetStatus(ActivityStatusCode.Ok); + } + } + + /// + /// Tracks the completion of an MCP tool execution with an exception. + /// + /// The activity instance. + /// The exception that occurred. + /// Optional error code for the failure. + public static void TrackMcpToolExecutionFinishedWithException( + this Activity activity, + Exception ex, + string? errorCode = null) + { + if (activity.IsAllDataRequested) + { + activity.SetStatus(ActivityStatusCode.Error, ex.Message); + activity.RecordException(ex); + activity.SetTag("error.type", ex.GetType().Name); + activity.SetTag("error.message", ex.Message); + + if (!string.IsNullOrEmpty(errorCode)) + { + activity.SetTag("error.code", errorCode); + } + } + } } } diff --git a/src/Service.Tests/Azure.DataApiBuilder.Service.Tests.csproj b/src/Service.Tests/Azure.DataApiBuilder.Service.Tests.csproj index d250822359..ae274a4dc2 100644 --- a/src/Service.Tests/Azure.DataApiBuilder.Service.Tests.csproj +++ b/src/Service.Tests/Azure.DataApiBuilder.Service.Tests.csproj @@ -97,6 +97,7 @@ + diff --git a/src/Service.Tests/UnitTests/McpTelemetryTests.cs b/src/Service.Tests/UnitTests/McpTelemetryTests.cs new file mode 100644 index 0000000000..5066c73f92 --- /dev/null +++ b/src/Service.Tests/UnitTests/McpTelemetryTests.cs @@ -0,0 +1,383 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +#nullable enable + +using System; +using System.Collections.Generic; +using System.Diagnostics; +using System.Linq; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using Azure.DataApiBuilder.Core.Telemetry; +using Azure.DataApiBuilder.Mcp.Model; +using Azure.DataApiBuilder.Mcp.Utils; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.VisualStudio.TestTools.UnitTesting; +using ModelContextProtocol.Protocol; +using static Azure.DataApiBuilder.Mcp.Model.McpEnums; + +namespace Azure.DataApiBuilder.Service.Tests.UnitTests +{ + /// + /// Tests for MCP telemetry functionality. + /// + [TestClass] + public class McpTelemetryTests + { + private static ActivityListener? _activityListener; + private static readonly List _recordedActivities = new(); + + /// + /// Initialize activity listener before all tests. + /// + [ClassInitialize] + public static void ClassInitialize(TestContext context) + { + _activityListener = new ActivityListener + { + ShouldListenTo = (activitySource) => activitySource.Name == "DataApiBuilder", + Sample = (ref ActivityCreationOptions options) => ActivitySamplingResult.AllDataAndRecorded, + ActivityStarted = activity => { }, + ActivityStopped = activity => + { + _recordedActivities.Add(activity); + } + }; + ActivitySource.AddActivityListener(_activityListener); + } + + /// + /// Cleanup activity listener after all tests. + /// + [ClassCleanup] + public static void ClassCleanup() + { + _activityListener?.Dispose(); + } + + /// + /// Clear recorded activities before each test. + /// + [TestInitialize] + public void TestInitialize() + { + _recordedActivities.Clear(); + } + + #region Helpers + + /// + /// Creates and starts a new MCP tool execution activity, asserting it was created. + /// + private static Activity CreateActivity() + { + Activity? activity = TelemetryTracesHelper.DABActivitySource.StartActivity("mcp.tool.execute"); + Assert.IsNotNull(activity, "Activity should be created"); + return activity; + } + + /// + /// Stops the activity and returns the first recorded activity, asserting it was captured. + /// + private static Activity StopAndGetRecordedActivity(Activity activity) + { + activity.Stop(); + Activity? recorded = _recordedActivities.FirstOrDefault(); + Assert.IsNotNull(recorded, "Activity should be recorded"); + return recorded; + } + + /// + /// Builds a minimal service provider for tests that don't need real services. + /// + private static IServiceProvider CreateServiceProvider() + { + return new ServiceCollection().BuildServiceProvider(); + } + + /// + /// Creates a CallToolResult with the given text and error state. + /// + private static CallToolResult CreateToolResult(string text = "result", bool isError = false) + { + return new CallToolResult + { + Content = new List { new TextContentBlock { Type = "text", Text = text } }, + IsError = isError + }; + } + + /// + /// Creates an exception instance from a type name string, for use with DataRow tests. + /// + private static Exception CreateExceptionByTypeName(string typeName) + { + return typeName switch + { + nameof(OperationCanceledException) => new OperationCanceledException(), + nameof(UnauthorizedAccessException) => new UnauthorizedAccessException(), + nameof(ArgumentException) => new ArgumentException(), + nameof(InvalidOperationException) => new InvalidOperationException(), + _ => new Exception() + }; + } + + #endregion + + #region TrackMcpToolExecutionStarted + + /// + /// Test that TrackMcpToolExecutionStarted sets the expected tags for various input combinations, + /// including when optional parameters are null. + /// + [DataTestMethod] + [DataRow("read_records", "books", "read", null, DisplayName = "Sets entity, operation; no procedure")] + [DataRow("custom_proc", "CustomEntity", "execute", "dbo.CustomProc", DisplayName = "Custom tool with all tags including db.procedure")] + [DataRow("describe_entities", null, "describe", null, DisplayName = "Describe tool with null entity")] + [DataRow("custom_tool", "MyEntity", "execute", "schema.MyStoredProc", DisplayName = "Sets all four tags")] + public void TrackMcpToolExecutionStarted_SetsExpectedTags( + string toolName, string? entityName, string? operation, string? dbProcedure) + { + // Arrange & Act + using Activity activity = CreateActivity(); + activity.TrackMcpToolExecutionStarted( + toolName: toolName, + entityName: entityName, + operation: operation, + dbProcedure: dbProcedure); + + Activity recorded = StopAndGetRecordedActivity(activity); + + // Assert — tool name is always set + Assert.AreEqual(toolName, recorded.GetTagItem("mcp.tool.name")); + + // Optional tags: present only when supplied + Assert.AreEqual(entityName, recorded.GetTagItem("dab.entity")); + Assert.AreEqual(operation, recorded.GetTagItem("dab.operation")); + Assert.AreEqual(dbProcedure, recorded.GetTagItem("db.procedure")); + } + + #endregion + + #region TrackMcpToolExecutionFinished + + /// + /// Test that TrackMcpToolExecutionFinished sets status to OK. + /// + [TestMethod] + public void TrackMcpToolExecutionFinished_SetsStatusToOk() + { + using Activity activity = CreateActivity(); + activity.TrackMcpToolExecutionStarted(toolName: "read_records"); + activity.TrackMcpToolExecutionFinished(); + + Activity recorded = StopAndGetRecordedActivity(activity); + Assert.AreEqual(ActivityStatusCode.Ok, recorded.Status); + } + + /// + /// Test that TrackMcpToolExecutionFinishedWithException records exception and sets error status. + /// + [TestMethod] + public void TrackMcpToolExecutionFinishedWithException_RecordsExceptionAndSetsErrorStatus() + { + using Activity activity = CreateActivity(); + activity.TrackMcpToolExecutionStarted(toolName: "read_records"); + + Exception testException = new InvalidOperationException("Test exception"); + activity.TrackMcpToolExecutionFinishedWithException(testException, errorCode: McpTelemetryErrorCodes.EXECUTION_FAILED); + + Activity recorded = StopAndGetRecordedActivity(activity); + Assert.AreEqual(ActivityStatusCode.Error, recorded.Status); + Assert.AreEqual("Test exception", recorded.StatusDescription); + Assert.AreEqual("InvalidOperationException", recorded.GetTagItem("error.type")); + Assert.AreEqual("Test exception", recorded.GetTagItem("error.message")); + Assert.AreEqual(McpTelemetryErrorCodes.EXECUTION_FAILED, recorded.GetTagItem("error.code")); + + ActivityEvent? exceptionEvent = recorded.Events.FirstOrDefault(e => e.Name == "exception"); + Assert.IsNotNull(exceptionEvent, "Exception event should be recorded"); + } + + #endregion + + #region InferOperationFromTool + + /// + /// Test that InferOperationFromTool returns the correct operation for built-in and custom tools. + /// Built-in tools are mapped by name; custom tools always return "execute". + /// + [DataTestMethod] + // Built-in DML tool names mapped to operations + [DataRow(ToolType.BuiltIn, "read_records", "read", DisplayName = "Built-in: read_records -> read")] + [DataRow(ToolType.BuiltIn, "create_record", "create", DisplayName = "Built-in: create_record -> create")] + [DataRow(ToolType.BuiltIn, "update_record", "update", DisplayName = "Built-in: update_record -> update")] + [DataRow(ToolType.BuiltIn, "delete_record", "delete", DisplayName = "Built-in: delete_record -> delete")] + [DataRow(ToolType.BuiltIn, "describe_entities", "describe", DisplayName = "Built-in: describe_entities -> describe")] + [DataRow(ToolType.BuiltIn, "execute_entity", "execute", DisplayName = "Built-in: execute_entity -> execute")] + [DataRow(ToolType.BuiltIn, "unknown_builtin", "execute", DisplayName = "Built-in: unknown -> execute (fallback)")] + // Custom tools always return "execute" + [DataRow(ToolType.Custom, "get_book", "execute", DisplayName = "Custom: get_book -> execute (stored proc)")] + [DataRow(ToolType.Custom, "read_users", "execute", DisplayName = "Custom: read_users -> execute (ignore name)")] + [DataRow(ToolType.Custom, "custom_proc", "execute", DisplayName = "Custom: custom_proc -> execute")] + public void InferOperationFromTool_ReturnsCorrectOperation(ToolType toolType, string toolName, string expectedOperation) + { + IMcpTool tool = new MockMcpTool(CreateToolResult(), toolType); + Assert.AreEqual(expectedOperation, McpTelemetryHelper.InferOperationFromTool(tool, toolName)); + } + + #endregion + + #region MapExceptionToErrorCode + + /// + /// Test that MapExceptionToErrorCode returns the correct error code for each exception type. + /// + [DataTestMethod] + [DataRow("OperationCanceledException", McpTelemetryErrorCodes.OPERATION_CANCELLED)] + [DataRow("UnauthorizedAccessException", McpTelemetryErrorCodes.AUTHORIZATION_FAILED)] + [DataRow("ArgumentException", McpTelemetryErrorCodes.INVALID_REQUEST)] + [DataRow("InvalidOperationException", McpTelemetryErrorCodes.EXECUTION_FAILED)] + [DataRow("Exception", McpTelemetryErrorCodes.EXECUTION_FAILED)] + public void MapExceptionToErrorCode_ReturnsCorrectCode(string exceptionTypeName, string expectedErrorCode) + { + Exception ex = CreateExceptionByTypeName(exceptionTypeName); + Assert.AreEqual(expectedErrorCode, McpTelemetryHelper.MapExceptionToErrorCode(ex)); + } + + #endregion + + #region ExecuteWithTelemetryAsync + + /// + /// Test that ExecuteWithTelemetryAsync sets Ok status and correct operation for all built-in DML tools. + /// + [DataTestMethod] + [DataRow("read_records", "read", DisplayName = "read_records -> read operation")] + [DataRow("create_record", "create", DisplayName = "create_record -> create operation")] + [DataRow("update_record", "update", DisplayName = "update_record -> update operation")] + [DataRow("delete_record", "delete", DisplayName = "delete_record -> delete operation")] + [DataRow("describe_entities", "describe", DisplayName = "describe_entities -> describe operation")] + [DataRow("execute_entity", "execute", DisplayName = "execute_entity -> execute operation")] + public async Task ExecuteWithTelemetryAsync_SetsOkStatusAndCorrectOperation_ForBuiltInTools( + string toolName, string expectedOperation) + { + CallToolResult expectedResult = CreateToolResult("success"); + IMcpTool tool = new MockMcpTool(expectedResult, ToolType.BuiltIn); + + CallToolResult result = await McpTelemetryHelper.ExecuteWithTelemetryAsync( + tool, toolName, arguments: null, CreateServiceProvider(), CancellationToken.None); + + Assert.AreSame(expectedResult, result); + Activity recorded = _recordedActivities.First(); + Assert.AreEqual(ActivityStatusCode.Ok, recorded.Status); + Assert.AreEqual(toolName, recorded.GetTagItem("mcp.tool.name")); + Assert.AreEqual(expectedOperation, recorded.GetTagItem("dab.operation")); + } + + /// + /// Test that ExecuteWithTelemetryAsync always sets operation to "execute" for custom tools (stored procedures). + /// + [TestMethod] + public async Task ExecuteWithTelemetryAsync_SetsExecuteOperation_ForCustomTools() + { + CallToolResult expectedResult = CreateToolResult("success"); + IMcpTool tool = new MockMcpTool(expectedResult, ToolType.Custom); + + CallToolResult result = await McpTelemetryHelper.ExecuteWithTelemetryAsync( + tool, "get_book", arguments: null, CreateServiceProvider(), CancellationToken.None); + + Assert.AreSame(expectedResult, result); + Activity recorded = _recordedActivities.First(); + Assert.AreEqual(ActivityStatusCode.Ok, recorded.Status); + Assert.AreEqual("get_book", recorded.GetTagItem("mcp.tool.name")); + Assert.AreEqual("execute", recorded.GetTagItem("dab.operation")); + } + + /// + /// Test that ExecuteWithTelemetryAsync sets Error status when tool returns IsError=true. + /// + [TestMethod] + public async Task ExecuteWithTelemetryAsync_SetsErrorStatus_WhenToolReturnsIsError() + { + CallToolResult errorResult = CreateToolResult("error occurred", isError: true); + IMcpTool tool = new MockMcpTool(errorResult, ToolType.BuiltIn); + + CallToolResult result = await McpTelemetryHelper.ExecuteWithTelemetryAsync( + tool, "create_record", arguments: null, CreateServiceProvider(), CancellationToken.None); + + Assert.AreSame(errorResult, result); + Activity recorded = _recordedActivities.First(); + Assert.AreEqual(ActivityStatusCode.Error, recorded.Status); + Assert.AreEqual(true, recorded.GetTagItem("mcp.tool.error")); + } + + /// + /// Test that ExecuteWithTelemetryAsync records exception and re-throws when tool throws. + /// + [TestMethod] + public async Task ExecuteWithTelemetryAsync_RecordsExceptionAndRethrows_WhenToolThrows() + { + InvalidOperationException expectedException = new("tool exploded"); + IMcpTool tool = new MockMcpTool(expectedException, ToolType.BuiltIn); + + InvalidOperationException thrownEx = await Assert.ThrowsExceptionAsync( + () => McpTelemetryHelper.ExecuteWithTelemetryAsync( + tool, "delete_record", arguments: null, CreateServiceProvider(), CancellationToken.None)); + + Assert.AreEqual("tool exploded", thrownEx.Message); + + Activity recorded = _recordedActivities.First(); + Assert.AreEqual(ActivityStatusCode.Error, recorded.Status); + Assert.AreEqual("InvalidOperationException", recorded.GetTagItem("error.type")); + Assert.AreEqual(McpTelemetryErrorCodes.EXECUTION_FAILED, recorded.GetTagItem("error.code")); + + ActivityEvent? exceptionEvent = recorded.Events.FirstOrDefault(e => e.Name == "exception"); + Assert.IsNotNull(exceptionEvent, "Exception event should be recorded"); + } + + #endregion + + #region Test Mocks + + /// + /// A minimal mock IMcpTool for testing ExecuteWithTelemetryAsync. + /// Returns a predetermined result or throws a predetermined exception. + /// + private class MockMcpTool : IMcpTool + { + private readonly CallToolResult? _result; + private readonly Exception? _exception; + private readonly ToolType _toolType; + + public MockMcpTool(CallToolResult result, ToolType toolType = ToolType.BuiltIn) + { + _result = result; + _toolType = toolType; + } + + public MockMcpTool(Exception exception, ToolType toolType = ToolType.BuiltIn) + { + _exception = exception; + _toolType = toolType; + } + + public ToolType ToolType => _toolType; + + public Tool GetToolMetadata() => new() { Name = "mock_tool", Description = "Mock tool for testing" }; + + public Task ExecuteAsync(JsonDocument? arguments, IServiceProvider serviceProvider, CancellationToken cancellationToken = default) + { + if (_exception != null) + { + throw _exception; + } + + return Task.FromResult(_result!); + } + } + + #endregion + } +}