SelfHost Websocket - Web API vs OWIN Middleware

Question

I'm running an OWIN self-hosted application that hosts a REST API and also allows websocket connectivity for real-time data. I'm using WebAPI to handle the routing and mapping of routes to controllers.

When I use Web API to handle the websocket routes the socket is closed as soon as the controller returns. However, if I create my own middleware the socket does not close.

I'd prefer to use Web API for all of my routes. But more importantly I want to understand what's going on. I don't like my production code to work without understanding why it's working.

Here is the relevant Web API code snippet:

public class WebServer : IDisposable
{
    public void Configuration(IAppBuilder app)
    {
        HttpConfiguration config = new HttpConfiguration();

        config.Routes.MapHttpRoute(
            name: "Websocket",
            routeTemplate: "ws/all",
            defaults: new { controller = "MyWebSocket", action = "Get" });
        app.UseWebApi(config);
    }
}

public class MyWebSocketController : System.Web.Http.ApiController
{
    public IHttpActionResult Get()
    {
        var owinContext = Request.GetOwinContext();

        var accept = owinContext.Get<Action<IDictionary<string, object>, Func<IDictionary<string, object>, Task>>>("websocket.Accept");

        accept(null, RunWebSocket);
        return Ok();
    }

    private async Task RunWebSocket(IDictionary<string, object> websocketContext)
    {
        WebSocket socket;
        if (websocketContext.TryGetValue(typeof(System.Net.WebSockets.WebSocketContext).FullName, out value))
        {
            socket = ((System.Net.WebSockets.WebSocketContext)value).WebSocket;
        }

        ArraySegment<Byte> buffer = new ArraySegment<byte>(new Byte[128]);
        WebSocketReceiveResult result = null;

        using (var ms = new MemoryStream())
        {
            while (socket.State == WebSocketState.Open)
            {
                ms.SetLength(0);
                do
                {
                    result = await socket.ReceiveAsync(buffer, CancellationToken.None);
                    ms.Write(buffer.Array, buffer.Offset, result.Count);
                }
                while (!result.EndOfMessage);

                if (result.MessageType == WebSocketMessageType.Close)
                {
                    // Close socket
                }
                ms.Seek(0, SeekOrigin.Begin);

                if (result.MessageType == WebSocketMessageType.Text)
                {
                    // Handle message
                }
            }
        }
    }
}

The call to ReceiveAsync schedules a continuation. The Get method returns back to the ApiController which closes the connection, which also closes the websocket.

Here is the relevant code for the OWIN middleware.

public class WebServer : IDisposable
{
    public void Configuration(IAppBuilder app)
    {
        app.Use<WebSocketMiddleware>();
    }
}

public class WebSocketMiddleware : OwinMiddleware
{
    public override Task Invoke(IOwinContext context)
    {
        var accept = context.Get<Action<IDictionary<string, object>, Func<IDictionary<string, object>, Task>>>("websocket.Accept");

        accept(null, RunWebSocket);
        return;
    }

    private async Task RunWebSocket(IDictionary<string, object> websocketContext)
    {
        // Same as Web API implementation
    }
}

Again the continuation is scheduled during the call to ReceiveAsync and the Invoke method returns. However, the connection remains open and I'm able to send and receive data through the websocket.

So, I have a solution, but I'd really love to understand what's going on. Any references would be greatly appreciated.

Edit: Actually the socket is closed in both cases. The web API version sends a RST from the server, as if the connection was abruptly closed, while the OWIN version experiences a normal FIN ACK. However the web API doesn't allow any further communication over the websocket, while the OWIN version does. So I'm not really sure how this is supposed to work.