Add the built-in tool run_command_in_terminal for AI to execute commands in the connected PowerShell session (#398)

Allow AIShell to run command in the connected PowerShell session and collect all output and error.

1. Add `Invoke-AICommand` cmdlet (alias: `airun`) to `AIShell` module. Commands sent from the sidecar AIShell will be executed through this command in the form of `airun { <command> }`. This command is designed to collect all output and error messages as they are displayed in the terminal, while preserving the streaming behavior as expected.

2. Add `RunCommand` and `PostResult` messages to the protocol.

3. Update the `Channel` class in `AIShell` module to support the `OnRunCommand` action. We already support posting command to the PowerShell's prompt, but it turns out not easy to make the command be accepted. On Windows, we have to call `AcceptLine` within an `OnIdle` event handler and it also requires changes to `PSReadLine`.
    - `AcceptLine` only set a flag in `PSReadLine` to indicate the line was accepted. The flag is checked in `InputLoop`, however, when `PSReadLine` is waiting for input, it's blocked in the `ReadKey` call within `InputLoop`, so even if the flag is set, `InputLoop` won't be able to check the flag until after `ReadKey` call is returned.
    - I need to change PSReadLine a bit: after it finishes handling the `OnIdle` event, it checks if the `_lineAccepted` flag is set. If it's set, it means `AcceptLine` got called within the `OnIdle` handler, and it throws a `LineAcceptedException` to break out from `ReadKey`. I catch this exception in `InputLoop` to continue with the flag check.
    - However, a problem with this change is: the "readkey thread" is still blocked on `Console.ReadKey` when the command is returned to PowerShell to execute. On Windows, this could cause minor issues if the command also calls `Console.ReadKey` -- 2 threads calling `Console.ReadKey` in parallel, so it's uncertain which will get the next keystroke input. On macOS and Linux, the problem is way much bigger -- any subsequent writing to the terminal may be blocked, because on Unix platforms, reading cursor position will be blocked if another thread is calling `Console.ReadKey`.
    - So, this approach can only work on Windows. On macOS, we will have to depend on `iTerm2`, which has a Python API server that allows to send keystrokes to a tab using the Python API, so we could possibly use that for macOS. But Windows Terminal doesn't support that, and thus we will have to use the above approach to accept the command on Windows.
    - On macOS, if the Python API approach works fine, then we could even consider using it for the `PostCode` action.

4. Add `run_command_in_terminal` and `get_command_output` tools and expose them to agents.

Author: daxian-dbw

shell/AIShell.Abstraction/NamedPipe.cs

@@ -33,6 +33,21 @@ public enum MessageType : int
     /// A message from AIShell to command-line shell to send code block.
     /// </summary>
     PostCode = 4,
+
+    /// <summary>
+    /// A message from AIShell to command-line shell to run a command.
+    /// </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 = 7,
 }
 
 /// <summary>
@@ -201,6 +216,95 @@ public PostCodeMessage(List<string> codeBlocks)
     }
 }
 
+/// <summary>
+/// Message for <see cref="MessageType.RunCommand"/>.
+/// </summary>
+public sealed class RunCommandMessage : PipeMessage
+{
+    /// <summary>
+    /// Gets the command to run.
+    /// </summary>
+    public string Command { get; }
+
+    /// <summary>
+    /// Gets whether the command should be run in blocking mode.
+    /// </summary>
+    public bool Blocking { get; }
+
+    /// <summary>
+    /// Creates an instance of <see cref="RunCommandMessage"/>.
+    /// </summary>
+    public RunCommandMessage(string command, bool blocking)
+        : base(MessageType.RunCommand)
+    {
+        ArgumentException.ThrowIfNullOrEmpty(command);
+
+        Command = command;
+        Blocking = 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>
+public sealed class PostResultMessage : PipeMessage
+{
+    /// <summary>
+    /// Gets the result of the command for a blocking 'run_command' too call.
+    /// Or, for a non-blocking call, gets the id for retrieving the result later.
+    /// </summary>
+    public string Output { get; }
+
+    /// <summary>
+    /// Gets whether the command execution had any error.
+    /// i.e. a native command returned a non-zero exit code, or a powershell command threw any errors.
+    /// </summary>
+    public bool HadError { get; }
+
+    /// <summary>
+    /// Gets a value indicating whether the operation was canceled by the user.
+    /// </summary>
+    public bool UserCancelled { get; }
+
+    /// <summary>
+    /// Gets the internal exception message that is thrown when trying to run the command.
+    /// </summary>
+    public string Exception { get; }
+
+    /// <summary>
+    /// Creates an instance of <see cref="PostResultMessage"/>.
+    /// </summary>
+    public PostResultMessage(string output, bool hadError, bool userCancelled, string exception)
+        : base(MessageType.PostResult)
+    {
+        Output = output;
+        HadError = hadError;
+        UserCancelled = userCancelled;
+        Exception = exception;
+    }
+}
+
 /// <summary>
 /// The base type for common pipe operations.
 /// </summary>
@@ -301,7 +405,7 @@ protected async Task<PipeMessage> GetMessageAsync(CancellationToken cancellation
             return null;
         }
 
-        if (type > (int)MessageType.PostCode)
+        if (type > (int)MessageType.PostResult)
         {
             _pipeStream.Close();
             throw new IOException($"Unknown message type received: {type}. Connection was dropped.");
@@ -344,9 +448,12 @@ private static PipeMessage DeserializePayload(int type, ReadOnlySpan<byte> bytes
         {
             (int)MessageType.PostQuery => JsonSerializer.Deserialize<PostQueryMessage>(bytes),
             (int)MessageType.AskConnection => JsonSerializer.Deserialize<AskConnectionMessage>(bytes),
-            (int)MessageType.PostContext => JsonSerializer.Deserialize<PostContextMessage>(bytes),
             (int)MessageType.AskContext => JsonSerializer.Deserialize<AskContextMessage>(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"),
         };
     }
@@ -465,6 +572,16 @@ public async Task StartProcessingAsync(int timeout, CancellationToken cancellati
                     InvokeOnPostCode((PostCodeMessage)message);
                     break;
 
+                case MessageType.RunCommand:
+                    var result = InvokeOnRunCommand((RunCommandMessage)message);
+                    SendMessage(result);
+                    break;
+
+                case MessageType.AskCommandOutput:
+                    var output = InvokeOnAskCommandOutput((AskCommandOutputMessage)message);
+                    SendMessage(output);
+                    break;
+
                 default:
                     // Log: unexpected messages ignored.
                     break;
@@ -537,6 +654,66 @@ 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)
+        {
+            // Log: event handler not set.
+            return new PostResultMessage(
+                output: "Command execution is not supported.",
+                hadError: true,
+                userCancelled: false,
+                exception: null);
+        }
+
+        try
+        {
+            return OnRunCommand(message);
+        }
+        catch (Exception e)
+        {
+            // Log: exception when invoking 'OnRunCommand'
+            return new PostResultMessage(
+                output: "Failed to execute the command due to an internal error.",
+                hadError: true,
+                userCancelled: false,
+                exception: e.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>
@@ -551,6 +728,16 @@ private PostContextMessage InvokeOnAskContext(AskContextMessage message)
     /// Event for handling the <see cref="MessageType.AskContext"/> message.
     /// </summary>
     public event Func<AskContextMessage, PostContextMessage> OnAskContext;
+
+    /// <summary>
+    /// 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>
@@ -771,4 +958,52 @@ 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.
+        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;
+    }
+
+    /// <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/AIShell.psd1

@@ -10,9 +10,9 @@
     PowerShellVersion = '7.4.6'
     PowerShellHostName = 'ConsoleHost'
     FunctionsToExport = @()
-    CmdletsToExport = @('Start-AIShell','Invoke-AIShell','Resolve-Error')
+    CmdletsToExport = @('Start-AIShell','Invoke-AIShell', 'Invoke-AICommand', 'Resolve-Error')
     VariablesToExport = '*'
-    AliasesToExport = @('aish', 'askai', 'fixit')
+    AliasesToExport = @('aish', 'askai', 'fixit', 'airun')
     HelpInfoURI = 'https://aka.ms/aishell-help'
     PrivateData = @{ PSData = @{ Prerelease = 'preview5'; ProjectUri = 'https://github.com/PowerShell/AIShell' } }
 }