diff --git a/core/Azure.Mcp.Core/tests/Azure.Mcp.Core.Tests/Areas/Server/Commands/ToolLoading/NamespaceToolLoaderTests.cs b/core/Azure.Mcp.Core/tests/Azure.Mcp.Core.Tests/Areas/Server/Commands/ToolLoading/NamespaceToolLoaderTests.cs index 7740464cb1..f09f4d3491 100644 --- a/core/Azure.Mcp.Core/tests/Azure.Mcp.Core.Tests/Areas/Server/Commands/ToolLoading/NamespaceToolLoaderTests.cs +++ b/core/Azure.Mcp.Core/tests/Azure.Mcp.Core.Tests/Areas/Server/Commands/ToolLoading/NamespaceToolLoaderTests.cs @@ -958,4 +958,115 @@ public async ValueTask DisposeAsync() GC.SuppressFinalize(this); } + + [Fact] + public void TryEmitSubscriptionTag_AddsTag_WhenSubscriptionParameterIsString() + { + // Arrange + const string expectedSubscription = "11111111-1111-1111-1111-111111111111"; + var parameters = new Dictionary(StringComparer.OrdinalIgnoreCase) + { + ["subscription"] = JsonSerializer.SerializeToElement(expectedSubscription), + ["vault"] = JsonSerializer.SerializeToElement("myvault"), + }; + + using var activity = new System.Diagnostics.Activity("test-activity"); + activity.Start(); + + // Act + NamespaceToolLoader.TryEmitSubscriptionTag(parameters, activity); + + // Assert + var tag = activity.GetTagItem(Microsoft.Mcp.Core.Services.Telemetry.AzureTagName.SubscriptionGuid); + Assert.Equal(expectedSubscription, tag); + } + + [Fact] + public void TryEmitSubscriptionTag_MatchesKeyCaseInsensitively() + { + // Arrange + const string expectedSubscription = "22222222-2222-2222-2222-222222222222"; + var parameters = new Dictionary(StringComparer.Ordinal) + { + ["Subscription"] = JsonSerializer.SerializeToElement(expectedSubscription), + }; + + using var activity = new System.Diagnostics.Activity("test-activity"); + activity.Start(); + + // Act + NamespaceToolLoader.TryEmitSubscriptionTag(parameters, activity); + + // Assert + var tag = activity.GetTagItem(Microsoft.Mcp.Core.Services.Telemetry.AzureTagName.SubscriptionGuid); + Assert.Equal(expectedSubscription, tag); + } + + [Fact] + public void TryEmitSubscriptionTag_DoesNothing_WhenNoSubscriptionParameter() + { + // Arrange + var parameters = new Dictionary + { + ["vault"] = JsonSerializer.SerializeToElement("myvault"), + }; + + using var activity = new System.Diagnostics.Activity("test-activity"); + activity.Start(); + + // Act + NamespaceToolLoader.TryEmitSubscriptionTag(parameters, activity); + + // Assert + Assert.Null(activity.GetTagItem(Microsoft.Mcp.Core.Services.Telemetry.AzureTagName.SubscriptionGuid)); + } + + [Fact] + public void TryEmitSubscriptionTag_EmitsEmptyString_MatchingMcpRuntimeBehavior() + { + // Arrange + // McpRuntime.CallToolHandler emits the AzSubscriptionGuid tag for any non-null string, + // including empty. This test guards against drift between the two emission sites. + var parameters = new Dictionary + { + ["subscription"] = JsonSerializer.SerializeToElement(string.Empty), + }; + + using var activity = new System.Diagnostics.Activity("test-activity"); + activity.Start(); + + // Act + NamespaceToolLoader.TryEmitSubscriptionTag(parameters, activity); + + // Assert + var tag = activity.GetTagItem(Microsoft.Mcp.Core.Services.Telemetry.AzureTagName.SubscriptionGuid); + Assert.Equal(string.Empty, tag); + } + + [Fact] + public void TryEmitSubscriptionTag_DoesNothing_WhenActivityIsNull() + { + // Arrange + var parameters = new Dictionary + { + ["subscription"] = JsonSerializer.SerializeToElement("abc"), + }; + + // Act & Assert (must not throw) + NamespaceToolLoader.TryEmitSubscriptionTag(parameters, activity: null); + } + + [Fact] + public void TryEmitSubscriptionTag_DoesNothing_WhenParametersAreNull() + { + // Arrange + using var activity = new System.Diagnostics.Activity("test-activity"); + activity.Start(); + + // Act + NamespaceToolLoader.TryEmitSubscriptionTag(parameters: null, activity); + + // Assert + Assert.Null(activity.GetTagItem(Microsoft.Mcp.Core.Services.Telemetry.AzureTagName.SubscriptionGuid)); + } } diff --git a/core/Microsoft.Mcp.Core/src/Areas/Server/Commands/ToolLoading/NamespaceToolLoader.cs b/core/Microsoft.Mcp.Core/src/Areas/Server/Commands/ToolLoading/NamespaceToolLoader.cs index 3a1d8c3452..d27cf39f76 100644 --- a/core/Microsoft.Mcp.Core/src/Areas/Server/Commands/ToolLoading/NamespaceToolLoader.cs +++ b/core/Microsoft.Mcp.Core/src/Areas/Server/Commands/ToolLoading/NamespaceToolLoader.cs @@ -13,6 +13,8 @@ using Microsoft.Mcp.Core.Helpers; using Microsoft.Mcp.Core.Models; using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Models.Option; +using Microsoft.Mcp.Core.Services.Telemetry; using ModelContextProtocol; using ModelContextProtocol.Protocol; @@ -439,6 +441,11 @@ private async Task InvokeChildToolAsync( .SetTag(TagName.ToolId, cmd.Id) .SetTag(TagName.IsServerCommandInvoked, true); + // In namespace mode the McpRuntime sees only the outer {intent, command, parameters} + // envelope, so it cannot emit AzSubscriptionGuid from the top-level arguments. Emit it + // here from the inner parameter dictionary so per-subscription telemetry is captured. + TryEmitSubscriptionTag(parameters, currentActivity); + var commandResponse = await cmd.ExecuteAsync(commandContext, commandOptions, cancellationToken); var jsonResponse = JsonSerializer.Serialize(commandResponse, ModelsJsonContext.Default.CommandResponse); var isError = commandResponse.Status < HttpStatusCode.OK || commandResponse.Status >= HttpStatusCode.Ambiguous; @@ -672,6 +679,39 @@ private static bool IsRawMcpToolInputOption(Option option) string.Equals(NameNormalization.NormalizeOptionName(alias), RawMcpToolInputOptionName, StringComparison.OrdinalIgnoreCase)); } + internal static void TryEmitSubscriptionTag(IDictionary? parameters, Activity? activity) + { + if (activity == null || parameters == null || parameters.Count == 0) + { + return; + } + + var normalizedSubscriptionName = NameNormalization.NormalizeOptionName(OptionDefinitions.Common.Subscription.Name); + foreach (var kvp in parameters) + { + if (!string.Equals(kvp.Key, normalizedSubscriptionName, StringComparison.OrdinalIgnoreCase)) + { + continue; + } + + if (kvp.Value.ValueKind != JsonValueKind.String) + { + return; + } + + // Match McpRuntime.CallToolHandler: emit the tag for any non-null string, + // including empty values, so telemetry is consistent between direct-tool and + // namespace-mode invocations. + var subscription = kvp.Value.GetString(); + if (subscription != null) + { + activity.AddTag(AzureTagName.SubscriptionGuid, subscription); + } + + return; + } + } + internal static Dictionary GetParametersFromArgs(IDictionary? args) { if (args == null)