Don't expose Activity directly in CommandContext

core/Microsoft.Mcp.Core/src/Areas/Server/Commands/ToolLoading/CommandFactoryToolLoader.cs

@@ -215,11 +215,7 @@ public override async ValueTask<CallToolResult> CallToolHandler(RequestContext<C
 
         _logger.LogTrace("Invoking '{Tool}'.", realCommand.Name);
 
-        if (commandContext.Activity != null)
-        {
-            var serviceArea = _commandFactory.GetServiceArea(toolName);
-            commandContext.Activity.SetTag(TagName.ToolArea, serviceArea);
-        }
+        commandContext.SetTelemetryTag(TagName.ToolArea, _commandFactory.GetServiceArea(toolName));
 
         try
         {

core/Microsoft.Mcp.Core/src/Commands/BaseCommand`2.cs

@@ -112,9 +112,9 @@ async Task<CommandResponse> IBaseCommand.ExecuteAsync(CommandContext context, Pa
 
     protected virtual void HandleException(CommandContext context, Exception ex)
     {
-        context.Activity?.SetStatus(ActivityStatusCode.Error)
-            ?.SetTag(TagName.ExceptionType, ex.GetType().ToString())
-            ?.SetTag(TagName.ExceptionStackTrace, ex.StackTrace);
+        context.SetTelemetryStatus(ActivityStatusCode.Error)
+            .SetTelemetryTag(TagName.ExceptionType, ex.GetType().ToString())
+            .SetTelemetryTag(TagName.ExceptionStackTrace, ex.StackTrace);
 
         var response = context.Response;
 
@@ -134,7 +134,7 @@ protected virtual void HandleException(CommandContext context, Exception ex)
 
             // Include the command validation exception message as it should be safe. Requires custom validators to
             // exclude any sensitive information from their error messages.
-            context.Activity?.SetTag(TagName.ExceptionMessage, response.Message);
+            context.SetTelemetryTag(TagName.ExceptionMessage, response.Message);
             response.Results = null;
             return;
         }
@@ -160,7 +160,7 @@ protected virtual void HandleException(CommandContext context, Exception ex)
             exceptionDetails.Add("CorrelationId", msalClientException.CorrelationId);
         }
 
-        context.Activity?.SetTag(TagName.ExceptionMessage, exceptionDetails);
+        context.SetTelemetryTag(TagName.ExceptionMessage, exceptionDetails);
 
         var result = new ExceptionResult(
             Message: ex.Message ?? string.Empty,

core/Microsoft.Mcp.Core/src/Models/Command/CommandContext.cs

@@ -21,7 +21,7 @@ public class CommandContext
     /// <summary>
     /// Current telemetry context if there is one available.
     /// </summary>
-    public Activity? Activity { get; }
+    private Activity? _activity { get; }
 
     /// <summary>
     /// The MCP server handling the current tool call. Used by commands that need to send
@@ -49,11 +49,47 @@ public class CommandContext
     /// </summary>
     public CommandContext(Activity? activity = default)
     {
-        Activity = activity;
+        _activity = activity;
         Response = new CommandResponse
         {
             Status = HttpStatusCode.OK,
             Message = "Success"
         };
     }
+
+    /// <summary>
+    /// Adds a telemetry tag to the current activity if one is available. This is a no-op if there is no current activity.
+    /// </summary>
+    /// <para>
+    /// Equivalent to <see cref="Activity.AddTag(string, object?)" />.
+    /// </para>
+    /// <param name="key">The telemetry tag name.</param>
+    /// <param name="value">The telemetry tag value to add.</param>
+    /// <returns>The CommandContext</returns>
+    public CommandContext AddTelemetryTag(string key, object? value)
+    {
+        _activity?.AddTag(key, value);
+        return this;
+    }
+
+    /// <summary>
+    /// Set a telemetry tag to the current activity if one is available. This is a no-op if there is no current activity.
+    /// </summary>
+    /// <para>
+    /// Equivalent to <see cref="Activity.SetTag(string, object?)" />.
+    /// </para>
+    /// <param name="key">The telemetry tag name.</param>
+    /// <param name="value">The telemetry tag value to set.</param>
+    /// <returns>The CommandContext</returns>
+    public CommandContext SetTelemetryTag(string key, object? value)
+    {
+        _activity?.SetTag(key, value);
+        return this;
+    }
+
+    internal CommandContext SetTelemetryStatus(ActivityStatusCode status)
+    {
+        _activity?.SetStatus(status);
+        return this;
+    }
 }

tools/Azure.Mcp.Tools.Advisor/src/Commands/Recommendation/RecommendationApplyCommand.cs

@@ -53,7 +53,7 @@ public override Task<CommandResponse> ExecuteAsync(CommandContext context, Recom
 
             context.Response.Results = ResponseResult.Create([recommendationApplyRules], AdvisorJsonContext.Default.ListString);
 
-            context.Activity?.AddTag("RecommendationRules_Resource", options.Resource);
+            context.AddTelemetryTag("RecommendationRules_Resource", options.Resource);
         }
         catch (Exception ex)
         {

tools/Azure.Mcp.Tools.AppService/src/Commands/Database/DatabaseAddCommand.cs

@@ -37,7 +37,7 @@ public override async Task<CommandResponse> ExecuteAsync(CommandContext context,
     {
         try
         {
-            context.Activity?.AddTag("subscription", options.Subscription);
+            context.AddTelemetryTag("subscription", options.Subscription);
 
             var connectionInfo = await _appServiceService.AddDatabaseAsync(
                 options.App,

tools/Azure.Mcp.Tools.AppService/src/Commands/Webapp/Deployment/DeploymentGetCommand.cs

@@ -39,7 +39,7 @@ public override async Task<CommandResponse> ExecuteAsync(CommandContext context,
     {
         try
         {
-            context.Activity?.AddTag("subscription", options.Subscription);
+            context.AddTelemetryTag("subscription", options.Subscription);
 
             var deployments = await _appServiceService.GetDeploymentsAsync(
                 options.Subscription!,

tools/Azure.Mcp.Tools.AppService/src/Commands/Webapp/Diagnostic/DetectorDiagnoseCommand.cs

@@ -51,7 +51,7 @@ public override async Task<CommandResponse> ExecuteAsync(CommandContext context,
     {
         try
         {
-            context.Activity?.AddTag("subscription", options.Subscription);
+            context.AddTelemetryTag("subscription", options.Subscription);
 
             var diagnoses = await _appServiceService.DiagnoseDetectorAsync(
                 options.Subscription!,

tools/Azure.Mcp.Tools.AppService/src/Commands/Webapp/Diagnostic/DetectorListCommand.cs

@@ -38,7 +38,7 @@ public override async Task<CommandResponse> ExecuteAsync(CommandContext context,
     {
         try
         {
-            context.Activity?.AddTag("subscription", options.Subscription);
+            context.AddTelemetryTag("subscription", options.Subscription);
 
             var detectors = await _appServiceService.ListDetectorsAsync(
                 options.Subscription!,

tools/Azure.Mcp.Tools.AppService/src/Commands/Webapp/Settings/AppSettingsGetCommand.cs

@@ -35,7 +35,7 @@ public override async Task<CommandResponse> ExecuteAsync(CommandContext context,
     {
         try
         {
-            context.Activity?.AddTag("subscription", options.Subscription);
+            context.AddTelemetryTag("subscription", options.Subscription);
 
             var appSettings = await _appServiceService.GetAppSettingsAsync(
                 options.Subscription!,

tools/Azure.Mcp.Tools.AppService/src/Commands/Webapp/Settings/AppSettingsUpdateCommand.cs

@@ -81,7 +81,7 @@ public override async Task<CommandResponse> ExecuteAsync(CommandContext context,
     {
         try
         {
-            context.Activity?.AddTag("subscription", options.Subscription);
+            context.AddTelemetryTag("subscription", options.Subscription);
 
             var updateResult = await _appServiceService.UpdateAppSettingsAsync(
                 options.Subscription!,

tools/Azure.Mcp.Tools.AppService/src/Commands/Webapp/WebappChangeStateCommand.cs

@@ -86,7 +86,7 @@ public override async Task<CommandResponse> ExecuteAsync(CommandContext context,
     {
         try
         {
-            context.Activity?.AddTag("subscription", options.Subscription);
+            context.AddTelemetryTag("subscription", options.Subscription);
 
             var stateChange = await appServiceService.ChangeWebAppStateAsync(
                 options.Subscription!,

tools/Azure.Mcp.Tools.AppService/src/Commands/Webapp/WebappGetCommand.cs

@@ -49,7 +49,7 @@ public override async Task<CommandResponse> ExecuteAsync(CommandContext context,
     {
         try
         {
-            context.Activity?.AddTag("subscription", options.Subscription);
+            context.AddTelemetryTag("subscription", options.Subscription);
 
             var webapps = await _appServiceService.GetWebAppsAsync(
                 options.Subscription!,

tools/Azure.Mcp.Tools.AzureBackup/src/Commands/Backup/BackupStatusCommand.cs

@@ -37,8 +37,8 @@ public sealed class BackupStatusCommand(ILogger<BackupStatusCommand> logger, IAz
 
     public override async Task<CommandResponse> ExecuteAsync(CommandContext context, BackupStatusOptions options, CancellationToken cancellationToken)
     {
-        AzureBackupTelemetryTags.AddSubscriptionTag(context.Activity, options.Subscription);
-        context.Activity?.AddTag(AzureBackupTelemetryTags.OperationScope, "status-check");
+        AzureBackupTelemetryTags.AddSubscriptionTag(context, options.Subscription);
+        context.AddTelemetryTag(AzureBackupTelemetryTags.OperationScope, "status-check");
 
         try
         {

tools/Azure.Mcp.Tools.AzureBackup/src/Commands/DisasterRecovery/DisasterRecoveryEnableCrrCommand.cs

@@ -31,8 +31,8 @@ public sealed class DisasterRecoveryEnableCrrCommand(ILogger<DisasterRecoveryEna
 
     public override async Task<CommandResponse> ExecuteAsync(CommandContext context, BaseAzureBackupOptions options, CancellationToken cancellationToken)
     {
-        AzureBackupTelemetryTags.AddSubscriptionTag(context.Activity, options.Subscription);
-        AzureBackupTelemetryTags.AddVaultTags(context.Activity, options.VaultType);
+        AzureBackupTelemetryTags.AddSubscriptionTag(context, options.Subscription);
+        AzureBackupTelemetryTags.AddVaultTags(context, options.VaultType);
 
         try
         {

tools/Azure.Mcp.Tools.AzureBackup/src/Commands/Governance/GovernanceFindUnprotectedCommand.cs

@@ -56,8 +56,8 @@ public sealed class GovernanceFindUnprotectedCommand(ILogger<GovernanceFindUnpro
 
     public override async Task<CommandResponse> ExecuteAsync(CommandContext context, GovernanceFindUnprotectedOptions options, CancellationToken cancellationToken)
     {
-        AzureBackupTelemetryTags.AddSubscriptionTag(context.Activity, options.Subscription);
-        context.Activity?.AddTag(AzureBackupTelemetryTags.OperationScope, "scan");
+        AzureBackupTelemetryTags.AddSubscriptionTag(context, options.Subscription);
+        context.AddTelemetryTag(AzureBackupTelemetryTags.OperationScope, "scan");
 
         try
         {

tools/Azure.Mcp.Tools.AzureBackup/src/Commands/Governance/GovernanceImmutabilityCommand.cs

@@ -47,8 +47,8 @@ public override void ValidateOptions(GovernanceImmutabilityOptions options, Vali
 
     public override async Task<CommandResponse> ExecuteAsync(CommandContext context, GovernanceImmutabilityOptions options, CancellationToken cancellationToken)
     {
-        AzureBackupTelemetryTags.AddSubscriptionTag(context.Activity, options.Subscription);
-        AzureBackupTelemetryTags.AddVaultTags(context.Activity, options.VaultType);
+        AzureBackupTelemetryTags.AddSubscriptionTag(context, options.Subscription);
+        AzureBackupTelemetryTags.AddVaultTags(context, options.VaultType);
 
         try
         {

tools/Azure.Mcp.Tools.AzureBackup/src/Commands/Governance/GovernanceSoftDeleteCommand.cs

@@ -54,8 +54,8 @@ public override void ValidateOptions(GovernanceSoftDeleteOptions options, Valida
 
     public override async Task<CommandResponse> ExecuteAsync(CommandContext context, GovernanceSoftDeleteOptions options, CancellationToken cancellationToken)
     {
-        AzureBackupTelemetryTags.AddSubscriptionTag(context.Activity, options.Subscription);
-        AzureBackupTelemetryTags.AddVaultTags(context.Activity, options.VaultType);
+        AzureBackupTelemetryTags.AddSubscriptionTag(context, options.Subscription);
+        AzureBackupTelemetryTags.AddVaultTags(context, options.VaultType);
 
         try
         {

tools/Azure.Mcp.Tools.AzureBackup/src/Commands/Job/JobGetCommand.cs

@@ -39,9 +39,9 @@ public sealed class JobGetCommand(ILogger<JobGetCommand> logger, IAzureBackupSer
 
     public override async Task<CommandResponse> ExecuteAsync(CommandContext context, JobGetOptions options, CancellationToken cancellationToken)
     {
-        AzureBackupTelemetryTags.AddSubscriptionTag(context.Activity, options.Subscription);
-        AzureBackupTelemetryTags.AddVaultTags(context.Activity, options.VaultType);
-        context.Activity?.AddTag(AzureBackupTelemetryTags.OperationScope, string.IsNullOrEmpty(options.Job) ? "list" : "single");
+        AzureBackupTelemetryTags.AddSubscriptionTag(context, options.Subscription);
+        AzureBackupTelemetryTags.AddVaultTags(context, options.VaultType);
+        context.AddTelemetryTag(AzureBackupTelemetryTags.OperationScope, string.IsNullOrEmpty(options.Job) ? "list" : "single");
 
         try
         {

tools/Azure.Mcp.Tools.AzureBackup/src/Commands/Policy/PolicyCreateCommand.cs

@@ -31,8 +31,8 @@ public sealed class PolicyCreateCommand(ILogger<PolicyCreateCommand> logger, IAz
 
     public override async Task<CommandResponse> ExecuteAsync(CommandContext context, PolicyCreateOptions options, CancellationToken cancellationToken)
     {
-        AzureBackupTelemetryTags.AddSubscriptionTag(context.Activity, options.Subscription);
-        AzureBackupTelemetryTags.AddVaultAndWorkloadTags(context.Activity, options.VaultType, options.WorkloadType);
+        AzureBackupTelemetryTags.AddSubscriptionTag(context, options.Subscription);
+        AzureBackupTelemetryTags.AddVaultAndWorkloadTags(context, options.VaultType, options.WorkloadType);
 
         var validation = Services.Policy.PolicyCreateValidator.Validate(options);
         if (!validation.IsValid)

tools/Azure.Mcp.Tools.AzureBackup/src/Commands/Policy/PolicyGetCommand.cs

@@ -39,9 +39,9 @@ public sealed class PolicyGetCommand(ILogger<PolicyGetCommand> logger, IAzureBac
 
     public override async Task<CommandResponse> ExecuteAsync(CommandContext context, PolicyGetOptions options, CancellationToken cancellationToken)
     {
-        AzureBackupTelemetryTags.AddSubscriptionTag(context.Activity, options.Subscription);
-        AzureBackupTelemetryTags.AddVaultTags(context.Activity, options.VaultType);
-        context.Activity?.AddTag(AzureBackupTelemetryTags.OperationScope, string.IsNullOrEmpty(options.Policy) ? "list" : "single");
+        AzureBackupTelemetryTags.AddSubscriptionTag(context, options.Subscription);
+        AzureBackupTelemetryTags.AddVaultTags(context, options.VaultType);
+        context.AddTelemetryTag(AzureBackupTelemetryTags.OperationScope, string.IsNullOrEmpty(options.Policy) ? "list" : "single");
 
         try
         {

tools/Azure.Mcp.Tools.AzureBackup/src/Commands/Policy/PolicyUpdateCommand.cs

@@ -31,8 +31,8 @@ public sealed class PolicyUpdateCommand(ILogger<PolicyUpdateCommand> logger, IAz
 
     public override async Task<CommandResponse> ExecuteAsync(CommandContext context, PolicyUpdateOptions options, CancellationToken cancellationToken)
     {
-        AzureBackupTelemetryTags.AddSubscriptionTag(context.Activity, options.Subscription);
-        AzureBackupTelemetryTags.AddVaultTags(context.Activity, options.VaultType);
+        AzureBackupTelemetryTags.AddSubscriptionTag(context, options.Subscription);
+        AzureBackupTelemetryTags.AddVaultTags(context, options.VaultType);
 
         try
         {

tools/Azure.Mcp.Tools.AzureBackup/src/Commands/ProtectableItem/ProtectableItemListCommand.cs

@@ -53,8 +53,8 @@ public override void ValidateOptions(ProtectableItemListOptions options, Validat
 
     public override async Task<CommandResponse> ExecuteAsync(CommandContext context, ProtectableItemListOptions options, CancellationToken cancellationToken)
     {
-        AzureBackupTelemetryTags.AddSubscriptionTag(context.Activity, options.Subscription);
-        AzureBackupTelemetryTags.AddVaultAndWorkloadTags(context.Activity, options.VaultType ?? "rsv", options.WorkloadType);
+        AzureBackupTelemetryTags.AddSubscriptionTag(context, options.Subscription);
+        AzureBackupTelemetryTags.AddVaultAndWorkloadTags(context, options.VaultType ?? "rsv", options.WorkloadType);
 
         try
         {

tools/Azure.Mcp.Tools.AzureBackup/src/Commands/ProtectedItem/ProtectedItemGetCommand.cs

@@ -41,9 +41,9 @@ public sealed class ProtectedItemGetCommand(ILogger<ProtectedItemGetCommand> log
 
     public override async Task<CommandResponse> ExecuteAsync(CommandContext context, BaseProtectedItemOptions options, CancellationToken cancellationToken)
     {
-        AzureBackupTelemetryTags.AddSubscriptionTag(context.Activity, options.Subscription);
-        AzureBackupTelemetryTags.AddVaultTags(context.Activity, options.VaultType);
-        context.Activity?.AddTag(AzureBackupTelemetryTags.OperationScope, string.IsNullOrEmpty(options.ProtectedItem) ? "list" : "single");
+        AzureBackupTelemetryTags.AddSubscriptionTag(context, options.Subscription);
+        AzureBackupTelemetryTags.AddVaultTags(context, options.VaultType);
+        context.AddTelemetryTag(AzureBackupTelemetryTags.OperationScope, string.IsNullOrEmpty(options.ProtectedItem) ? "list" : "single");
 
         try
         {

tools/Azure.Mcp.Tools.AzureBackup/src/Commands/ProtectedItem/ProtectedItemProtectCommand.cs

@@ -89,9 +89,9 @@ public override void ValidateOptions(ProtectedItemProtectOptions options, Valida
 
     public override async Task<CommandResponse> ExecuteAsync(CommandContext context, ProtectedItemProtectOptions options, CancellationToken cancellationToken)
     {
-        AzureBackupTelemetryTags.AddSubscriptionTag(context.Activity, options.Subscription);
-        AzureBackupTelemetryTags.AddVaultTags(context.Activity, options.VaultType);
-        context.Activity?.AddTag(AzureBackupTelemetryTags.DatasourceType, AzureBackupTelemetryTags.NormalizeWorkloadType(options.DatasourceType));
+        AzureBackupTelemetryTags.AddSubscriptionTag(context, options.Subscription);
+        AzureBackupTelemetryTags.AddVaultTags(context, options.VaultType);
+        context.AddTelemetryTag(AzureBackupTelemetryTags.DatasourceType, AzureBackupTelemetryTags.NormalizeWorkloadType(options.DatasourceType));
 
         try
         {

tools/Azure.Mcp.Tools.AzureBackup/src/Commands/ProtectedItem/ProtectedItemUndeleteCommand.cs

@@ -38,8 +38,8 @@ public sealed class ProtectedItemUndeleteCommand(ILogger<ProtectedItemUndeleteCo
 
     public override async Task<CommandResponse> ExecuteAsync(CommandContext context, ProtectedItemUndeleteOptions options, CancellationToken cancellationToken)
     {
-        AzureBackupTelemetryTags.AddSubscriptionTag(context.Activity, options.Subscription);
-        AzureBackupTelemetryTags.AddVaultTags(context.Activity, options.VaultType);
+        AzureBackupTelemetryTags.AddSubscriptionTag(context, options.Subscription);
+        AzureBackupTelemetryTags.AddVaultTags(context, options.VaultType);
 
         try
         {

tools/Azure.Mcp.Tools.AzureBackup/src/Commands/RecoveryPoint/RecoveryPointGetCommand.cs

@@ -39,9 +39,9 @@ public sealed class RecoveryPointGetCommand(ILogger<RecoveryPointGetCommand> log
 
     public override async Task<CommandResponse> ExecuteAsync(CommandContext context, RecoveryPointGetOptions options, CancellationToken cancellationToken)
     {
-        AzureBackupTelemetryTags.AddSubscriptionTag(context.Activity, options.Subscription);
-        AzureBackupTelemetryTags.AddVaultTags(context.Activity, options.VaultType);
-        context.Activity?.AddTag(AzureBackupTelemetryTags.OperationScope, string.IsNullOrEmpty(options.RecoveryPoint) ? "list" : "single");
+        AzureBackupTelemetryTags.AddSubscriptionTag(context, options.Subscription);
+        AzureBackupTelemetryTags.AddVaultTags(context, options.VaultType);
+        context.AddTelemetryTag(AzureBackupTelemetryTags.OperationScope, string.IsNullOrEmpty(options.RecoveryPoint) ? "list" : "single");
 
         try
         {

tools/Azure.Mcp.Tools.AzureBackup/src/Commands/Security/SecurityConfigureEncryptionCommand.cs

@@ -82,8 +82,8 @@ public override void ValidateOptions(SecurityConfigureEncryptionOptions options,
 
     public override async Task<CommandResponse> ExecuteAsync(CommandContext context, SecurityConfigureEncryptionOptions options, CancellationToken cancellationToken)
     {
-        AzureBackupTelemetryTags.AddSubscriptionTag(context.Activity, options.Subscription);
-        AzureBackupTelemetryTags.AddVaultTags(context.Activity, options.VaultType);
+        AzureBackupTelemetryTags.AddSubscriptionTag(context, options.Subscription);
+        AzureBackupTelemetryTags.AddVaultTags(context, options.VaultType);
         _lastVaultType = options.VaultType;
 
         try

tools/Azure.Mcp.Tools.AzureBackup/src/Commands/Security/SecurityConfigureMuaCommand.cs

@@ -48,8 +48,8 @@ public override void ValidateOptions(SecurityConfigureMuaOptions options, Valida
 
     public override async Task<CommandResponse> ExecuteAsync(CommandContext context, SecurityConfigureMuaOptions options, CancellationToken cancellationToken)
     {
-        AzureBackupTelemetryTags.AddSubscriptionTag(context.Activity, options.Subscription);
-        AzureBackupTelemetryTags.AddVaultTags(context.Activity, options.VaultType);
+        AzureBackupTelemetryTags.AddSubscriptionTag(context, options.Subscription);
+        AzureBackupTelemetryTags.AddVaultTags(context, options.VaultType);
 
         try
         {

tools/Azure.Mcp.Tools.AzureBackup/src/Commands/Vault/VaultCreateCommand.cs

@@ -58,8 +58,8 @@ public override void ValidateOptions(VaultCreateOptions options, ValidationResul
 
     public override async Task<CommandResponse> ExecuteAsync(CommandContext context, VaultCreateOptions options, CancellationToken cancellationToken)
     {
-        AzureBackupTelemetryTags.AddSubscriptionTag(context.Activity, options.Subscription);
-        AzureBackupTelemetryTags.AddVaultTags(context.Activity, options.VaultType);
+        AzureBackupTelemetryTags.AddSubscriptionTag(context, options.Subscription);
+        AzureBackupTelemetryTags.AddVaultTags(context, options.VaultType);
 
         try
         {