From 786ca72a4cac269e52edeef3b6889cfc4db6c202 Mon Sep 17 00:00:00 2001
From: Ionite <dev@ionite.io>
Date: Fri, 14 Jul 2023 21:00:08 -0400
Subject: [PATCH] Add input sending for launch page console

---
 .../Assets/sitecustomize.py                   | 49 +++++++++++++
 .../ViewModels/LaunchPageViewModel.cs         | 70 ++++++++++++++++++-
 .../Views/LaunchPageView.axaml                | 62 +++++++++++++++-
 .../Models/Packages/BaseGitPackage.cs         | 26 ++++++-
 StabilityMatrix.Core/Processes/AnsiProcess.cs |  6 ++
 StabilityMatrix.Core/Processes/ApcMessage.cs  | 12 ++++
 StabilityMatrix.Core/Processes/ApcParser.cs   | 55 +++++++++++++++
 StabilityMatrix.Core/Processes/ApcType.cs     | 11 +++
 .../Processes/AsyncStreamReader.cs            | 69 +++++++++++++++---
 .../Processes/ProcessOutput.cs                | 27 +++++--
 10 files changed, 369 insertions(+), 18 deletions(-)
 create mode 100644 StabilityMatrix.Avalonia/Assets/sitecustomize.py
 create mode 100644 StabilityMatrix.Core/Processes/ApcMessage.cs
 create mode 100644 StabilityMatrix.Core/Processes/ApcParser.cs
 create mode 100644 StabilityMatrix.Core/Processes/ApcType.cs

diff --git a/StabilityMatrix.Avalonia/Assets/sitecustomize.py b/StabilityMatrix.Avalonia/Assets/sitecustomize.py
new file mode 100644
index 00000000..c51abc9d
--- /dev/null
+++ b/StabilityMatrix.Avalonia/Assets/sitecustomize.py
@@ -0,0 +1,49 @@
+"""
+Startup site customization for Stability Matrix.
+
+Currently this installs an audit hook to notify the parent process when input() is called,
+so we can prompt the user to enter something.
+"""
+
+import sys
+
+# Application Program Command escape sequence
+# This wraps messages sent to the parent process.
+esc_apc = "\x9F"
+esc_prefix = "[SM;"
+esc_st = "\x9C"
+
+
+def send_apc(msg: str):
+    """Send an Application Program Command to the parent process."""
+    sys.stdout.flush()
+    sys.stdout.write(esc_apc + esc_prefix + msg + esc_st)
+    sys.stdout.flush()
+
+
+def send_apc_input(prompt: str):
+    """Apc message for input() prompt."""
+    send_apc('{"type":"input","data":"' + str(prompt) + '"}')
+
+
+def audit(event: str, *args):
+    """Main audit hook function."""
+    # https://docs.python.org/3/library/functions.html#input
+    # input() raises audit event `builtins.input` with args (prompt: str) *before* reading from stdin.
+    # `builtins.input/result` raised after reading from stdin.
+
+    if event == "builtins.input":
+        try:
+            prompts = args[0] if args else ()
+            prompt = "".join(prompts)
+            send_apc_input(prompt)
+        except Exception:
+            pass
+
+
+# Reconfigure stdout to UTF-8
+# noinspection PyUnresolvedReferences
+sys.stdout.reconfigure(encoding="utf-8")
+
+# Install the audit hook
+sys.addaudithook(audit)
diff --git a/StabilityMatrix.Avalonia/ViewModels/LaunchPageViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/LaunchPageViewModel.cs
index 19b8aa83..428d380b 100644
--- a/StabilityMatrix.Avalonia/ViewModels/LaunchPageViewModel.cs
+++ b/StabilityMatrix.Avalonia/ViewModels/LaunchPageViewModel.cs
@@ -28,6 +28,7 @@ using StabilityMatrix.Core.Models.Packages;
 using StabilityMatrix.Core.Processes;
 using StabilityMatrix.Core.Python;
 using StabilityMatrix.Core.Services;
+using ThreadState = System.Diagnostics.ThreadState;
 
 namespace StabilityMatrix.Avalonia.ViewModels;
 
@@ -65,6 +66,10 @@ public partial class LaunchPageViewModel : PageViewModelBase, IDisposable
     
     // private bool clearingPackages;
     private string webUiUrl = string.Empty;
+    
+    // Input info-bars
+    [ObservableProperty] private bool showManualInputPrompt;
+    [ObservableProperty] private bool showConfirmInputPrompt;
 
     public LaunchPageViewModel(ILogger<LaunchPageViewModel> logger, ISettingsManager settingsManager, IPackageFactory packageFactory,
         IPyRunner pyRunner, INotificationService notificationService, ServiceManager<ViewModelBase> dialogFactory)
@@ -248,12 +253,27 @@ public partial class LaunchPageViewModel : PageViewModelBase, IDisposable
     
     private async Task BeginUpdateConsole(CancellationToken ct)
     {
+        // This should be run in the UI thread
+        Dispatcher.UIThread.CheckAccess();
         try
         {
             while (true)
             {
                 ct.ThrowIfCancellationRequested();
                 var output = await consoleUpdateBuffer.ReceiveAsync(ct);
+                // Check for Apc messages
+                if (output.ApcMessage is not null)
+                {
+                    // Handle Apc message, for now just input audit events
+                    var message = output.ApcMessage.Value;
+                    if (message.Type == ApcType.Input)
+                    {
+                        ShowConfirmInputPrompt = true;
+                    }
+                    // Ignore further processing
+                    continue;
+                }
+                
                 using var update = ConsoleDocument.RunUpdate();
                 // Handle remove
                 if (output.ClearLines > 0)
@@ -274,6 +294,54 @@ public partial class LaunchPageViewModel : PageViewModelBase, IDisposable
         }
     }
     
+    // Send user input to running package
+    public async Task SendInput(string input)
+    {
+        if (RunningPackage is BaseGitPackage package)
+        {
+            var venv = package.VenvRunner;
+            var process = venv?.Process;
+            if (process is not null)
+            {
+                await process.StandardInput.WriteLineAsync(input);
+            }
+            else
+            {
+                logger.LogWarning("Attempted to write input but Process is null");
+            }
+        }
+    }
+
+    [RelayCommand]
+    private async Task SendConfirmInput(bool value)
+    {
+        // This must be on the UI thread
+        Dispatcher.UIThread.CheckAccess();
+        // Also send input to our own console
+        if (value)
+        {
+            consoleUpdateBuffer.Post(new ProcessOutput { Text = "y\n" });
+            await SendInput("y\n");
+        }
+        else
+        {
+            consoleUpdateBuffer.Post(new ProcessOutput { Text = "n\n" });
+            await SendInput("n\n");
+        }
+
+        ShowConfirmInputPrompt = false;
+    }
+    
+    // Handle user input requests
+    public async Task HandleApcMessage(ApcMessage message)
+    {
+        // Handle inputs by prompting
+        if (message.Type == ApcType.Input)
+        {
+            ShowConfirmInputPrompt = true;
+        }
+    }
+    
     public async Task Stop()
     {
         if (RunningPackage is null) return;
@@ -335,7 +403,7 @@ public partial class LaunchPageViewModel : PageViewModelBase, IDisposable
         consoleUpdateBuffer.Post(output);
         EventManager.Instance.OnScrollToBottomRequested();
     }
-    
+
     private void OnOneClickInstallFinished(object? sender, bool e)
     {
         OnLoaded();
diff --git a/StabilityMatrix.Avalonia/Views/LaunchPageView.axaml b/StabilityMatrix.Avalonia/Views/LaunchPageView.axaml
index 6b882f10..774e38fa 100644
--- a/StabilityMatrix.Avalonia/Views/LaunchPageView.axaml
+++ b/StabilityMatrix.Avalonia/Views/LaunchPageView.axaml
@@ -10,15 +10,21 @@
     xmlns:models="clr-namespace:StabilityMatrix.Core.Models;assembly=StabilityMatrix.Core"
     xmlns:ui="using:FluentAvalonia.UI.Controls"
     xmlns:vm="clr-namespace:StabilityMatrix.Avalonia.ViewModels"
+    xmlns:system="clr-namespace:System;assembly=System.Runtime"
     d:DataContext="{x:Static mocks:DesignData.LaunchPageViewModel}"
     d:DesignHeight="450"
     d:DesignWidth="700"
     x:CompileBindings="True"
     x:DataType="vm:LaunchPageViewModel"
     mc:Ignorable="d">
+    
+    <controls:UserControlBase.Resources>
+        <system:Boolean x:Key="True">True</system:Boolean>
+        <system:Boolean x:Key="False">False</system:Boolean>
+    </controls:UserControlBase.Resources>
 
     <Grid RowDefinitions="Auto,*,Auto">
-        <Grid ColumnDefinitions="Auto,*"
+        <Grid ColumnDefinitions="Auto,*,Auto"
               Margin="0,8,0, 8">
             <Grid ColumnDefinitions="0.8*,0.2*">
                 <Button
@@ -59,7 +65,7 @@
                 x:Name="SelectPackageComboBox"
                 Grid.Row="0"
                 Grid.Column="1"
-                Margin="8,8,16,0"
+                Margin="8,8,0,0"
                 HorizontalAlignment="Stretch"
                 VerticalAlignment="Top"
                 ItemsSource="{Binding InstalledPackages}"
@@ -82,6 +88,19 @@
                     </DataTemplate>
                 </ComboBox.ItemTemplate>
             </ComboBox>
+            
+            <!-- Keyboard button to show manual input info bar -->
+            <ToggleButton
+                Grid.Column="2"
+                Width="48"
+                Margin="8,8,16,0"
+                IsChecked="{Binding ShowManualInputPrompt}"
+                ToolTip.Tip="Send Input"
+                HorizontalAlignment="Left"
+                VerticalAlignment="Stretch"
+                FontSize="16">
+                <ui:SymbolIcon FontSize="18" Symbol="Keyboard" />
+            </ToggleButton>
         </Grid>
 
         <avaloniaEdit:TextEditor
@@ -96,6 +115,45 @@
             VerticalScrollBarVisibility="Auto"
             WordWrap="True" />
         
+        <Grid Grid.Row="1" ColumnDefinitions="0.5*,*">
+            <StackPanel Grid.Column="1" Margin="8" Spacing="4">
+                <!-- Info bar for manual input -->
+                <ui:InfoBar 
+                    Title="Input"
+                    IsIconVisible="False"
+                    IsClosable="False"
+                    Severity="Informational"
+                    IsOpen="{Binding ShowManualInputPrompt, Mode=TwoWay}" 
+                    Margin="0">
+                    <ui:InfoBar.ActionButton>
+                        <Grid ColumnDefinitions="*,auto">
+                            <TextBox Name="ManualInputBox" Margin="0,0,8,0" />
+                            <Button Grid.Column="1" Margin="0,0,8,0" Content="Send"
+                                    Command="{Binding SendInput}"
+                                    CommandParameter="{Binding #ManualInputBox.Text}"/>
+                        </Grid>
+                    </ui:InfoBar.ActionButton>
+                </ui:InfoBar>
+                <!-- Info bar for auto prompts -->
+                <ui:InfoBar 
+                    Title="Input required" 
+                    Severity="Informational"
+                    IsOpen="{Binding ShowConfirmInputPrompt, Mode=TwoWay}" 
+                    Margin="0"
+                    Message="Confirm?">
+                    <ui:InfoBar.ActionButton>
+                        <StackPanel Orientation="Horizontal" Spacing="8">
+                            <Button Classes="accent" Content="Yes" 
+                                    Command="{Binding SendConfirmInputCommand}"
+                                    CommandParameter="{StaticResource True}"/>
+                            <Button Content="No" Command="{Binding SendConfirmInputCommand}"
+                                    CommandParameter="{StaticResource False}"/>
+                        </StackPanel>
+                    </ui:InfoBar.ActionButton>
+                </ui:InfoBar>
+            </StackPanel>
+        </Grid>
+
         <Button
             Grid.Row="2"
             Grid.ColumnSpan="2"
diff --git a/StabilityMatrix.Core/Models/Packages/BaseGitPackage.cs b/StabilityMatrix.Core/Models/Packages/BaseGitPackage.cs
index 05e4b6ed..779f9b95 100644
--- a/StabilityMatrix.Core/Models/Packages/BaseGitPackage.cs
+++ b/StabilityMatrix.Core/Models/Packages/BaseGitPackage.cs
@@ -22,7 +22,7 @@ public abstract class BaseGitPackage : BasePackage
     protected readonly ISettingsManager SettingsManager;
     protected readonly IDownloadService DownloadService;
     protected readonly IPrerequisiteHelper PrerequisiteHelper;
-    protected PyVenvRunner? VenvRunner;
+    public PyVenvRunner? VenvRunner;
     
     /// <summary>
     /// URL of the hosted web page on launch
@@ -215,7 +215,29 @@ public abstract class BaseGitPackage : BasePackage
             return latestCommit.Sha;
         }
     }
-
+    
+    // Send input to the running process.
+    public virtual void SendInput(string input)
+    {
+        var process = VenvRunner?.Process;
+        if (process == null)
+        {
+            Logger.Warn("No process running for {Name}", Name);
+            return;
+        }
+        process.StandardInput.WriteLine(input);
+    }
+    
+    public virtual async Task SendInputAsync(string input)
+    {
+        var process = VenvRunner?.Process;
+        if (process == null)
+        {
+            Logger.Warn("No process running for {Name}", Name);
+            return;
+        }
+        await process.StandardInput.WriteLineAsync(input);
+    }
 
     public override Task Shutdown()
     {
diff --git a/StabilityMatrix.Core/Processes/AnsiProcess.cs b/StabilityMatrix.Core/Processes/AnsiProcess.cs
index d5f1783d..8738e7d8 100644
--- a/StabilityMatrix.Core/Processes/AnsiProcess.cs
+++ b/StabilityMatrix.Core/Processes/AnsiProcess.cs
@@ -1,4 +1,5 @@
 using System.Diagnostics;
+using System.Text;
 
 namespace StabilityMatrix.Core.Processes;
 
@@ -17,6 +18,11 @@ public class AnsiProcess : Process
         StartInfo.RedirectStandardOutput = true;
         StartInfo.RedirectStandardInput = true;
         StartInfo.RedirectStandardError = true;
+        
+        // Need this to parse ANSI escape sequences correctly
+        StartInfo.StandardOutputEncoding = Encoding.UTF8;
+        StartInfo.StandardErrorEncoding = Encoding.UTF8;
+        StartInfo.StandardInputEncoding = Encoding.UTF8;
     }
 
     /// <summary>
diff --git a/StabilityMatrix.Core/Processes/ApcMessage.cs b/StabilityMatrix.Core/Processes/ApcMessage.cs
new file mode 100644
index 00000000..7a528f9a
--- /dev/null
+++ b/StabilityMatrix.Core/Processes/ApcMessage.cs
@@ -0,0 +1,12 @@
+using System.Text.Json.Serialization;
+
+namespace StabilityMatrix.Core.Processes;
+
+public readonly struct ApcMessage
+{
+    [JsonPropertyName("type")]
+    public required ApcType Type { get; init; }
+    
+    [JsonPropertyName("data")]
+    public required string Data { get; init; }
+}
diff --git a/StabilityMatrix.Core/Processes/ApcParser.cs b/StabilityMatrix.Core/Processes/ApcParser.cs
new file mode 100644
index 00000000..28f471dd
--- /dev/null
+++ b/StabilityMatrix.Core/Processes/ApcParser.cs
@@ -0,0 +1,55 @@
+using System.Diagnostics;
+using System.Diagnostics.CodeAnalysis;
+using System.Text;
+using System.Text.Json;
+
+namespace StabilityMatrix.Core.Processes;
+
+/// <summary>
+/// Parse escaped messages from subprocess
+/// The message standard:
+///   - Message events are prefixed with char 'APC' (9F)
+///   - Followed by '[SM;'
+///   - Json dict string of 2 strings, 'type' and 'data'
+///   - Ends with char 'ST' (9C)
+/// </summary>
+[SuppressMessage("ReSharper", "MemberCanBePrivate.Global")]
+internal static class ApcParser
+{
+    public const char ApcEscape = (char) 0x9F;
+    public const string IdPrefix = "[SM;";
+    public const char StEscape = (char) 0x9C;
+    
+    /// <summary>
+    /// Attempts to extract an APC message from the given text
+    /// </summary>
+    /// <returns>ApcMessage struct</returns>
+    public static bool TryParse(string text, out ApcMessage? message)
+    {
+        message = null;
+        var startIndex = text.IndexOf(ApcEscape);
+        if (startIndex == -1) return false;
+        
+        // Check the IdPrefix follows the ApcEscape
+        var idIndex = text.IndexOf(IdPrefix, startIndex + 1, StringComparison.Ordinal);
+        if (idIndex == -1) return false;
+        
+        // Get the end index (ST escape)
+        var stIndex = text.IndexOf(StEscape, idIndex + IdPrefix.Length);
+        if (stIndex == -1) return false;
+
+        // Extract the json string (between idIndex and stIndex)
+        var json = text.Substring(idIndex + IdPrefix.Length, stIndex - idIndex - IdPrefix.Length);
+        
+        try
+        {
+            message = JsonSerializer.Deserialize<ApcMessage>(json);
+            return true;
+        }
+        catch (Exception e)
+        {
+            Debug.WriteLine($"Failed to parse APC message: {e.Message}");
+            return false;
+        }
+    }
+}
diff --git a/StabilityMatrix.Core/Processes/ApcType.cs b/StabilityMatrix.Core/Processes/ApcType.cs
new file mode 100644
index 00000000..dfe88c38
--- /dev/null
+++ b/StabilityMatrix.Core/Processes/ApcType.cs
@@ -0,0 +1,11 @@
+using System.Runtime.Serialization;
+using System.Text.Json.Serialization;
+
+namespace StabilityMatrix.Core.Processes;
+
+[JsonConverter(typeof(JsonStringEnumConverter))]
+public enum ApcType
+{
+    [EnumMember(Value = "input")]
+    Input = 1,
+}
diff --git a/StabilityMatrix.Core/Processes/AsyncStreamReader.cs b/StabilityMatrix.Core/Processes/AsyncStreamReader.cs
index 0379a72e..c440b1d3 100644
--- a/StabilityMatrix.Core/Processes/AsyncStreamReader.cs
+++ b/StabilityMatrix.Core/Processes/AsyncStreamReader.cs
@@ -4,6 +4,7 @@
 
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
+using System.Runtime.CompilerServices;
 using System.Runtime.ExceptionServices;
 using System.Text;
 
@@ -11,7 +12,9 @@ namespace StabilityMatrix.Core.Processes;
 
 /// <summary>
 /// Modified from System.Diagnostics.AsyncStreamReader to support progress bars,
-/// preserving '\r' instead of parsing as a line break. 
+/// preserving '\r' instead of parsing as a line break.
+/// This will also parse Apc escaped messages.
+/// <seealso cref="ApcParser"/>
 /// </summary>
 [SuppressMessage("ReSharper", "InconsistentNaming")]
 internal sealed class AsyncStreamReader : IDisposable
@@ -35,7 +38,6 @@ internal sealed class AsyncStreamReader : IDisposable
 
     // Cache the last position scanned in sb when searching for lines.
     private int _currentLinePos;
-
     // Creates a new AsyncStreamReader for the given stream. The
     // character encoding is set by encoding and the buffer size,
     // in number of 16-bit characters, is set by bufferSize.
@@ -131,6 +133,32 @@ internal sealed class AsyncStreamReader : IDisposable
 
         FlushMessageQueue(rethrowInNewThread: true);
     }
+    
+    // Send remaining buffer
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    private void SendRemainingBuffer()
+    {
+        lock (_messageQueue)
+        {
+            if (_sb!.Length == 0) return;
+            
+            _messageQueue.Enqueue(_sb.ToString());
+            _sb.Length = 0;
+        }
+    }
+    
+    // Send remaining buffer from index
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    private void SendRemainingBuffer(int startIndex)
+    {
+        lock (_messageQueue)
+        {
+            if (_sb!.Length == 0) return;
+            
+            _messageQueue.Enqueue(_sb.ToString(startIndex, _sb.Length - startIndex));
+            _sb.Length = 0;
+        }
+    }
 
     // Read lines stored in StringBuilder and the buffer we just read into.
     // A line is defined as a sequence of characters followed by
@@ -148,12 +176,7 @@ internal sealed class AsyncStreamReader : IDisposable
         // For progress bars
         if (len > 0 && _sb[0] == '\r' && (len == 1 || _sb[1] != '\n'))
         {
-            lock (_messageQueue)
-            {
-                _messageQueue.Enqueue(_sb.ToString());
-                _sb.Length = 0;
-            }
-
+            SendRemainingBuffer();
             return;
         }
 
@@ -211,6 +234,36 @@ internal sealed class AsyncStreamReader : IDisposable
                     // otherwise we ignore \r and treat it as normal char
                     break;
                 }
+                // Additional handling for Apc escape messages
+                case ApcParser.ApcEscape:
+                {
+                    // Unconditionally consume until StEscape
+                    // Look for index of StEscape
+                    var searchIndex = currentIndex;
+                    while (searchIndex < len && _sb[searchIndex] != ApcParser.StEscape)
+                    {
+                        searchIndex++;
+                    }
+                    
+                    // If we found StEscape, we have a complete APC message
+                    if (searchIndex < len)
+                    {
+                        // Include the StEscape as part of line.
+                        var line = _sb.ToString(lineStart, searchIndex - lineStart + 1);
+                        lock (_messageQueue)
+                        {
+                            _messageQueue.Enqueue(line);
+                        }
+                        // Advance currentIndex and lineStart to StEscape
+                        // lineStart = searchIndex + 1;
+                        currentIndex = searchIndex;
+                        // Also send the rest of the buffer immediately
+                        SendRemainingBuffer(currentIndex + 1);
+                        return;
+                    }
+                    // Otherwise continue without any other changes
+                    break;
+                }
             }
             currentIndex++;
         }
diff --git a/StabilityMatrix.Core/Processes/ProcessOutput.cs b/StabilityMatrix.Core/Processes/ProcessOutput.cs
index 031edb8e..b6b8fa69 100644
--- a/StabilityMatrix.Core/Processes/ProcessOutput.cs
+++ b/StabilityMatrix.Core/Processes/ProcessOutput.cs
@@ -3,14 +3,14 @@
 public readonly record struct ProcessOutput
 {
     /// <summary>
-    /// Raw output
+    /// Parsed text with escape sequences and line endings removed
     /// </summary>
-    public string RawText { get; init; }
-    
+    public required string Text { get; init; }
+
     /// <summary>
-    /// Parsed text with escape sequences and line endings removed
+    /// Raw output
     /// </summary>
-    public string Text { get; init; }
+    public string? RawText { get; init; }
     
     /// <summary>
     /// True if output from stderr, false for stdout
@@ -26,6 +26,11 @@ public readonly record struct ProcessOutput
     /// Instruction to clear last n lines
     /// </summary>
     public int ClearLines { get; init; }
+    
+    /// <summary>
+    /// Apc message sent from the subprocess
+    /// </summary>
+    public ApcMessage? ApcMessage { get; init; }
 
     public static ProcessOutput FromStdOutLine(string text)
     {
@@ -39,6 +44,18 @@ public readonly record struct ProcessOutput
 
     private static ProcessOutput FromLine(string text, bool isStdErr)
     {
+        // Parse APC message
+        if (ApcParser.TryParse(text, out var message))
+        {
+            // Override and return
+            return new ProcessOutput
+            {
+                RawText = text,
+                Text = text,
+                IsStdErr = isStdErr,
+                ApcMessage = message
+            };
+        }
         // If text contains newlines, split it first
         if (text.Contains(Environment.NewLine))
         {