[dotnet] [bidi] Properly handle websocket close handshake

dotnet/src/webdriver/BiDi/Broker.cs

@@ -113,7 +113,7 @@ public async ValueTask DisposeAsync()
 
         _receiveMessagesCancellationTokenSource.Dispose();
 
-        _transport.Dispose();
+        await _transport.DisposeAsync().ConfigureAwait(false);
 
         GC.SuppressFinalize(this);
     }
@@ -269,6 +269,15 @@ private async Task ReceiveMessagesAsync(CancellationToken cancellationToken)
                 _logger.Error($"Unhandled error occurred while receiving remote messages: {ex}");
             }
 
+            // Fail all pending commands, as the connection is likely broken if we failed to receive messages.
+            foreach (var id in _pendingCommands.Keys)
+            {
+                if (_pendingCommands.TryRemove(id, out var pendingCommand))
+                {
+                    pendingCommand.TaskCompletionSource.TrySetException(ex);
+                }
+            }
+
             throw;
         }
     }

dotnet/src/webdriver/BiDi/ITransport.cs

@@ -19,7 +19,7 @@
 
 namespace OpenQA.Selenium.BiDi;
 
-interface ITransport : IDisposable
+interface ITransport : IAsyncDisposable
 {
     Task<byte[]> ReceiveAsync(CancellationToken cancellationToken);
 

dotnet/src/webdriver/BiDi/WebSocketTransport.cs

@@ -24,7 +24,7 @@
 
 namespace OpenQA.Selenium.BiDi;
 
-sealed class WebSocketTransport(ClientWebSocket webSocket) : ITransport, IDisposable
+sealed class WebSocketTransport(ClientWebSocket webSocket) : ITransport
 {
     private readonly static ILogger _logger = Internal.Logging.Log.GetLogger<WebSocketTransport>();
 
@@ -67,6 +67,14 @@ public async Task<byte[]> ReceiveAsync(CancellationToken cancellationToken)
             {
                 result = await _webSocket.ReceiveAsync(segment, cancellationToken).ConfigureAwait(false);
 
+                if (result.MessageType == WebSocketMessageType.Close)
+                {
+                    await _webSocket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, string.Empty, CancellationToken.None).ConfigureAwait(false);
+
+                    throw new WebSocketException(WebSocketError.ConnectionClosedPrematurely,
+                        $"The remote end closed the WebSocket connection. Status: {result.CloseStatus}, Description: {result.CloseStatusDescription}");
+                }
+
                 _sharedMemoryStream.Write(receiveBuffer, 0, result.Count);
             }
             while (!result.EndOfMessage);
@@ -107,26 +115,31 @@ public async Task SendAsync(byte[] data, CancellationToken cancellationToken)
 
     private bool _disposed;
 
-    public void Dispose()
+    public async ValueTask DisposeAsync()
     {
-        Dispose(true);
-        GC.SuppressFinalize(this);
-    }
+        if (_disposed) return;
 
-    private void Dispose(bool disposing)
-    {
-        if (_disposed)
+        if (_webSocket.State == WebSocketState.Open)
         {
-            return;
+            try
+            {
+                await _webSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, string.Empty, CancellationToken.None).ConfigureAwait(false);
+            }
+            catch (Exception ex)
+            {
+                if (_logger.IsEnabled(LogEventLevel.Warn))
+                {
+                    _logger.Warn($"Error closing WebSocket gracefully: {ex.Message}");
+                }
+            }
         }
 
-        if (disposing)
-        {
-            _webSocket.Dispose();
-            _sharedMemoryStream.Dispose();
-            _socketSendSemaphoreSlim.Dispose();
-        }
+        _webSocket.Dispose();
+        _sharedMemoryStream.Dispose();
+        _socketSendSemaphoreSlim.Dispose();
 
         _disposed = true;
+
+        GC.SuppressFinalize(this);
     }
 }