Changes to learning flow

core/Azure.Mcp.Core/src/Areas/Server/Commands/Runtime/McpRuntime.cs

@@ -98,6 +98,7 @@ public async ValueTask<CallToolResult> CallToolHandler(RequestContext<CallToolRe
         try
         {
             callTool = await _toolLoader.CallToolHandler(request!, cancellationToken);
+            activity?.SetTag(TagName.IsServerCommandInvoked, request.Items.TryGetValue(TagName.IsServerCommandInvoked, out var isInvoked) && isInvoked is true);
 
             var isSuccessful = !callTool.IsError.HasValue || !callTool.IsError.Value;
             if (isSuccessful)

core/Azure.Mcp.Core/src/Areas/Server/Commands/ToolLoading/BaseToolLoader.cs

@@ -1,7 +1,6 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-using System.Text.Json;
 using Microsoft.Extensions.Logging;
 using ModelContextProtocol.Client;
 using ModelContextProtocol.Protocol;

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

@@ -8,7 +8,6 @@
 using Azure.Mcp.Core.Commands;
 using Azure.Mcp.Core.Helpers;
 using Azure.Mcp.Core.Models.Elicitation;
-using Azure.Mcp.Core.Services.Telemetry;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
 using ModelContextProtocol.Protocol;

core/Azure.Mcp.Core/src/Areas/Server/Commands/ToolLoading/NamespaceToolLoader.cs

@@ -10,7 +10,6 @@
 using Azure.Mcp.Core.Commands;
 using Azure.Mcp.Core.Helpers;
 using Azure.Mcp.Core.Models.Elicitation;
-using Azure.Mcp.Core.Services.Telemetry;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
 using ModelContextProtocol;
@@ -175,7 +174,7 @@ public override async ValueTask<CallToolResult> CallToolHandler(RequestContext<C
         {
             var activity = Activity.Current;
 
-            if (learn && string.IsNullOrEmpty(command))
+            if (learn)
             {
                 return await InvokeToolLearn(request, intent ?? "", tool, cancellationToken);
             }
@@ -273,7 +272,6 @@ private async Task<CallToolResult> InvokeChildToolAsync(
             };
         }
 
-        Activity.Current?.SetTag(TagName.IsServerCommandInvoked, true);
         IReadOnlyDictionary<string, IBaseCommand> namespaceCommands;
         try
         {
@@ -404,6 +402,7 @@ private async Task<CallToolResult> InvokeChildToolAsync(
             // this case, which will be executed.
             currentActivity?.SetTag(TagName.ToolName, command);
 
+            request.Items[TagName.IsServerCommandInvoked] = true;
             var commandResponse = await cmd.ExecuteAsync(commandContext, commandOptions);
             var jsonResponse = JsonSerializer.Serialize(commandResponse, ModelsJsonContext.Default.CommandResponse);
             var isError = commandResponse.Status < HttpStatusCode.OK || commandResponse.Status >= HttpStatusCode.Ambiguous;
@@ -469,36 +468,53 @@ There was an error finding or calling tool and command.
 
     private async Task<CallToolResult> InvokeToolLearn(RequestContext<CallToolRequestParams> request, string? intent, string namespaceName, CancellationToken cancellationToken)
     {
-        Activity.Current?.SetTag(TagName.IsServerCommandInvoked, false);
-        var toolsJson = GetChildToolListJson(request, namespaceName);
+        request.Items[TagName.IsServerCommandInvoked] = false;
+        var availableTools = GetChildToolList(request, namespaceName);
 
-        var learnResponse = new CallToolResult
+        if (SupportsSampling(request.Server) && !string.IsNullOrWhiteSpace(intent))
         {
-            Content =
-            [
+            (string? commandName, IReadOnlyDictionary<string, JsonElement> parameters) = await GetCommandAndParametersFromIntentAsync(request, intent, namespaceName, availableTools, cancellationToken);
+            if (commandName != null)
+            {
+                var commandTool = availableTools.Where(tool => string.Equals(commandName, tool.Name)).FirstOrDefault();
+                if (commandTool != null)
+                {
+                    return new CallToolResult
+                    {
+                        Content = [
+                            new TextContentBlock {
+                                Text = $"""
+                                Here is the most suitable command and its parameters based on the intent '{intent}'.
+                                If you do not find the command suitable, run again with the "learn=true" to get a list of available commands and their parameters.
+                                Next, identify the command you want to execute and run again with the "command" and "parameters" arguments.
+
+                                {JsonSerializer.Serialize(commandTool, ServerJsonContext.Default.Tool)}
+                                """
+                            }
+                        ],
+                        IsError = false
+                    };
+                }
+            }
+        }
+
+        // Wasn't able to infer a specific tool command, either intent was missing, sampling wasn't support, or sampling failed.
+        // Return the list of available tools for the namespace.
+        return new CallToolResult
+        {
+            Content = [
                 new TextContentBlock {
                     Text = $"""
-                        Here are the available command and their parameters for '{namespaceName}' tool.
+                        Here are the available commands and their parameters for '{namespaceName}' tool.
                         If you do not find a suitable command, run again with the "learn=true" to get a list of available commands and their parameters.
                         Next, identify the command you want to execute and run again with the "command" and "parameters" arguments.
 
-                        {toolsJson}
+                        {JsonSerializer.Serialize(availableTools, ServerJsonContext.Default.ListTool)}
                         """
                 }
             ],
             IsError = false
         };
-        var response = learnResponse;
-        if (SupportsSampling(request.Server) && !string.IsNullOrWhiteSpace(intent))
-        {
-            var availableTools = GetChildToolList(request, namespaceName);
-            (string? commandName, IReadOnlyDictionary<string, JsonElement> parameters) = await GetCommandAndParametersFromIntentAsync(request, intent, namespaceName, availableTools, cancellationToken);
-            if (commandName != null)
-            {
-                response = await InvokeChildToolAsync(request, intent, namespaceName, commandName, parameters, cancellationToken);
-            }
-        }
-        return response;
     }
 
     /// <summary>
@@ -542,12 +558,6 @@ private List<Tool> GetChildToolList(RequestContext<CallToolRequestParams> reques
         return list;
     }
 
-    private string GetChildToolListJson(RequestContext<CallToolRequestParams> request, string namespaceName)
-    {
-        var listTools = GetChildToolList(request, namespaceName);
-        return JsonSerializer.Serialize(listTools, ServerJsonContext.Default.ListTool);
-    }
-
     private string GetChildToolJson(RequestContext<CallToolRequestParams> request, string namespaceName, string commandName)
     {
         var tools = GetChildToolList(request, namespaceName);

core/Azure.Mcp.Core/src/Areas/Server/Commands/ToolLoading/ServerToolLoader.cs

@@ -245,7 +245,7 @@ private async Task<CallToolResult> InvokeChildToolAsync(RequestContext<CallToolR
             }
 
             // At this point we should always have a valid command (child tool) call to invoke.
-            Activity.Current?.SetTag(TagName.IsServerCommandInvoked, true);
+            request.Items[TagName.IsServerCommandInvoked] = true;
             await NotifyProgressAsync(request, $"Calling {tool} {command}...", cancellationToken);
             var toolCallResponse = await client.CallToolAsync(command, parameters, cancellationToken: cancellationToken);
             if (toolCallResponse.IsError is true)
@@ -321,7 +321,7 @@ There was an error finding or calling tool and command.
 
     private async Task<CallToolResult> InvokeToolLearn(RequestContext<CallToolRequestParams> request, string? intent, string tool, CancellationToken cancellationToken)
     {
-        Activity.Current?.SetTag(TagName.IsServerCommandInvoked, false);
+        request.Items[TagName.IsServerCommandInvoked] = false;
         var toolsJson = await GetChildToolListJsonAsync(request, tool);
 
         var learnResponse = new CallToolResult

core/Azure.Mcp.Core/src/Areas/Server/Commands/ToolLoading/SingleProxyToolLoader.cs

@@ -1,9 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-using System.Diagnostics;
 using Azure.Mcp.Core.Areas.Server.Commands.Discovery;
-using Azure.Mcp.Core.Services.Telemetry;
 using Microsoft.Extensions.Logging;
 using ModelContextProtocol;
 using ModelContextProtocol.Client;
@@ -133,11 +131,11 @@ public override async ValueTask<CallToolResult> CallToolHandler(RequestContext<C
             learn = true;
         }
 
-        if (learn && string.IsNullOrEmpty(tool) && string.IsNullOrEmpty(command))
+        if (learn && string.IsNullOrEmpty(tool))
         {
             return await RootLearnModeAsync(request, intent ?? "", cancellationToken);
         }
-        else if (learn && !string.IsNullOrEmpty(tool) && string.IsNullOrEmpty(command))
+        else if (learn && !string.IsNullOrEmpty(tool))
         {
             return await ToolLearnModeAsync(request, intent ?? "", tool!, cancellationToken);
         }
@@ -205,7 +203,7 @@ private async Task<string> GetToolListJsonAsync(RequestContext<CallToolRequestPa
 
     private async Task<CallToolResult> RootLearnModeAsync(RequestContext<CallToolRequestParams> request, string intent, CancellationToken cancellationToken)
     {
-        Activity.Current?.SetTag(TagName.IsServerCommandInvoked, false);
+        request.Items[TagName.IsServerCommandInvoked] = false;
         var toolsJson = await GetRootToolsJsonAsync();
         var learnResponse = new CallToolResult
         {
@@ -236,7 +234,7 @@ Here are the available list of tools.
 
     private async Task<CallToolResult> ToolLearnModeAsync(RequestContext<CallToolRequestParams> request, string intent, string tool, CancellationToken cancellationToken)
     {
-        Activity.Current?.SetTag(TagName.IsServerCommandInvoked, false);
+        request.Items[TagName.IsServerCommandInvoked] = false;
         var toolsJson = await GetToolListJsonAsync(request, tool);
         if (string.IsNullOrEmpty(toolsJson))
         {
@@ -273,7 +271,7 @@ private async Task<CallToolResult> ToolLearnModeAsync(RequestContext<CallToolReq
 
     private async Task<CallToolResult> CommandModeAsync(RequestContext<CallToolRequestParams> request, string intent, string tool, string command, Dictionary<string, object?> parameters, CancellationToken cancellationToken)
     {
-        Activity.Current?.SetTag(TagName.IsServerCommandInvoked, true);
+        request.Items[TagName.IsServerCommandInvoked] = true;
         McpClient? client;
 
         try

core/Azure.Mcp.Core/src/Commands/CommandFactory.cs

@@ -201,7 +201,7 @@ private void ConfigureCommandHandler(Command command, IBaseCommand implementatio
 
                 if (response.Status == HttpStatusCode.OK && response.Results == null)
                 {
-                    response.Results = ResponseResult.Create(new List<string>(), JsonSourceGenerationContext.Default.ListString);
+                    response.Results = ResponseResult.Create([], JsonSourceGenerationContext.Default.ListString);
                 }
 
                 var isServiceStartCommand = implementation is Areas.Server.Commands.ServiceStartCommand;