[CmdletBinding()]
[OutputType([string])]
param (
    [Parameter(Mandatory = $true, Position = 0)]
    [System.Net.WebSockets.ClientWebSocket]
    $WebSocket
)

begin {
    $cts = [System.Threading.CancellationTokenSource]::new();

    if ($WebSocket.State -ne [System.Net.WebSockets.WebSocketState]::Open) {
        Write-Error "Can only receive from Open WebSockets, was $($WebSocket.State)";
        $abort = $true;
    }
}

process {
    if ($abort) {
        return;
    }

    $endOfMessage = $false;
    [string]$message = "";
    $memory = [System.Memory[byte]]::new([byte[]]::new(1024));
    while (-not $endOfMessage) {
        $task = $WebSocket.ReceiveAsync($memory, $cts.Token);
        while ($null -eq $result) {
            if ($task.IsCompletedSuccessfully) {
                $result = $task.Result
            }
            elseif ($task.IsFaulted) {
                Write-Error "Receiving faulted";
                return;
            }
            elseif ($task.IsCanceled) {
                return;
            }
        }
        $part = [System.Text.Encoding]::UTF8.GetString($memory.Slice(0, $result.Count).ToArray());
        $message += $part;
        $endOfMessage = $result.EndOfMessage;
        $result = $null;
    }
    return $message;
}

clean {
    $cts.Cancel();
}