【Azure APIM】在 Azure API Management 中配置 SSE Transport 的 MCP Server：从 App Service 到 APIM 的一次实践

using System.Collections.Concurrent;
using System.Text.Json;
using System.Text.Json.Nodes;
using System.Threading.Channels;
var builder = WebApplication.CreateBuilder(args);
builder.Services.AddCors(options =>
{
    options.AddDefaultPolicy(policy => policy
        .AllowAnyOrigin()
        .AllowAnyHeader()
        .AllowAnyMethod());
});
var app = builder.Build();
var sessions = new ConcurrentDictionary<string, SseClient>();
var jsonOptions = new JsonSerializerOptions(JsonSerializerDefaults.Web)
{
    WriteIndented = false
};
app.UseCors();
app.MapGet("/", () => Results.Json(new
{
    name = "MCP SSE Demo",
    endpoints = new
    {
        sse = "/sse",
        messages = "/messages?sessionId={sessionId}",
        hello = "/api/hello?name=World",
        add = "/api/add"
    }
}));
app.MapGet("/health", () => Results.Ok(new { status = "ok", time = DateTimeOffset.UtcNow }));
app.MapGet("/api/hello", (string? name) => Results.Ok(new
{
    message = $"Hello, {(string.IsNullOrWhiteSpace(name) ? "World" : name)}!",
    time = DateTimeOffset.UtcNow
}));
app.MapPost("/api/add", (AddRequest request) => Results.Ok(new
{
    request.A,
    request.B,
    sum = request.A + request.B
}));
app.MapGet("/sse", async (HttpContext context) =>
{
    var sessionId = Guid.NewGuid().ToString("N");
    var client = new SseClient(Channel.CreateUnbounded<string>());
    sessions[sessionId] = client;
    context.Response.Headers.CacheControl = "no-cache";
    context.Response.Headers.Connection = "keep-alive";
    context.Response.Headers.ContentType = "text/event-stream";
    // Use a relative message endpoint so reverse proxies/APIM can resolve it
    // against the public MCP SSE URL instead of leaking the backend host.
    var endpoint = $"/messages?sessionId={sessionId}";
        
    context.Response.Headers.CacheControl = "no-cache, no-transform";
    context.Response.Headers["X-Mcp-Endpoint-Source"] = "relative-v2";
    context.Response.Headers["X-Mcp-Endpoint-Value"] = endpoint;
    Console.WriteLine($"Writing MCP endpoint event: {endpoint}");
    await WriteSseAsync(context.Response, endpoint, "endpoint", context.RequestAborted);
    try
    {
        while (!context.RequestAborted.IsCancellationRequested)
        {
            var messageAvailable = client.Messages.Reader.WaitToReadAsync(context.RequestAborted).AsTask();
            var heartbeat = Task.Delay(TimeSpan.FromSeconds(15), context.RequestAborted);
            var completed = await Task.WhenAny(messageAvailable, heartbeat);
            if (completed == heartbeat)
            {
                await WriteSseAsync(context.Response, DateTimeOffset.UtcNow.ToString("O"), "ping", context.RequestAborted);
                continue;
            }
            if (!await messageAvailable)
            {
                break;
            }
            while (client.Messages.Reader.TryRead(out var payload))
            {
                await WriteSseAsync(context.Response, payload, "message", context.RequestAborted);
            }
        }
    }
    catch (OperationCanceledException)
    {
        // Client disconnected.
    }
    finally
    {
        sessions.TryRemove(sessionId, out _);
    }
});
app.MapPost("/messages", async (HttpContext context) =>
{
    var sessionId = context.Request.Query["sessionId"].ToString();
    if (string.IsNullOrWhiteSpace(sessionId) || !sessions.TryGetValue(sessionId, out var client))
    {
        return Results.NotFound(new { error = "Unknown or expired sessionId. Connect to /sse first." });
    }
    JsonNode? rpc;
    try
    {
        rpc = await JsonNode.ParseAsync(context.Request.Body, cancellationToken: context.RequestAborted);
    }
    catch (JsonException ex)
    {
        await client.Messages.Writer.WriteAsync(CreateError(null, -32700, $"Parse error: {ex.Message}"), context.RequestAborted);
        return Results.Accepted();
    }
    if (rpc is not JsonObject request)
    {
        await client.Messages.Writer.WriteAsync(CreateError(null, -32600, "Invalid JSON-RPC request."), context.RequestAborted);
        return Results.Accepted();
    }
    var response = HandleJsonRpc(request);
    if (response is not null)
    {
        await client.Messages.Writer.WriteAsync(response, context.RequestAborted);
    }
    return Results.Accepted();
});
app.Run();
string? HandleJsonRpc(JsonObject request)
{
    var id = request["id"];
    var method = request["method"]?.GetValue<string>();
    if (string.IsNullOrWhiteSpace(method))
    {
        return CreateError(id, -32600, "Missing JSON-RPC method.");
    }
    if (method.StartsWith("notifications/", StringComparison.Ordinal))
    {
        return null;
    }
    return method switch
    {
        "initialize" => CreateResult(id, new JsonObject
        {
            ["protocolVersion"] = "2024-11-05",
            ["capabilities"] = new JsonObject
            {
                ["tools"] = new JsonObject
                {
                    ["listChanged"] = false
                }
            },
            ["serverInfo"] = new JsonObject
            {
                ["name"] = "dotnet-mcp-sse-demo",
                ["version"] = "1.0.0"
            }
        }),
        "tools/list" => CreateResult(id, BuildToolsList()),
        "tools/call" => CreateResult(id, CallTool(request["params"] as JsonObject)),
        "resources/list" => CreateResult(id, new JsonObject { ["resources"] = new JsonArray() }),
        "prompts/list" => CreateResult(id, new JsonObject { ["prompts"] = new JsonArray() }),
        _ => CreateError(id, -32601, $"Method not found: {method}")
    };
}
JsonObject BuildToolsList() => new()
{
    ["tools"] = new JsonArray
    {
        new JsonObject
        {
            ["name"] = "echo",
            ["description"] = "Return the provided text.",
            ["inputSchema"] = new JsonObject
            {
                ["type"] = "object",
                ["properties"] = new JsonObject
                {
                    ["text"] = new JsonObject
                    {
                        ["type"] = "string",
                        ["description"] = "Text to echo."
                    }
                },
                ["required"] = new JsonArray("text")
            }
        },
        new JsonObject
        {
            ["name"] = "server_time",
            ["description"] = "Get the current server UTC time.",
            ["inputSchema"] = new JsonObject
            {
                ["type"] = "object",
                ["properties"] = new JsonObject()
            }
        },
        new JsonObject
        {
            ["name"] = "add",
            ["description"] = "Add two numbers.",
            ["inputSchema"] = new JsonObject
            {
                ["type"] = "object",
                ["properties"] = new JsonObject
                {
                    ["a"] = new JsonObject { ["type"] = "number" },
                    ["b"] = new JsonObject { ["type"] = "number" }
                },
                ["required"] = new JsonArray("a", "b")
            }
        }
    }
};
JsonObject CallTool(JsonObject? parameters)
{
    var name = parameters?["name"]?.GetValue<string>();
    var arguments = parameters?["arguments"] as JsonObject ?? new JsonObject();
    var text = name switch
    {
        "echo" => arguments["text"]?.GetValue<string>() ?? string.Empty,
        "server_time" => DateTimeOffset.UtcNow.ToString("O"),
        "add" => Add(arguments),
        _ => $"Unknown tool: {name}"
    };
    return new JsonObject
    {
        ["content"] = new JsonArray
        {
            new JsonObject
            {
                ["type"] = "text",
                ["text"] = text
            }
        }
    };
}
static string Add(JsonObject arguments)
{
    var a = arguments["a"]?.GetValue<double>() ?? 0;
    var b = arguments["b"]?.GetValue<double>() ?? 0;
    return $"{a} + {b} = {a + b}";
}
string CreateResult(JsonNode? id, JsonNode result)
{
    return new JsonObject
    {
        ["jsonrpc"] = "2.0",
        ["id"] = CloneNode(id),
        ["result"] = result
    }.ToJsonString(jsonOptions);
}
string CreateError(JsonNode? id, int code, string message)
{
    return new JsonObject
    {
        ["jsonrpc"] = "2.0",
        ["id"] = CloneNode(id),
        ["error"] = new JsonObject
        {
            ["code"] = code,
            ["message"] = message
        }
    }.ToJsonString(jsonOptions);
}
static JsonNode? CloneNode(JsonNode? node) => node is null ? null : JsonNode.Parse(node.ToJsonString());
static async Task WriteSseAsync(HttpResponse response, string data, string? eventName, CancellationToken cancellationToken)
{
    if (!string.IsNullOrWhiteSpace(eventName))
    {
        await response.WriteAsync($"event: {eventName}\n", cancellationToken);
    }
    var lines = data.Replace("\r\n", "\n", StringComparison.Ordinal).Split('\n');
    foreach (var line in lines)
    {
        await response.WriteAsync($"data: {line}\n", cancellationToken);
    }
    await response.WriteAsync("\n", cancellationToken);
    await response.Body.FlushAsync(cancellationToken);
}
public sealed record AddRequest(double A, double B);
public sealed record SseClient(Channel<string> Messages);