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..acb8804172 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpServerConfiguration.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpServerConfiguration.cs @@ -1,8 +1,11 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +using System.Diagnostics; using System.Text.Json; +using Azure.DataApiBuilder.Core.Telemetry; using Azure.DataApiBuilder.Mcp.Model; +using Azure.DataApiBuilder.Mcp.Utils; using Microsoft.Extensions.DependencyInjection; using ModelContextProtocol; using ModelContextProtocol.Protocol; @@ -60,23 +63,75 @@ internal static IServiceCollection ConfigureMcpServer(this IServiceCollection se throw new McpException($"Unknown tool: '{toolName}'"); } + // Start OpenTelemetry activity for MCP tool execution + using Activity? activity = TelemetryTracesHelper.DABActivitySource.StartActivity("mcp.tool.execute"); + JsonDocument? arguments = null; - if (request.Params?.Arguments != null) + try { - // Convert IReadOnlyDictionary to JsonDocument - Dictionary jsonObject = new(); - foreach (KeyValuePair kvp in request.Params.Arguments) + // Extract entity name from arguments for telemetry + string? entityName = null; + string? operation = null; + string? dbProcedure = null; + + 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; + + // Extract entity name if present + if (kvp.Key == "entity" && kvp.Value.ValueKind == JsonValueKind.String) + { + entityName = kvp.Value.GetString(); + } + } + + string json = JsonSerializer.Serialize(jsonObject); + arguments = JsonDocument.Parse(json); } - string json = JsonSerializer.Serialize(jsonObject); - arguments = JsonDocument.Parse(json); - } + // Determine operation based on tool name + operation = McpTelemetryHelper.InferOperationFromToolName(toolName); - try + // For custom tools (DynamicCustomTool), extract stored procedure information + if (tool is DynamicCustomTool customTool) + { + // Get entity name and procedure from the custom tool + (entityName, dbProcedure) = McpTelemetryHelper.ExtractCustomToolMetadata(customTool, request.Services!); + } + + // 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, request.Services!, ct); + + // Track successful completion + activity?.TrackMcpToolExecutionFinished(); + + return result; + } + catch (Exception ex) { - return await tool!.ExecuteAsync(arguments, request.Services!, ct); + // Track exception in telemetry with specific error code based on exception type + string errorCode = ex switch + { + OperationCanceledException => McpTelemetryErrorCodes.OPERATION_CANCELLED, + UnauthorizedAccessException => McpTelemetryErrorCodes.AUTHENTICATION_FAILED, + System.Data.Common.DbException => McpTelemetryErrorCodes.DATABASE_ERROR, + ArgumentException => McpTelemetryErrorCodes.INVALID_REQUEST, + _ => McpTelemetryErrorCodes.EXECUTION_FAILED + }; + + activity?.TrackMcpToolExecutionFinishedWithException(ex, errorCode: errorCode); + throw; } finally { diff --git a/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs b/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs index 51d8295068..461491208f 100644 --- a/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs +++ b/src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs @@ -1,4 +1,5 @@ using System.Collections; +using System.Diagnostics; using System.Reflection; using System.Security.Claims; using System.Text; @@ -6,7 +7,9 @@ using Azure.DataApiBuilder.Config.ObjectModel; using Azure.DataApiBuilder.Core.AuthenticationHelpers.AuthenticationSimulator; using Azure.DataApiBuilder.Core.Configurations; +using Azure.DataApiBuilder.Core.Telemetry; using Azure.DataApiBuilder.Mcp.Model; +using Azure.DataApiBuilder.Mcp.Utils; using Microsoft.AspNetCore.Http; using Microsoft.Extensions.Configuration; using Microsoft.Extensions.DependencyInjection; @@ -270,20 +273,51 @@ private async Task HandleCallToolAsync(JsonElement? id, JsonElement root, Cancel return; } + // Start OpenTelemetry activity for MCP tool execution + using Activity? activity = TelemetryTracesHelper.DABActivitySource.StartActivity("mcp.tool.execute"); + JsonDocument? argsDoc = null; try { + // Extract entity name and operation from arguments for telemetry + string? entityName = null; + string? operation = null; + string? dbProcedure = null; + if (@params.TryGetProperty("arguments", out JsonElement argsEl) && argsEl.ValueKind == JsonValueKind.Object) { string rawArgs = argsEl.GetRawText(); Console.Error.WriteLine($"[MCP DEBUG] callTool → tool: {toolName}, args: {rawArgs}"); argsDoc = JsonDocument.Parse(rawArgs); + + // Extract entity name if present in arguments + if (argsDoc.RootElement.TryGetProperty("entity", out JsonElement entityEl) && entityEl.ValueKind == JsonValueKind.String) + { + entityName = entityEl.GetString(); + } } else { Console.Error.WriteLine($"[MCP DEBUG] callTool → tool: {toolName}, args: "); } + // Determine operation based on tool name + operation = McpTelemetryHelper.InferOperationFromToolName(toolName!); + + // For custom tools (DynamicCustomTool), extract stored procedure information + if (tool is DynamicCustomTool customTool) + { + // Get entity name and procedure from the custom tool + (entityName, dbProcedure) = McpTelemetryHelper.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. // 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 @@ -335,12 +369,30 @@ private async Task HandleCallToolAsync(JsonElement? id, JsonElement root, Cancel callResult = await tool.ExecuteAsync(argsDoc, _serviceProvider, ct); } + // Track successful completion + activity?.TrackMcpToolExecutionFinished(); + // Normalize to MCP content blocks (array). We try to pass through if a 'Content' property exists, // otherwise we wrap into a single text block. object[] content = CoerceToMcpContentBlocks(callResult); WriteResult(id, new { content }); } + catch (Exception ex) + { + // Track exception in telemetry with specific error code based on exception type + string errorCode = ex switch + { + OperationCanceledException => McpTelemetryErrorCodes.OPERATION_CANCELLED, + UnauthorizedAccessException => McpTelemetryErrorCodes.AUTHENTICATION_FAILED, + System.Data.Common.DbException => McpTelemetryErrorCodes.DATABASE_ERROR, + ArgumentException => McpTelemetryErrorCodes.INVALID_REQUEST, + _ => McpTelemetryErrorCodes.EXECUTION_FAILED + }; + + activity?.TrackMcpToolExecutionFinishedWithException(ex, errorCode: errorCode); + throw; + } finally { argsDoc?.Dispose(); diff --git a/src/Azure.DataApiBuilder.Mcp/Utils/McpTelemetryErrorCodes.cs b/src/Azure.DataApiBuilder.Mcp/Utils/McpTelemetryErrorCodes.cs new file mode 100644 index 0000000000..47380061e4 --- /dev/null +++ b/src/Azure.DataApiBuilder.Mcp/Utils/McpTelemetryErrorCodes.cs @@ -0,0 +1,36 @@ +// 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 or authorization failure error code. + /// + public const string AUTHENTICATION_FAILED = "AuthenticationFailed"; + + /// + /// 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..d219dbdc99 --- /dev/null +++ b/src/Azure.DataApiBuilder.Mcp/Utils/McpTelemetryHelper.cs @@ -0,0 +1,71 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using Azure.DataApiBuilder.Config.ObjectModel; +using Azure.DataApiBuilder.Core.Configurations; +using Microsoft.Extensions.DependencyInjection; + +namespace Azure.DataApiBuilder.Mcp.Utils +{ + /// + /// Utility class for MCP telemetry operations. + /// + internal static class McpTelemetryHelper + { + /// + /// Infers the operation type from the tool name. + /// + /// The name of the tool. + /// The inferred operation type. + public static string InferOperationFromToolName(string toolName) + { + return toolName.ToLowerInvariant() switch + { + string s when s.Contains("read") || s.Contains("get") || s.Contains("list") || s.Contains("describe") => "read", + string s when s.Contains("create") || s.Contains("insert") => "create", + string s when s.Contains("update") || s.Contains("modify") => "update", + string s when s.Contains("delete") || s.Contains("remove") => "delete", + string s when s.Contains("execute") => "execute", + _ => "execute" + }; + } + + /// + /// Extracts metadata from a custom tool for telemetry purposes. + /// + /// The custom tool instance. + /// The service provider. + /// A tuple containing the entity name and database procedure name. + public static (string? entityName, string? dbProcedure) ExtractCustomToolMetadata(Core.DynamicCustomTool customTool, IServiceProvider serviceProvider) + { + try + { + // Access public properties instead of reflection + string? entityName = customTool.EntityName; + + if (entityName != null) + { + // 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); + } + } + } + + return (entityName, null); + } + catch (Exception ex) when (ex is InvalidOperationException || ex is ArgumentException) + { + // If configuration access fails due to invalid state or arguments, return null values + // This is expected during startup or configuration changes + return (null, 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/Mcp/McpTelemetryTests.cs b/src/Service.Tests/Mcp/McpTelemetryTests.cs new file mode 100644 index 0000000000..f86b338917 --- /dev/null +++ b/src/Service.Tests/Mcp/McpTelemetryTests.cs @@ -0,0 +1,229 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +#nullable enable + +using System; +using System.Collections.Generic; +using System.Diagnostics; +using System.Linq; +using Azure.DataApiBuilder.Core.Telemetry; +using Microsoft.VisualStudio.TestTools.UnitTesting; + +namespace Azure.DataApiBuilder.Service.Tests.Mcp +{ + /// + /// Tests for MCP telemetry functionality. + /// + [TestClass] + public class McpTelemetryTests + { + private static ActivityListener? _activityListener; + private static 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(); + } + + /// + /// Test that TrackMcpToolExecutionStarted sets the correct tags. + /// + [TestMethod] + public void TrackMcpToolExecutionStarted_SetsCorrectTags() + { + // Arrange + using Activity? activity = TelemetryTracesHelper.DABActivitySource.StartActivity("mcp.tool.execute"); + Assert.IsNotNull(activity, "Activity should be created"); + + // Act + activity.TrackMcpToolExecutionStarted( + toolName: "read_records", + entityName: "books", + operation: "read", + dbProcedure: null); + + activity.Stop(); + + // Assert + Activity? recordedActivity = _recordedActivities.FirstOrDefault(); + Assert.IsNotNull(recordedActivity, "Activity should be recorded"); + Assert.AreEqual("read_records", recordedActivity.GetTagItem("mcp.tool.name")); + Assert.AreEqual("books", recordedActivity.GetTagItem("dab.entity")); + Assert.AreEqual("read", recordedActivity.GetTagItem("dab.operation")); + } + + /// + /// Test that TrackMcpToolExecutionStarted sets db.procedure tag when provided. + /// + [TestMethod] + public void TrackMcpToolExecutionStarted_SetsDbProcedureTag_WhenProvided() + { + // Arrange + using Activity? activity = TelemetryTracesHelper.DABActivitySource.StartActivity("mcp.tool.execute"); + Assert.IsNotNull(activity, "Activity should be created"); + + // Act + activity.TrackMcpToolExecutionStarted( + toolName: "get_book", + entityName: "GetBook", + operation: "execute", + dbProcedure: "dbo.GetBookById"); + + activity.Stop(); + + // Assert + Activity? recordedActivity = _recordedActivities.FirstOrDefault(); + Assert.IsNotNull(recordedActivity, "Activity should be recorded"); + Assert.AreEqual("get_book", recordedActivity.GetTagItem("mcp.tool.name")); + Assert.AreEqual("GetBook", recordedActivity.GetTagItem("dab.entity")); + Assert.AreEqual("execute", recordedActivity.GetTagItem("dab.operation")); + Assert.AreEqual("dbo.GetBookById", recordedActivity.GetTagItem("db.procedure")); + } + + /// + /// Test that TrackMcpToolExecutionFinished sets status to OK. + /// + [TestMethod] + public void TrackMcpToolExecutionFinished_SetsStatusToOk() + { + // Arrange + using Activity? activity = TelemetryTracesHelper.DABActivitySource.StartActivity("mcp.tool.execute"); + Assert.IsNotNull(activity, "Activity should be created"); + + activity.TrackMcpToolExecutionStarted(toolName: "read_records"); + + // Act + activity.TrackMcpToolExecutionFinished(); + activity.Stop(); + + // Assert + Activity? recordedActivity = _recordedActivities.FirstOrDefault(); + Assert.IsNotNull(recordedActivity, "Activity should be recorded"); + Assert.AreEqual(ActivityStatusCode.Ok, recordedActivity.Status); + } + + /// + /// Test that TrackMcpToolExecutionFinishedWithException records exception and sets error status. + /// + [TestMethod] + public void TrackMcpToolExecutionFinishedWithException_RecordsExceptionAndSetsErrorStatus() + { + // Arrange + using Activity? activity = TelemetryTracesHelper.DABActivitySource.StartActivity("mcp.tool.execute"); + Assert.IsNotNull(activity, "Activity should be created"); + + activity.TrackMcpToolExecutionStarted(toolName: "read_records"); + + Exception testException = new InvalidOperationException("Test exception"); + + // Act + activity.TrackMcpToolExecutionFinishedWithException(testException, errorCode: "ExecutionFailed"); + activity.Stop(); + + // Assert + Activity? recordedActivity = _recordedActivities.FirstOrDefault(); + Assert.IsNotNull(recordedActivity, "Activity should be recorded"); + Assert.AreEqual(ActivityStatusCode.Error, recordedActivity.Status); + Assert.AreEqual("Test exception", recordedActivity.StatusDescription); + Assert.AreEqual("InvalidOperationException", recordedActivity.GetTagItem("error.type")); + Assert.AreEqual("Test exception", recordedActivity.GetTagItem("error.message")); + Assert.AreEqual("ExecutionFailed", recordedActivity.GetTagItem("error.code")); + + // Check that exception was recorded + ActivityEvent? exceptionEvent = recordedActivity.Events.FirstOrDefault(e => e.Name == "exception"); + Assert.IsNotNull(exceptionEvent, "Exception event should be recorded"); + } + + /// + /// Test that TrackMcpToolExecutionStarted handles null optional parameters gracefully. + /// + [TestMethod] + public void TrackMcpToolExecutionStarted_HandlesNullOptionalParameters() + { + // Arrange + using Activity? activity = TelemetryTracesHelper.DABActivitySource.StartActivity("mcp.tool.execute"); + Assert.IsNotNull(activity, "Activity should be created"); + + // Act - only provide tool name, all others are null + activity.TrackMcpToolExecutionStarted(toolName: "describe_entities"); + activity.Stop(); + + // Assert + Activity? recordedActivity = _recordedActivities.FirstOrDefault(); + Assert.IsNotNull(recordedActivity, "Activity should be recorded"); + Assert.AreEqual("describe_entities", recordedActivity.GetTagItem("mcp.tool.name")); + Assert.IsNull(recordedActivity.GetTagItem("dab.entity")); + Assert.IsNull(recordedActivity.GetTagItem("dab.operation")); + Assert.IsNull(recordedActivity.GetTagItem("db.procedure")); + } + + /// + /// Test that multiple tags can be set on the same activity. + /// + [TestMethod] + public void TrackMcpToolExecutionStarted_SupportsMultipleTags() + { + // Arrange + using Activity? activity = TelemetryTracesHelper.DABActivitySource.StartActivity("mcp.tool.execute"); + Assert.IsNotNull(activity, "Activity should be created"); + + // Act + activity.TrackMcpToolExecutionStarted( + toolName: "custom_tool", + entityName: "MyEntity", + operation: "execute", + dbProcedure: "schema.MyStoredProc"); + + activity.Stop(); + + // Assert + Activity? recordedActivity = _recordedActivities.FirstOrDefault(); + Assert.IsNotNull(recordedActivity, "Activity should be recorded"); + + // Verify all tags are present + Assert.AreEqual(4, recordedActivity.Tags.Count(t => + t.Key == "mcp.tool.name" || + t.Key == "dab.entity" || + t.Key == "dab.operation" || + t.Key == "db.procedure")); + + Assert.AreEqual("custom_tool", recordedActivity.GetTagItem("mcp.tool.name")); + Assert.AreEqual("MyEntity", recordedActivity.GetTagItem("dab.entity")); + Assert.AreEqual("execute", recordedActivity.GetTagItem("dab.operation")); + Assert.AreEqual("schema.MyStoredProc", recordedActivity.GetTagItem("db.procedure")); + } + } +}