Alchegos.MCP/Tools/TerminalTool.cs

using System.Collections.Concurrent;
using System.ComponentModel;
using System.Diagnostics;
using System.Text;
using System.Text.RegularExpressions;
using ModelContextProtocol.Server;

namespace Alchegos.MCP.Tools;

[McpServerToolType, Description("Manage multiple persistent terminal sessions, enabling creation, closure, command execution, and output retrieval.")]
public static class TerminalTool
{
    private static readonly ConcurrentDictionary<string, TerminalSession> _sessions = new();

    public record SessionInfo(
        [property: Description("Unique identifier for the session.")] string SessionId,
        [property: Description("Human-readable label for the session.")] string Label
    );

    public record ListSessionsResult(
        [property: Description("List of all active sessions with their IDs and labels.")] List<SessionInfo> Sessions
    );

    public record CreateSessionResult(
        [property: Description("Status of the operation: 'success' or 'fail'.")] string Status,
        [property: Description("Details of the created session if successful; otherwise, with empty fields.")] SessionInfo Session
    );

    public record ExitSessionResult(
        [property: Description("Status of the operation: 'success' or 'fail'.")] string Status
    );

    public record SendSignalResult(
        [property: Description("Status of the operation: 'success' or 'fail'.")] string Status
    );

    public record SendInputWithExceptResult(
        [property: Description("Indicates how the response was triggered: 'except' when a regex matched, or 'timeout' when the timeout was reached.")] string OutputBy,
        [property: Description("The regex pattern that matched, or null if none matched before timeout.")] string? MatchedExcept,
        [property: Description("Accumulated standard output from the session after sending input.")] string Stdout,
        [property: Description("Accumulated standard error from the session after sending input.")] string Stderr
    );

    public record SendInputAndForgetResult(
        [property: Description("Status of the operation: 'success' or 'fail'.")] string Status
    );

    public record OneOffCommandResult(
        [property: Description("How the command completed: 'finish' when the process exited normally, or 'timeout' when the timeout was reached.")] string ReturnBy,
        [property: Description("Standard output produced by the command.")] string Stdout,
        [property: Description("Standard error produced by the command.")] string Stderr
    );

    public record GetCachedResultResult(
        [property: Description("Cached standard output of the session.")] string Stdout,
        [property: Description("Cached standard error of the session.")] string Stderr
    );

    public record ClearResult(
        [property: Description("Status of the operation: 'success' or 'fail'.")] string Status
    );

    private class TerminalSession : IDisposable
    {
        public string Id { get; }
        public string Label { get; }
        public Process Process { get; }
        
        private readonly StringBuilder _stdoutBuffer = new();
        private readonly StringBuilder _stderrBuffer = new();
        private readonly object _lockObject = new();
        
        private bool _isDisposed;

        public TerminalSession(string id, string label)
        {
            Id = id;
            Label = label;
            
            Process = new Process
            {
                StartInfo = new ProcessStartInfo
                {
                    FileName = GetShellPath(),
                    Arguments = GetShellArguments(),
                    RedirectStandardInput = true,
                    RedirectStandardOutput = true,
                    RedirectStandardError = true,
                    UseShellExecute = false,
                    CreateNoWindow = true
                },
                EnableRaisingEvents = true
            };
            
            Process.OutputDataReceived += (sender, e) =>
            {
                if (e.Data != null)
                {
                    lock (_lockObject)
                    {
                        _stdoutBuffer.AppendLine(e.Data);
                    }
                }
            };
            
            Process.ErrorDataReceived += (sender, e) =>
            {
                if (e.Data != null)
                {
                    lock (_lockObject)
                    {
                        _stderrBuffer.AppendLine(e.Data);
                    }
                }
            };
            
            Process.Start();
            Process.BeginOutputReadLine();
            Process.BeginErrorReadLine();
        }

        private static string GetShellPath()
        {
            if (OperatingSystem.IsWindows())
                return "cmd.exe";
            return "/bin/bash";
        }

        private static string GetShellArguments()
        {
            if (OperatingSystem.IsWindows())
                return "/q";
            return "-i";
        }

        public void SendInput(string input)
        {
            if (!Process.HasExited)
            {
                Process.StandardInput.WriteLine(input);
                Process.StandardInput.Flush();
            }
        }

        public bool SendSignal(int signal)
        {
            try
            {
                if (OperatingSystem.IsWindows())
                {
                    if (signal == 2)
                        return NativeMethods.SendCtrlC(Process);
                    return false;
                }
                return NativeMethods.SendSignal(Process.Id, signal);
            }
            catch
            {
                return false;
            }
        }

        public (string stdout, string stderr) GetAndClearBuffers()
        {
            lock (_lockObject)
            {
                string stdout = _stdoutBuffer.ToString();
                string stderr = _stderrBuffer.ToString();
                
                _stdoutBuffer.Clear();
                _stderrBuffer.Clear();
                
                return (stdout, stderr);
            }
        }

        public (string stdout, string stderr) GetBuffers()
        {
            lock (_lockObject)
            {
                return (_stdoutBuffer.ToString(), _stderrBuffer.ToString());
            }
        }

        public void ClearBuffers()
        {
            lock (_lockObject)
            {
                _stdoutBuffer.Clear();
                _stderrBuffer.Clear();
            }
        }

        public void Dispose()
        {
            if (_isDisposed)
                return;

            _isDisposed = true;
            
            try
            {
                if (!Process.HasExited)
                {
                    Process.StandardInput.WriteLine("exit");
                    Process.StandardInput.Flush();
                    
                    if (!Process.WaitForExit(1000))
                        Process.Kill();
                }
            }
            catch
            {
                // Best effort only
            }
            finally
            {
                Process.Dispose();
            }
        }
    }

    private static class NativeMethods
    {
        public static bool SendCtrlC(Process process)
        {
            try
            {
                if (process.HasExited)
                    return false;
                
                if (!AttachConsole((uint)process.Id))
                    return false;
                
                SetConsoleCtrlHandler(null, true);
                try
                {
                    if (!GenerateConsoleCtrlEvent(CtrlTypes.CTRL_C_EVENT, 0))
                        return false;
                    return true;
                }
                finally
                {
                    SetConsoleCtrlHandler(null, false);
                    FreeConsole();
                }
            }
            catch
            {
                return false;
            }
        }

        public static bool SendSignal(int pid, int signal)
        {
            try
            {
                if (OperatingSystem.IsLinux() || OperatingSystem.IsMacOS())
                {
                    var killProcess = Process.Start(new ProcessStartInfo
                    {
                        FileName = "kill",
                        Arguments = $"-{signal} {pid}",
                        UseShellExecute = false,
                        CreateNoWindow = true
                    });
                    
                    killProcess?.WaitForExit();
                    return killProcess?.ExitCode == 0;
                }
                
                return false;
            }
            catch
            {
                return false;
            }
        }

        // Windows Console API imports
        [System.Runtime.InteropServices.DllImport("kernel32.dll")]
        public static extern bool AttachConsole(uint dwProcessId);

        [System.Runtime.InteropServices.DllImport("kernel32.dll")]
        public static extern bool FreeConsole();

        [System.Runtime.InteropServices.DllImport("kernel32.dll")]
        public static extern bool SetConsoleCtrlHandler(ConsoleCtrlHandler? HandlerRoutine, bool Add);

        [System.Runtime.InteropServices.DllImport("kernel32.dll")]
        public static extern bool GenerateConsoleCtrlEvent(CtrlTypes dwCtrlEvent, uint dwProcessGroupId);

        public enum CtrlTypes
        {
            CTRL_C_EVENT = 0,
            CTRL_BREAK_EVENT = 1,
            CTRL_CLOSE_EVENT = 2,
            CTRL_LOGOFF_EVENT = 5,
            CTRL_SHUTDOWN_EVENT = 6
        }

        public delegate bool ConsoleCtrlHandler(CtrlTypes ctrlType);
    }

    //——— Tool Methods ———//

    [McpServerTool, Description("""
        List all active terminal sessions.
        Return Format:
        {
            sessions: [ { session_id: string, label: string }, ... ]
        }
        """)]
    public static ListSessionsResult ListSessions()
    {
        var sessions = _sessions.Values.Select(s => new SessionInfo(s.Id, s.Label)).ToList();
        return new ListSessionsResult(sessions);
    }

    [McpServerTool, Description("""
        Create a new terminal session.
        Return Format:
        {
            status: 'success' | 'fail',
            session: { session_id: string, label: string }
        }
        """)]
    public static CreateSessionResult CreateSession(
        [Description("Memorable label for the new session.")] string label
    )
    {
        try
        {
            string sessionId = Guid.NewGuid().ToString("N");
            var session = new TerminalSession(sessionId, label);
            
            if (_sessions.TryAdd(sessionId, session))
            {
                return new CreateSessionResult("success", new SessionInfo(sessionId, label));
            }
            
            session.Dispose();
            return new CreateSessionResult("fail", new SessionInfo(string.Empty, string.Empty));
        }
        catch (Exception)
        {
            return new CreateSessionResult("fail", new SessionInfo(string.Empty, string.Empty));
        }
    }

    [McpServerTool, Description("""
        Exit a terminal session. Fails if the session ID does not exist.
        Return Format:
        {
            status: 'success' | 'fail'
        }
        """)]
    public static ExitSessionResult ExitSession(
        [Description("ID of the session to exit.")] string sessionId
    )
    {
        if (_sessions.TryRemove(sessionId, out var session))
        {
            session.Dispose();
            return new ExitSessionResult("success");
        }
        
        return new ExitSessionResult("fail");
    }

    [McpServerTool, Description("""
        Send a signal to a terminal session (e.g., SIGINT=2).
        Return Format:
        {
            status: 'success' | 'fail'
        }
        """)]
    public static SendSignalResult SendSignal(
        [Description("Session ID to send the signal to.")] string sessionId,
        [Description("Signal number to send (e.g., 2 for SIGINT)." )] int signal
    )
    {
        if (_sessions.TryGetValue(sessionId, out var session))
        {
            bool success = session.SendSignal(signal);
            return new SendSignalResult(success ? "success" : "fail");
        }
        
        return new SendSignalResult("fail");
    }

    [McpServerTool, Description("""
        Send input to a terminal session and wait until the output matches one of the exception regexes or timeout.
        Use cases: interactive prompts (e.g., confirmation prompts, password input).
        Return Format:
        {
            output_by: 'except' | 'timeout',
            matched_except: string | null,
            stdout: string,
            stderr: string
        }
        Note: cached stdout/stderr is cleared after this call.
        """)]
    public static SendInputWithExceptResult SendInputWithExcept(
        [Description("Session ID to send input to.")] string sessionId,
        [Description("Input text to send (without newline)." )] string input,
        [Description("Timeout in milliseconds.")] int timeout,
        [Description("Regex patterns to match against output.")] params string[] excepts
    )
    {
        if (!_sessions.TryGetValue(sessionId, out var session))
        {
            return new SendInputWithExceptResult("timeout", null, string.Empty, string.Empty);
        }
        
        session.SendInput(input);
        
        var regexPatterns = excepts.Select(pattern => new Regex(pattern, RegexOptions.Compiled)).ToArray();
        
        string matchedPattern = null;
        bool matched = false;
        var watch = Stopwatch.StartNew();
        
        while (watch.ElapsedMilliseconds < timeout && !matched)
        {
            var (stdout, stderr) = session.GetBuffers();
            string combinedOutput = stdout + stderr;
            
            for (int i = 0; i < regexPatterns.Length; i++)
            {
                if (regexPatterns[i].IsMatch(combinedOutput))
                {
                    matchedPattern = excepts[i];
                    matched = true;
                    break;
                }
            }
            
            if (!matched)
            {
                Thread.Sleep(10);
            }
        }
        
        watch.Stop();
        
        var (finalStdout, finalStderr) = session.GetAndClearBuffers();
        
        return new SendInputWithExceptResult(
            matched ? "except" : "timeout",
            matchedPattern,
            finalStdout,
            finalStderr
        );
    }

    [McpServerTool, Description("""
        Send input to a terminal session without waiting for output. Output is cached for later retrieval.
        Use cases: long-running or background commands. e.g. ssh
        Return Format:
        {
            status: 'success' | 'fail'
        }
        """)]
    public static SendInputAndForgetResult SendInputAndForget(
        [Description("Session ID to send input to.")] string sessionId,
        [Description("Input text to send (without newline)." )] string input
    )
    {
        if (_sessions.TryGetValue(sessionId, out var session))
        {
            try
            {
                session.SendInput(input);
                return new SendInputAndForgetResult("success");
            }
            catch
            {
                return new SendInputAndForgetResult("fail");
            }
        }
        
        return new SendInputAndForgetResult("fail");
    }

    [McpServerTool, Description("""
        Execute a one-off command in the specified session and return stdout/stderr when done or timed out.
        Use cases: quick commands that complete shortly (e.g., ls, echo).
        Return Format:
        {
            return_by: 'finish' | 'timeout',
            stdout: string,
            stderr: string
        }
        """)]
    public static OneOffCommandResult OneOffCommand(
        [Description("Session ID to execute the command in.")] string sessionId,
        [Description("Command to execute, including arguments.")] string command,
        [Description("Timeout in milliseconds.")] int timeout
    )
    {
        if (!_sessions.TryGetValue(sessionId, out var session))
        {
            return new OneOffCommandResult("timeout", string.Empty, string.Empty);
        }
        
        session.ClearBuffers();
        
        string marker = $"CMD_COMPLETE_{Guid.NewGuid():N}";
        
        if (OperatingSystem.IsWindows())
        {
            session.SendInput($"{command} & echo {marker}");
        }
        else
        {
            session.SendInput($"{command}; echo {marker}");
        }
       
        bool completed = false;
        var watch = Stopwatch.StartNew();
        
        while (watch.ElapsedMilliseconds < timeout && !completed)
        {
            var (stdout, _) = session.GetBuffers();
            
            if (stdout.Contains(marker))
            {
                completed = true;
            }
            else
            {
                Thread.Sleep(10);
            }
        }
        
        watch.Stop();
        
        var (finalStdout, finalStderr) = session.GetAndClearBuffers();
        
        if (completed)
        {
            finalStdout = Regex.Replace(finalStdout, $"{marker}\r?\n?$", "");
        }
        
        return new OneOffCommandResult(
            completed ? "finish" : "timeout",
            finalStdout,
            finalStderr
        );
    }

    [McpServerTool, Description("""
        Get and clear cached stdout/stderr of the specified session.
        Return Format:
        {
            stdout: string,
            stderr: string
        }
        """)]
    public static GetCachedResultResult GetCachedResult(
        [Description("Session ID to retrieve output from.")] string sessionId
    )
    {
        if (_sessions.TryGetValue(sessionId, out var session))
        {
            var (stdout, stderr) = session.GetAndClearBuffers();
            return new GetCachedResultResult(stdout, stderr);
        }
        
        return new GetCachedResultResult(string.Empty, string.Empty);
    }

    [McpServerTool, Description("""
        Clear cached stdout/stderr of the specified session.
        Return Format:
        {
            status: 'success' | 'fail'
        }
        """)]
    public static ClearResult Clear(
        [Description("Session ID to clear cache for.")] string sessionId
    )
    {
        if (_sessions.TryGetValue(sessionId, out var session))
        {
            session.ClearBuffers();
            return new ClearResult("success");
        }
        
        return new ClearResult("fail");
    }
}