Enable 'run_command_in_terminal' and 'get_command_output' tools

Author: daxian-dbw

shell/AIShell.Abstraction/NamedPipe.cs

@@ -39,10 +39,15 @@ public enum MessageType : int
     /// </summary>
     RunCommand = 5,
 
+    /// <summary>
+    /// A message from AIShell to command-line shell to ask for the result of a previous command run.
+    /// </summary>
+    AskCommandOutput = 6,
+
     /// <summary>
     /// A message from command-line shell to AIShell to post the result of a command.
     /// </summary>
-    PostResult = 6,
+    PostResult = 7,
 }
 
 /// <summary>
@@ -239,6 +244,27 @@ public RunCommandMessage(string command, bool blocking)
     }
 }
 
+/// <summary>
+/// Message for <see cref="MessageType.AskCommandOutput"/>.
+/// </summary>
+public sealed class AskCommandOutputMessage : PipeMessage
+{
+    /// <summary>
+    /// Gets the id of the command to retrieve the output for.
+    /// </summary>
+    public string CommandId { get; }
+
+    /// <summary>
+    /// Creates an instance of <see cref="AskCommandOutputMessage"/>.
+    /// </summary>
+    public AskCommandOutputMessage(string commandId)
+        : base(MessageType.AskCommandOutput)
+    {
+        ArgumentException.ThrowIfNullOrEmpty(commandId);
+        CommandId = commandId;
+    }
+}
+
 /// <summary>
 /// Message for <see cref="MessageType.PostResult"/>.
 /// </summary>
@@ -426,6 +452,7 @@ private static PipeMessage DeserializePayload(int type, ReadOnlySpan<byte> bytes
             (int)MessageType.PostContext => JsonSerializer.Deserialize<PostContextMessage>(bytes),
             (int)MessageType.PostCode => JsonSerializer.Deserialize<PostCodeMessage>(bytes),
             (int)MessageType.RunCommand => JsonSerializer.Deserialize<RunCommandMessage>(bytes),
+            (int)MessageType.AskCommandOutput => JsonSerializer.Deserialize<AskCommandOutputMessage>(bytes),
             (int)MessageType.PostResult => JsonSerializer.Deserialize<PostResultMessage>(bytes),
             _ => throw new NotSupportedException("Unreachable code"),
         };
@@ -550,6 +577,11 @@ public async Task StartProcessingAsync(int timeout, CancellationToken cancellati
                     SendMessage(result);
                     break;
 
+                case MessageType.AskCommandOutput:
+                    var output = InvokeOnAskCommandOutput((AskCommandOutputMessage)message);
+                    SendMessage(output);
+                    break;
+
                 default:
                     // Log: unexpected messages ignored.
                     break;
@@ -622,6 +654,9 @@ private PostContextMessage InvokeOnAskContext(AskContextMessage message)
         return null;
     }
 
+    /// <summary>
+    /// Helper to invoke the <see cref="OnRunCommand"/> event.
+    /// </summary>
     private PostResultMessage InvokeOnRunCommand(RunCommandMessage message)
     {
         if (OnRunCommand is null)
@@ -649,6 +684,36 @@ private PostResultMessage InvokeOnRunCommand(RunCommandMessage message)
         }
     }
 
+    /// <summary>
+    /// Helper to invoke the <see cref="OnAskCommandOutput"/> event.
+    /// </summary>
+    private PostResultMessage InvokeOnAskCommandOutput(AskCommandOutputMessage message)
+    {
+        if (OnAskCommandOutput is null)
+        {
+            // Log: event handler not set.
+            return new PostResultMessage(
+                output: "Retrieving command output is not supported.",
+                hadError: true,
+                userCancelled: false,
+                exception: null);
+        }
+
+        try
+        {
+            return OnAskCommandOutput(message);
+        }
+        catch (Exception e)
+        {
+            // Log: exception when invoking 'OnAskCommandOutput'
+            return new PostResultMessage(
+                output: "Failed to retrieve the command output due to an internal error.",
+                hadError: true,
+                userCancelled: false,
+                exception: e.Message);
+        }
+    }
+
     /// <summary>
     /// Event for handling the <see cref="MessageType.PostCode"/> message.
     /// </summary>
@@ -668,6 +733,11 @@ private PostResultMessage InvokeOnRunCommand(RunCommandMessage message)
     /// Event for handling the <see cref="MessageType.RunCommand"/> message.
     /// </summary>
     public event Func<RunCommandMessage, PostResultMessage> OnRunCommand;
+
+    /// <summary>
+    /// Event for handling the <see cref="MessageType.AskCommandOutput"/> message.
+    /// </summary>
+    public event Func<AskCommandOutputMessage, PostResultMessage> OnAskCommandOutput;
 }
 
 /// <summary>
@@ -889,6 +959,13 @@ public async Task<PostContextMessage> AskContext(AskContextMessage message, Canc
         return postContext;
     }
 
+    /// <summary>
+    /// Run a command in the connected PowerShell session.
+    /// </summary>
+    /// <param name="message">The <see cref="MessageType.RunCommand"/> message.</param>
+    /// <param name="cancellationToken">A cancellation token.</param>
+    /// <returns>A <see cref="MessageType.PostResult"/> message as the response.</returns>
+    /// <exception cref="IOException">Throws when the pipe is closed by the other side.</exception>
     public async Task<PostResultMessage> RunCommand(RunCommandMessage message, CancellationToken cancellationToken)
     {
         // Send the request message to the shell.
@@ -905,4 +982,28 @@ public async Task<PostResultMessage> RunCommand(RunCommandMessage message, Cance
 
         return postResult;
     }
+
+    /// <summary>
+    /// Ask for the output of a previously run command in the connected PowerShell session.
+    /// </summary>
+    /// <param name="message">The <see cref="MessageType.AskCommandOutput"/> message.</param>
+    /// <param name="cancellationToken">A cancellation token.</param>
+    /// <returns>A <see cref="MessageType.PostResult"/> message as the response.</returns>
+    /// <exception cref="IOException">Throws when the pipe is closed by the other side.</exception>
+    public async Task<PostResultMessage> AskCommandOutput(AskCommandOutputMessage message, CancellationToken cancellationToken)
+    {
+        // Send the request message to the shell.
+        SendMessage(message);
+
+        // Receiving response from the shell.
+        var response = await GetMessageAsync(cancellationToken);
+        if (response is not PostResultMessage postResult)
+        {
+            // Log: unexpected message. drop connection.
+            _client.Close();
+            throw new IOException($"Expecting '{MessageType.PostResult}' response, but received '{message.Type}' message.");
+        }
+
+        return postResult;
+    }
 }

shell/AIShell.Integration/Channel.cs

@@ -151,6 +151,7 @@ public string StartChannelSetup()
         _serverPipe.OnAskContext += OnAskContext;
         _serverPipe.OnPostCode += OnPostCode;
         _serverPipe.OnRunCommand += OnRunCommand;
+        _serverPipe.OnAskCommandOutput += OnAskCommandOutput;
 
         _serverThread = new Thread(ThreadProc)
         {
@@ -284,6 +285,7 @@ private void Reset()
             _serverPipe.OnAskContext -= OnAskContext;
             _serverPipe.OnPostCode -= OnPostCode;
             _serverPipe.OnRunCommand -= OnRunCommand;
+            _serverPipe.OnAskCommandOutput -= OnAskCommandOutput;
         }
 
         _serverPipe = null;
@@ -488,12 +490,6 @@ private void OnAskConnection(ShellClientPipe clientPipe, Exception exception)
         _connSetupWaitHandler.Set();
     }
 
-    private PostContextMessage OnAskContext(AskContextMessage askContextMessage)
-    {
-        // Not implemented yet.
-        return null;
-    }
-
     private PostResultMessage OnRunCommand(RunCommandMessage runCommandMessage)
     {
         // Ignore 'run_command' request when a code posting operation is on-going.
@@ -544,14 +540,13 @@ private PostResultMessage OnRunCommand(RunCommandMessage runCommandMessage)
             _runCommandRequest.Event.Wait();
             RunCommandResult result = _runCommandRequest.Result;
 
-            _pwsh ??= PowerShell.Create();
-            _pwsh.Commands.Clear();
             string output = result.ErrorAndOutput.Count is 0
                 ? string.Empty
-                : _pwsh.AddCommand("Out-String")
+                : (_pwsh ??= PowerShell.Create())
+                    .AddCommand("Out-String")
                     .AddParameter("InputObject", result.ErrorAndOutput)
                     .AddParameter("Width", 120)
-                    .Invoke<string>()[0];
+                    .InvokeAndCleanup<string>()[0];
 
             PostResultMessage response = new(
                 output: output,
@@ -568,6 +563,57 @@ private PostResultMessage OnRunCommand(RunCommandMessage runCommandMessage)
         return new PostResultMessage(output: _runCommandRequest.Id, hadError: false, userCancelled: false, exception: null);
     }
 
+    private PostResultMessage OnAskCommandOutput(AskCommandOutputMessage askOutputMessage)
+    {
+        if (_runCommandRequest is null)
+        {
+            return new PostResultMessage(
+                output: "No command was previously run in background, or the output of a background command was already retrieved.",
+                hadError: true,
+                userCancelled: false,
+                exception: null);
+        }
+
+        string commandId = askOutputMessage.CommandId;
+        if (!string.Equals(commandId, _runCommandRequest.Id, StringComparison.OrdinalIgnoreCase))
+        {
+            return new PostResultMessage(
+                output: $"The specified command id '{commandId}' cannot be found.",
+                hadError: true,
+                userCancelled: false,
+                exception: null);
+        }
+
+        if (_runCommandRequest.Result is null)
+        {
+            return new PostResultMessage(
+                output: "Command output is not yet available.",
+                hadError: true,
+                userCancelled: false,
+                exception: null);
+        }
+
+        RunCommandResult result = _runCommandRequest.Result;
+        string output = result.ErrorAndOutput.Count is 0
+            ? string.Empty
+            : (_pwsh ??= PowerShell.Create())
+                .AddCommand("Out-String")
+                .AddParameter("InputObject", result.ErrorAndOutput)
+                .AddParameter("Width", 120)
+                .InvokeAndCleanup<string>()[0];
+
+        PostResultMessage response = new(
+                output: output,
+                hadError: result.HadErrors,
+                userCancelled: result.UserCancelled,
+                exception: null);
+
+        _runCommandRequest.Dispose();
+        _runCommandRequest = null;
+
+        return response;
+    }
+
     private void PSRLInsert(string text)
     {
         using var _ = new NoWindowResizingCheck();

shell/AIShell.Kernel/Command/CodeCommand.cs

@@ -36,13 +36,6 @@ public CodeCommand()
         copy.SetHandler(CopyAction, nth);
         save.SetHandler(SaveAction, file, append);
         post.SetHandler(PostAction, nth);
-
-
-        var run = new Command("run", "Run the specified command.");
-        var command = new Argument<string>("command", "command to run");
-        run.AddArgument(command);
-        run.SetHandler(RunAction, command);
-        AddCommand(run);
     }
 
     private static string GetCodeText(Shell shell, int index)
@@ -188,13 +181,4 @@ private void PostAction(int nth)
             host.WriteErrorLine(e.Message);
         }
     }
-
-    private async Task RunAction(string command)
-    {
-        var shell = (Shell)Shell;
-        var host = shell.Host;
-        var result = await shell.Channel.RunCommand(new RunCommandMessage(command, blocking: true), shell.CancellationToken);
-
-        host.WriteLine($"HadError: {result.HadError}\nUserCancelled: {result.UserCancelled}\nOutput: {result.Output}\nException: {result.Exception}");
-    }
 }

shell/AIShell.Kernel/Host.cs

@@ -570,9 +570,9 @@ internal void RenderReferenceText(string header, string content)
     /// </summary>
     /// <param name="tool">The MCP tool.</param>
     /// <param name="jsonArgs">The arguments in JSON form to be sent for the tool call.</param>
-    internal void RenderToolCallRequest(McpTool tool, string jsonArgs)
+    internal void RenderMcpToolCallRequest(McpTool tool, string jsonArgs)
     {
-        RequireStdoutOrStderr(operation: "render tool call request");
+        RequireStdoutOrStderr(operation: "render MCP tool call request");
         IAnsiConsole ansiConsole = _outputRedirected ? _stderrConsole : AnsiConsole.Console;
 
         bool hasArgs = !string.IsNullOrEmpty(jsonArgs);
@@ -610,6 +610,44 @@ internal void RenderToolCallRequest(McpTool tool, string jsonArgs)
         FancyStreamRender.ConsoleUpdated();
     }
 
+    /// <summary>
+    /// Render the built-in tool call request.
+    /// </summary>
+    internal void RenderBuiltInToolCallRequest(string toolName, string description, Tuple<string, string> argument)
+    {
+        RequireStdoutOrStderr(operation: "render built-in tool call request");
+        IAnsiConsole ansiConsole = _outputRedirected ? _stderrConsole : AnsiConsole.Console;
+
+        bool hasArgs = argument is not null;
+        string argLine = hasArgs ? $"{argument.Item1}:" : $"Input: <none>";
+        IRenderable content = new Markup($"""
+
+            [bold]Run [olive]{toolName}[/] from [olive]{McpManager.BuiltInServerName}[/] (Built-in tool)[/]
+
+            {description}
+
+            {argLine}
+            """);
+
+        if (hasArgs)
+        {
+            content = new Grid()
+                .AddColumn(new GridColumn())
+                .AddRow(content)
+                .AddRow(argument.Item2.EscapeMarkup());
+        }
+
+        var panel = new Panel(content)
+            .Expand()
+            .RoundedBorder()
+            .Header("[green]  Tool Call Request  [/]")
+            .BorderColor(Color.Grey);
+
+        ansiConsole.WriteLine();
+        ansiConsole.Write(panel);
+        FancyStreamRender.ConsoleUpdated();
+    }
+
     /// <summary>
     /// Render a table with information about available MCP servers and tools.
     /// </summary>
@@ -672,7 +710,13 @@ internal void RenderMcpServersAndTools(McpManager mcpManager)
             toolTable.AddRow($"[olive underline]{McpManager.BuiltInServerName}[/]", "[green]\u2713 Ready[/]", string.Empty);
             foreach (var item in mcpManager.BuiltInTools)
             {
-                toolTable.AddRow(string.Empty, item.Key.EscapeMarkup(), item.Value.Description.EscapeMarkup());
+                string description = item.Value.Description;
+                int index = description.IndexOf('\n');
+                if (index > 0)
+                {
+                    description = description[..index].Trim();
+                }
+                toolTable.AddRow(string.Empty, item.Key.EscapeMarkup(), description.EscapeMarkup());
             }
         }