From b889206fc6517b43044d9212c1eefce403b3d5f0 Mon Sep 17 00:00:00 2001 From: alzimmermsft <48699787+alzimmermsft@users.noreply.github.com> Date: Thu, 19 Mar 2026 17:40:35 -0400 Subject: [PATCH] Start a host when calling commands via CLI to capture telemetry --- .../src/Commands/CommandFactory.cs | 30 ++++++ .../src/Extensions/OpenTelemetryExtensions.cs | 2 +- servers/Azure.Mcp.Server/src/Program.cs | 97 ++++++++++++++++++- 3 files changed, 124 insertions(+), 5 deletions(-) diff --git a/core/Microsoft.Mcp.Core/src/Commands/CommandFactory.cs b/core/Microsoft.Mcp.Core/src/Commands/CommandFactory.cs index 670e2d733a..495be79c53 100644 --- a/core/Microsoft.Mcp.Core/src/Commands/CommandFactory.cs +++ b/core/Microsoft.Mcp.Core/src/Commands/CommandFactory.cs @@ -2,6 +2,7 @@ // Licensed under the MIT License. using System.CommandLine.Help; +using System.CommandLine.Parsing; using System.Diagnostics; using System.Net; using System.Reflection; @@ -191,6 +192,7 @@ private void ConfigureCommandHandler(Command command, IBaseCommand implementatio using var activity = _telemetryService.StartActivity(ActivityName.CommandExecuted); activity?.SetTag(TagName.ToolId, implementation.Id); + InjectToolAreaAndName(activity, parseResult); var cmdContext = new CommandContext(_serviceProvider, activity); var startTime = DateTime.UtcNow; try @@ -327,6 +329,34 @@ private void CustomizeHelpOption(Command command) } } + /// + /// Injects tool area and name tags into the activity based on the command being executed. The full command name + /// is parsed to determine the tool area and name, which are then added as tags to the activity. + /// + /// The activity to inject tool area and name tags into. + /// The parsing result for the command. + private void InjectToolAreaAndName(Activity? activity, ParseResult parseResult) + { + if (activity == null) + { + return; + } + + var commandResult = parseResult.CommandResult; + var fullCommandName = commandResult.Command.Name; + while (commandResult.Parent is not null + && commandResult?.Parent is CommandResult + && (commandResult.Parent as CommandResult)?.Command.Name != RootCommand.Name) + { + commandResult = (commandResult.Parent as CommandResult)!; + fullCommandName = commandResult.Command.Name + "_" + fullCommandName; + } + + var index = fullCommandName.IndexOf('_'); + activity.SetTag(TagName.ToolArea, index == -1 ? fullCommandName : fullCommandName.Substring(0, index)) + .SetTag(TagName.ToolName, fullCommandName); + } + /// /// Creates a command dictionary. Each sibling and child of the root node is created without using its name as a prefix. /// diff --git a/core/Microsoft.Mcp.Core/src/Extensions/OpenTelemetryExtensions.cs b/core/Microsoft.Mcp.Core/src/Extensions/OpenTelemetryExtensions.cs index dd0f7aa524..339ae90802 100644 --- a/core/Microsoft.Mcp.Core/src/Extensions/OpenTelemetryExtensions.cs +++ b/core/Microsoft.Mcp.Core/src/Extensions/OpenTelemetryExtensions.cs @@ -103,7 +103,7 @@ private static void EnableAzureMonitor(this IServiceCollection services) #endif var enableOtlp = Environment.GetEnvironmentVariable("AZURE_MCP_ENABLE_OTLP_EXPORTER"); - if (!string.IsNullOrEmpty(enableOtlp) && bool.TryParse(enableOtlp, out var shouldEnable) && shouldEnable) + if (true) { otelBuilder.WithTracing(tracing => tracing.AddOtlpExporter()) .WithMetrics(metrics => metrics.AddOtlpExporter()) diff --git a/servers/Azure.Mcp.Server/src/Program.cs b/servers/Azure.Mcp.Server/src/Program.cs index f75961ef4e..f81902249d 100644 --- a/servers/Azure.Mcp.Server/src/Program.cs +++ b/servers/Azure.Mcp.Server/src/Program.cs @@ -5,6 +5,7 @@ using Azure.Mcp.Core.Commands; using Azure.Mcp.Core.Helpers; using Azure.Mcp.Core.Services.Azure; +using Azure.Mcp.Core.Services.Azure.Authentication; using Azure.Mcp.Core.Services.Azure.ResourceGroup; using Azure.Mcp.Core.Services.Azure.Subscription; using Azure.Mcp.Core.Services.Azure.Tenant; @@ -12,6 +13,7 @@ using Azure.Mcp.Core.Services.ProcessExecution; using Azure.Mcp.Core.Services.Time; using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Hosting; using Microsoft.Extensions.Logging; using Microsoft.Extensions.Options; using Microsoft.Mcp.Core.Areas; @@ -50,7 +52,7 @@ private static async Task Main(string[] args) ServiceCollection services = new(); - ConfigureServices(services); + ConfigureServices(services, [.. Areas.Where(a => a is ServerSetup)]); services.AddLogging(builder => { @@ -64,7 +66,40 @@ private static async Task Main(string[] args) var commandFactory = serviceProvider.GetRequiredService(); var rootCommand = commandFactory.RootCommand; var parseResult = rootCommand.Parse(args); - var status = await parseResult.InvokeAsync(); + int status = 0; + + if (parseResult.Errors.Count > 0) + { + // Command wasn't one of the registered ServerSetup commands, so bind up a Host of all the services + // to run the command. + var builder = Host.CreateApplicationBuilder(); + builder.Logging.ClearProviders(); + builder.Logging.AddEventSourceLogger(); + builder.Services.AddSingleIdentityTokenCredentialProvider(); + ConfigureServices(builder.Services); + builder.Services.AddAzureMcpServer(new ServiceStartOptions() + { + Transport = TransportTypes.StdIo + }); + + using var host = builder.Build(); + + await InitializeServicesAsync(host.Services); + await host.StartAsync(); + + commandFactory = host.Services.GetRequiredService(); + rootCommand = commandFactory.RootCommand; + parseResult = rootCommand.Parse(args); + + status = await parseResult.InvokeAsync(); + + await host.StopAsync(); + await host.WaitForShutdownAsync(); + } + else + { + status = await parseResult.InvokeAsync(); + } if (status == 0) { @@ -214,7 +249,60 @@ private static void WriteResponse(CommandResponse response) /// /// /// A service collection. - internal static void ConfigureServices(IServiceCollection services) + internal static void ConfigureServices(IServiceCollection services) => ConfigureServices(services, Areas); + + /// + /// + /// Configures services for dependency injection. + /// + /// + /// WARNING: This method is being used for TWO DEPENDENCY INJECTION CONTAINERS: + /// + /// + /// + /// 's command picking: The container used to populate instances of + /// and selected by + /// based on the command line input. This container is a local variable in + /// , and it is not tied to + /// Microsoft.Extensions.Hosting.IHostBuilder (stdio) nor any + /// Microsoft.AspNetCore.Hosting.IWebHostBuilder (http). + /// + /// + /// 's execution: The container is created by some + /// dynamically created Microsoft.Extensions.Hosting.IHostBuilder (stdio) or + /// Microsoft.AspNetCore.Hosting.IWebHostBuilder (http). While the + /// instance of + /// is created by the first container, this second container it creates and runs is + /// built separately during . Thus, this + /// container is built and this method is called sometime + /// during that method execution. + /// + /// + /// + /// DUE TO THIS DUAL USAGE, PLEASE BE VERY CAREFUL WHEN MODIFYING THIS METHOD. This + /// method may have some expectations, but it and all methods it calls must be safe for + /// both the stdio and http transport modes. + /// + /// + /// For example, most instances take an indirect dependency + /// on or , both of which have + /// transport-specific implementations. This method can add the stdio-specific + /// implementation to allow the first container (used for command picking) to work, + /// but such transport-specific registrations must be overridden within + /// with the appropriate + /// transport-specific implementation based on command line arguments. + /// + /// + /// This large doc comment is copy/pasta in each Program.cs file of this repo, so if + /// you're reading this, please keep them in sync and/or add specific warnings per + /// project if needed. Below is the list of known differences: + /// + /// + /// No differences. This is also copy/pasta as a placeholder for this project. + /// + /// + /// A service collection. + private static void ConfigureServices(IServiceCollection services, IAreaSetup[] areas) { var thisAssembly = typeof(Program).Assembly; @@ -236,7 +324,7 @@ internal static void ConfigureServices(IServiceCollection services) services.AddAzureTenantService(); services.AddSingleUserCliCacheService(); - foreach (var area in Areas) + foreach (var area in areas) { services.AddSingleton(area); area.ConfigureServices(services); @@ -254,6 +342,7 @@ internal static void ConfigureServices(IServiceCollection services) ActivatorUtilities.CreateInstance(sp, thisAssembly, $"allowed-plugin-file-references.json")); } + internal static async Task InitializeServicesAsync(IServiceProvider serviceProvider) { ServiceStartOptions? options = serviceProvider.GetService>()?.Value;