Playa/tests/Volcanic.AgenticCore.Extensions.Tests.FakeMcpServer/FakeServerProgram.cs

using System.Collections;
using System.Text;
using System.Text.Json;
using ModelContextProtocol.Protocol;
using ModelContextProtocol.Server;

namespace Volcanic.AgenticCore.Extensions.Tests.FakeMcpServer;

/// <summary>Test MCP stdio server. <c>FAKE_MCP_MODE</c> (comma-separated) selects behaviour; <c>FAKE_MCP_LOG</c> receives JSON lines.</summary>
public static class FakeServerProgram
{
    public const string LongToolName = "a_tool_whose_name_is_much_longer_than_forty_eight_characters";
    public const string Png1x1 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=";
    public const string EchoSchema = """{"type":"object","properties":{"text":{"type":"string","description":"Text to echo"}},"required":["text"]}""";
    public const int DieExitCode = 3;

    public static async Task<int> Main(string[] args)
    {
        var modes = (Environment.GetEnvironmentVariable("FAKE_MCP_MODE") is { Length: > 0 } m ? m : "normal")
            .Split(',', StringSplitOptions.RemoveEmptyEntries | StringSplitOptions.TrimEntries)
            .ToHashSet(StringComparer.Ordinal);
        var log = Environment.GetEnvironmentVariable("FAKE_MCP_LOG");
        if (modes.Contains("count-starts"))
            Log(log, new { @event = "start", pid = Environment.ProcessId, cwd = Directory.GetCurrentDirectory(), args, env = modes.Contains("env") ? Snapshot() : null });

        if (modes.Contains("crash-on-start"))
        {
            await Console.Error.WriteLineAsync("Traceback (most recent call last):");
            await Console.Error.WriteLineAsync("  File \"fake_server.py\", line 1, in <module>");
            await Console.Error.WriteLineAsync("RuntimeError: fake crash on start");
            return 1;
        }
        if (modes.Contains("hang-on-start"))
        {
            await Console.In.ReadToEndAsync(); // swallow initialize, never answer; exit once the client closes stdin
            return 0;
        }

        var playwrightLike = args.Contains("--output-dir");
        var calls = 0;
        var options = new McpServerOptions
        {
            ServerInfo = new Implementation { Name = "fake-mcp", Version = "1.2.3" },
            Capabilities = new ServerCapabilities { Tools = new ToolsCapability() },
            Handlers = new McpServerHandlers
            {
                ListToolsHandler = (_, _) => ValueTask.FromResult(new ListToolsResult { Tools = Tools(modes, playwrightLike) }),
                CallToolHandler = async (request, ct) =>
                {
                    var result = await CallAsync(request.Params!, log, ct);
                    if (modes.Contains("die-after-first-call") && Interlocked.Increment(ref calls) == 1) _ = DieSoonAsync();
                    return result;
                },
            },
        };
        await using var server = McpServer.Create(new StdioServerTransport(options), options);
        await server.RunAsync();
        return 0;
    }

    private static List<Tool> Tools(HashSet<string> modes, bool playwrightLike)
    {
        var tools = new List<Tool>
        {
            Make("echo", "Echoes text back", EchoSchema),
            Make("image", "Returns a 1x1 PNG"),
            Make("fail", "Always fails"),
            Make("slow", "Waits until cancelled"),
            Make("weird.name/x", "Has an awkward name"),
            Make("weird_name_x", "Collides with weird.name/x after sanitising"),
            new() { Name = "noschema", Description = "Has no input schema" },
            Make(LongToolName, "Has a long name"),
        };
        if (modes.Contains("env")) tools.Add(Make("env", "Returns the server's environment"));
        if (modes.Contains("builtin-clash")) tools.Add(Make("read", "Clashes with the built-in read tool"));
        if (playwrightLike)
        {
            tools.Add(Make("browser_navigate", "Navigate to a URL", """{"type":"object","properties":{"url":{"type":"string"}},"required":["url"]}"""));
            tools.Add(Make("browser_take_screenshot", "Take a screenshot", """{"type":"object","properties":{"filename":{"type":"string"}}}"""));
        }
        return tools;
    }

    private static Tool Make(string name, string description, string schema = """{"type":"object","properties":{}}""") =>
        new() { Name = name, Description = description, InputSchema = JsonSerializer.Deserialize<JsonElement>(schema) };

    private static async ValueTask<CallToolResult> CallAsync(CallToolRequestParams call, string? log, CancellationToken ct)
    {
        string Arg(string name) =>
            call.Arguments is { } a && a.TryGetValue(name, out var v) && v.ValueKind == JsonValueKind.String ? v.GetString()! : "";
        switch (call.Name)
        {
            case "echo":
                return Text($"echo: {Arg("text")}");
            case "image":
                return new CallToolResult { Content = [Image()] };
            case "fail":
                return new CallToolResult { IsError = true, Content = [new TextContentBlock { Text = "fake failure: nothing works" }, Image()] };
            case "slow":
                try { await Task.Delay(Timeout.Infinite, ct); }
                catch (OperationCanceledException) { Log(log, new { @event = "cancelled", tool = "slow" }); throw; }
                return Text("slow finished");
            case "env":
                return Text(string.Join('\n', Snapshot().OrderBy(kv => kv.Key, StringComparer.Ordinal).Select(kv => $"{kv.Key}={kv.Value}")));
            case "browser_navigate":
                return Text($"navigated to {Arg("url")}");
            case "browser_take_screenshot":
                var filename = Arg("filename");
                if (filename.Length == 0) return new CallToolResult { Content = [Image()] };
                var path = Path.GetFullPath(filename, Directory.GetCurrentDirectory());
                Directory.CreateDirectory(Path.GetDirectoryName(path)!);
                await File.WriteAllBytesAsync(path, Convert.FromBase64String(Png1x1), ct);
                return Text($"saved {filename}");
            default:
                return Text($"called {call.Name}");
        }
    }

    private static CallToolResult Text(string text) => new() { Content = [new TextContentBlock { Text = text }] };

    private static ImageContentBlock Image() => new() { Data = Encoding.UTF8.GetBytes(Png1x1), MimeType = "image/png" };

    private static Dictionary<string, string?> Snapshot()
    {
        var env = new Dictionary<string, string?>(StringComparer.Ordinal);
        foreach (DictionaryEntry e in Environment.GetEnvironmentVariables()) env[(string)e.Key] = (string?)e.Value;
        return env;
    }

    private static async Task DieSoonAsync()
    {
        await Task.Delay(200);
        await Console.Error.WriteLineAsync($"dying: {Environment.GetEnvironmentVariable("FAKE_MCP_DIE_MESSAGE")}");
        await Console.Error.FlushAsync();
        Environment.Exit(DieExitCode);
    }

    private static void Log(string? path, object entry)
    {
        if (string.IsNullOrEmpty(path)) return;
        var bytes = Encoding.UTF8.GetBytes(JsonSerializer.Serialize(entry) + "\n");
        for (var attempt = 0; ; attempt++)
        {
            try
            {
                using var stream = new FileStream(path, FileMode.Append, FileAccess.Write, FileShare.ReadWrite);
                stream.Write(bytes);
                return;
            }
            catch (IOException) when (attempt < 20) { Thread.Sleep(25); }
        }
    }
}