[dotnet] [bidi] Modules as extensions

dotnet/src/webdriver/BiDi/BiDi.cs

@@ -55,7 +55,7 @@ internal Session.SessionModule SessionModule
             if (_sessionModule is not null) return _sessionModule;
             lock (_moduleLock)
             {
-                _sessionModule ??= new Session.SessionModule(_broker);
+                _sessionModule ??= Module.Create<Session.SessionModule>(_broker);
             }
             return _sessionModule;
         }
@@ -68,7 +68,7 @@ public BrowsingContext.BrowsingContextModule BrowsingContext
             if (_browsingContextModule is not null) return _browsingContextModule;
             lock (_moduleLock)
             {
-                _browsingContextModule ??= new BrowsingContext.BrowsingContextModule(_broker);
+                _browsingContextModule ??= Module.Create<BrowsingContext.BrowsingContextModule>(_broker);
             }
             return _browsingContextModule;
         }
@@ -81,7 +81,7 @@ public Browser.BrowserModule Browser
             if (_browserModule is not null) return _browserModule;
             lock (_moduleLock)
             {
-                _browserModule ??= new Browser.BrowserModule(_broker);
+                _browserModule ??= Module.Create<Browser.BrowserModule>(_broker);
             }
             return _browserModule;
         }
@@ -94,7 +94,7 @@ public Network.NetworkModule Network
             if (_networkModule is not null) return _networkModule;
             lock (_moduleLock)
             {
-                _networkModule ??= new Network.NetworkModule(_broker);
+                _networkModule ??= Module.Create<Network.NetworkModule>(_broker);
             }
             return _networkModule;
         }
@@ -107,7 +107,7 @@ internal Input.InputModule InputModule
             if (_inputModule is not null) return _inputModule;
             lock (_moduleLock)
             {
-                _inputModule ??= new Input.InputModule(_broker);
+                _inputModule ??= Module.Create<Input.InputModule>(_broker);
             }
             return _inputModule;
         }
@@ -120,7 +120,7 @@ public Script.ScriptModule Script
             if (_scriptModule is not null) return _scriptModule;
             lock (_moduleLock)
             {
-                _scriptModule ??= new Script.ScriptModule(_broker);
+                _scriptModule ??= Module.Create<Script.ScriptModule>(_broker);
             }
             return _scriptModule;
         }
@@ -133,7 +133,7 @@ public Log.LogModule Log
             if (_logModule is not null) return _logModule;
             lock (_moduleLock)
             {
-                _logModule ??= new Log.LogModule(_broker);
+                _logModule ??= Module.Create<Log.LogModule>(_broker);
             }
             return _logModule;
         }
@@ -146,7 +146,7 @@ public Storage.StorageModule Storage
             if (_storageModule is not null) return _storageModule;
             lock (_moduleLock)
             {
-                _storageModule ??= new Storage.StorageModule(_broker);
+                _storageModule ??= Module.Create<Storage.StorageModule>(_broker);
             }
             return _storageModule;
         }
@@ -159,7 +159,7 @@ public WebExtension.WebExtensionModule WebExtension
             if (_webExtensionModule is not null) return _webExtensionModule;
             lock (_moduleLock)
             {
-                _webExtensionModule ??= new WebExtension.WebExtensionModule(_broker);
+                _webExtensionModule ??= Module.Create<WebExtension.WebExtensionModule>(_broker);
             }
             return _webExtensionModule;
         }
@@ -172,7 +172,7 @@ public Emulation.EmulationModule Emulation
             if (_emulationModule is not null) return _emulationModule;
             lock (_moduleLock)
             {
-                _emulationModule ??= new Emulation.EmulationModule(_broker);
+                _emulationModule ??= Module.Create<Emulation.EmulationModule>(_broker);
             }
             return _emulationModule;
         }

dotnet/src/webdriver/BiDi/Browser/BrowserModule.cs

@@ -17,12 +17,14 @@
 // under the License.
 // </copyright>
 
-using System.Threading.Tasks;
 using OpenQA.Selenium.BiDi.Communication;
+using System.Collections.Generic;
+using System.Text.Json.Serialization;
+using System.Threading.Tasks;
 
 namespace OpenQA.Selenium.BiDi.Browser;
 
-public sealed class BrowserModule(Broker broker) : Module(broker)
+public sealed class BrowserModule : Module
 {
     public async Task<EmptyResult> CloseAsync(CloseOptions? options = null)
     {
@@ -73,4 +75,22 @@ public async Task<EmptyResult> SetDownloadBehaviorDeniedAsync(SetDownloadBehavio
 
         return await Broker.ExecuteCommandAsync<SetDownloadBehaviorCommand, EmptyResult>(new SetDownloadBehaviorCommand(@params), options).ConfigureAwait(false);
     }
+
+    protected internal override void Initialize(Broker broker)
+    {
+        broker.ConfigureJsonContext(opts => opts.TypeInfoResolverChain.Add(BrowserModuleJsonSerializerContext.Default));
+    }
 }
+
+[JsonSerializable(typeof(CloseCommand))]
+[JsonSerializable(typeof(CreateUserContextCommand))]
+[JsonSerializable(typeof(GetUserContextsCommand))]
+[JsonSerializable(typeof(GetUserContextsResult))]
+[JsonSerializable(typeof(RemoveUserContextCommand))]
+[JsonSerializable(typeof(GetClientWindowsCommand))]
+[JsonSerializable(typeof(GetClientWindowsResult))]
+[JsonSerializable(typeof(SetDownloadBehaviorCommand))]
+[JsonSerializable(typeof(UserContextInfo))]
+[JsonSerializable(typeof(IReadOnlyList<UserContextInfo>))]
+[JsonSerializable(typeof(IReadOnlyList<ClientWindowInfo>))]
+internal partial class BrowserModuleJsonSerializerContext : JsonSerializerContext;

dotnet/src/webdriver/BiDi/BrowsingContext/BrowsingContextModule.cs

@@ -23,7 +23,7 @@
 
 namespace OpenQA.Selenium.BiDi.BrowsingContext;
 
-public sealed class BrowsingContextModule(Broker broker) : Module(broker)
+public sealed class BrowsingContextModule : Module
 {
     public async Task<BrowsingContext> CreateAsync(ContextType type, CreateOptions? options = null)
     {
@@ -250,4 +250,9 @@ public async Task<Subscription> OnUserPromptClosedAsync(Action<UserPromptClosedE
     {
         return await Broker.SubscribeAsync("browsingContext.userPromptClosed", handler, options).ConfigureAwait(false);
     }
+
+    protected internal override void Initialize(Broker broker)
+    {
+
+    }
 }

dotnet/src/webdriver/BiDi/Communication/Broker.cs

@@ -53,14 +53,14 @@ public sealed class Broker : IAsyncDisposable
     private Task? _eventEmitterTask;
     private CancellationTokenSource? _receiveMessagesCancellationTokenSource;
 
-    private readonly BiDiJsonSerializerContext _jsonSerializerContext;
+    private readonly JsonSerializerOptions _jsonSerializerOptions;
 
     internal Broker(BiDi bidi, Uri url)
     {
         _bidi = bidi;
         _transport = new WebSocketTransport(url);
 
-        var jsonSerializerOptions = new JsonSerializerOptions
+        _jsonSerializerOptions = new JsonSerializerOptions
         {
             PropertyNameCaseInsensitive = true,
             PropertyNamingPolicy = JsonNamingPolicy.CamelCase,
@@ -108,7 +108,14 @@ internal Broker(BiDi bidi, Uri url)
             }
         };
 
-        _jsonSerializerContext = new BiDiJsonSerializerContext(jsonSerializerOptions);
+        // Add base BiDi generated context resolver; keep options mutable for module contexts
+        _jsonSerializerOptions.TypeInfoResolverChain.Add(BiDiJsonSerializerContext.Default);
+    }
+
+    public void ConfigureJsonContext(Action<JsonSerializerOptions> action)
+    {
+        // Keep options mutable; do not create a context bound to them (avoids InvalidOperationException)
+        action(_jsonSerializerOptions);
     }
 
     public async Task ConnectAsync(CancellationToken cancellationToken)
@@ -205,21 +212,14 @@ private async Task<EmptyResult> ExecuteCommandCoreAsync<TCommand>(TCommand comma
         where TCommand : Command
     {
         command.Id = Interlocked.Increment(ref _currentCommandId);
-
         var tcs = new TaskCompletionSource<EmptyResult>(TaskCreationOptions.RunContinuationsAsynchronously);
-
         var timeout = options?.Timeout ?? TimeSpan.FromSeconds(30);
-
         using var cts = new CancellationTokenSource(timeout);
-
         cts.Token.Register(() => tcs.TrySetCanceled(cts.Token));
-
         _pendingCommands[command.Id] = new(command.Id, command.ResultType, tcs);
 
-        var data = JsonSerializer.SerializeToUtf8Bytes(command, typeof(TCommand), _jsonSerializerContext);
-
+        var data = JsonSerializer.SerializeToUtf8Bytes(command, typeof(TCommand), _jsonSerializerOptions);
         await _transport.SendAsync(data, cts.Token).ConfigureAwait(false);
-
         return await tcs.Task.ConfigureAwait(false);
     }
 
@@ -341,11 +341,11 @@ private void ProcessReceivedMessage(byte[]? data)
                     break;
 
                 case "result":
-                    resultReader = reader; // cloning reader with current position
+                    resultReader = reader; // snapshot
                     break;
 
                 case "params":
-                    paramsReader = reader; // cloning reader with current position
+                    paramsReader = reader; // snapshot
                     break;
 
                 case "error":
@@ -368,7 +368,7 @@ private void ProcessReceivedMessage(byte[]? data)
 
                 if (_pendingCommands.TryGetValue(id.Value, out var successCommand))
                 {
-                    var messageSuccess = JsonSerializer.Deserialize(ref resultReader, successCommand.ResultType, _jsonSerializerContext)!;
+                    var messageSuccess = JsonSerializer.Deserialize(ref resultReader, successCommand.ResultType, _jsonSerializerOptions)!;
                     successCommand.TaskCompletionSource.SetResult((EmptyResult)messageSuccess);
                     _pendingCommands.TryRemove(id.Value, out _);
                 }
@@ -384,7 +384,7 @@ private void ProcessReceivedMessage(byte[]? data)
 
                 if (_eventTypesMap.TryGetValue(method, out var eventType))
                 {
-                    var eventArgs = (EventArgs)JsonSerializer.Deserialize(ref paramsReader, eventType, _jsonSerializerContext)!;
+                    var eventArgs = (EventArgs)JsonSerializer.Deserialize(ref paramsReader, eventType, _jsonSerializerOptions)!;
 
                     var messageEvent = new MessageEvent(method, eventArgs);
                     _pendingEvents.Add(messageEvent);

dotnet/src/webdriver/BiDi/Emulation/EmulationModule.cs

@@ -23,7 +23,7 @@
 
 namespace OpenQA.Selenium.BiDi.Emulation;
 
-public sealed class EmulationModule(Broker broker) : Module(broker)
+public sealed class EmulationModule : Module
 {
     public async Task<EmptyResult> SetTimezoneOverrideAsync(string? timezone, SetTimezoneOverrideOptions? options = null)
     {
@@ -89,4 +89,9 @@ public async Task<EmptyResult> SetGeolocationPositionErrorOverrideAsync(SetGeolo
 
         return await Broker.ExecuteCommandAsync<SetGeolocationOverrideCommand, EmptyResult>(new SetGeolocationOverrideCommand(@params), options).ConfigureAwait(false);
     }
+
+    protected internal override void Initialize(Broker broker)
+    {
+        throw new NotImplementedException();
+    }
 }

dotnet/src/webdriver/BiDi/Input/InputModule.cs

@@ -23,7 +23,7 @@
 
 namespace OpenQA.Selenium.BiDi.Input;
 
-public sealed class InputModule(Broker broker) : Module(broker)
+public sealed class InputModule : Module
 {
     public async Task<EmptyResult> PerformActionsAsync(BrowsingContext.BrowsingContext context, IEnumerable<SourceActions> actions, PerformActionsOptions? options = null)
     {
@@ -45,4 +45,9 @@ public async Task<EmptyResult> SetFilesAsync(BrowsingContext.BrowsingContext con
 
         return await Broker.ExecuteCommandAsync<SetFilesCommand, EmptyResult>(new SetFilesCommand(@params), options).ConfigureAwait(false);
     }
+
+    protected internal override void Initialize(Broker broker)
+    {
+        throw new System.NotImplementedException();
+    }
 }

dotnet/src/webdriver/BiDi/Log/LogModule.cs

@@ -23,7 +23,7 @@
 
 namespace OpenQA.Selenium.BiDi.Log;
 
-public sealed class LogModule(Broker broker) : Module(broker)
+public sealed class LogModule : Module
 {
     public async Task<Subscription> OnEntryAddedAsync(Func<LogEntry, Task> handler, SubscriptionOptions? options = null)
     {
@@ -34,4 +34,9 @@ public async Task<Subscription> OnEntryAddedAsync(Action<LogEntry> handler, Subs
     {
         return await Broker.SubscribeAsync("log.entryAdded", handler, options).ConfigureAwait(false);
     }
+
+    protected internal override void Initialize(Broker broker)
+    {
+        throw new NotImplementedException();
+    }
 }

dotnet/src/webdriver/BiDi/Module.cs

@@ -21,7 +21,17 @@
 
 namespace OpenQA.Selenium.BiDi;
 
-public abstract class Module(Broker broker)
+public abstract class Module
 {
-    protected Broker Broker { get; } = broker;
+    protected Broker Broker { get; private set; }
+
+    protected internal abstract void Initialize(Broker broker);
+
+    public static TModule Create<TModule>(Broker broker) where TModule : Module, new()
+    {
+        TModule module = new();
+        module.Broker = broker;
+        module.Initialize(broker);
+        return module;
+    }
 }

dotnet/src/webdriver/BiDi/Network/NetworkModule.cs

@@ -24,7 +24,7 @@
 
 namespace OpenQA.Selenium.BiDi.Network;
 
-public sealed partial class NetworkModule(Broker broker) : Module(broker)
+public sealed partial class NetworkModule : Module
 {
     public async Task<Collector> AddDataCollectorAsync(IEnumerable<DataType> DataTypes, int MaxEncodedDataSize, AddDataCollectorOptions? options = null)
     {
@@ -173,4 +173,9 @@ public async Task<Subscription> OnAuthRequiredAsync(Action<AuthRequiredEventArgs
     {
         return await Broker.SubscribeAsync("network.authRequired", handler, options).ConfigureAwait(false);
     }
+
+    protected internal override void Initialize(Broker broker)
+    {
+        throw new System.NotImplementedException();
+    }
 }

dotnet/src/webdriver/BiDi/Script/ScriptModule.cs

@@ -23,7 +23,7 @@
 
 namespace OpenQA.Selenium.BiDi.Script;
 
-public sealed class ScriptModule(Broker broker) : Module(broker)
+public sealed class ScriptModule : Module
 {
     public async Task<EvaluateResult> EvaluateAsync(string expression, bool awaitPromise, Target target, EvaluateOptions? options = null)
     {
@@ -105,4 +105,9 @@ public async Task<Subscription> OnRealmDestroyedAsync(Action<RealmDestroyedEvent
     {
         return await Broker.SubscribeAsync("script.realmDestroyed", handler, options).ConfigureAwait(false);
     }
+
+    protected internal override void Initialize(Broker broker)
+    {
+        throw new NotImplementedException();
+    }
 }

dotnet/src/webdriver/BiDi/Session/SessionModule.cs

@@ -23,7 +23,7 @@
 
 namespace OpenQA.Selenium.BiDi.Session;
 
-internal sealed class SessionModule(Broker broker) : Module(broker)
+internal sealed class SessionModule : Module
 {
     public async Task<StatusResult> StatusAsync(StatusOptions? options = null)
     {
@@ -55,4 +55,9 @@ public async Task<EmptyResult> EndAsync(EndOptions? options = null)
     {
         return await Broker.ExecuteCommandAsync<EndCommand, EmptyResult>(new EndCommand(), options).ConfigureAwait(false);
     }
+
+    protected internal override void Initialize(Broker broker)
+    {
+        throw new System.NotImplementedException();
+    }
 }

dotnet/src/webdriver/BiDi/Storage/StorageModule.cs

@@ -22,7 +22,7 @@
 
 namespace OpenQA.Selenium.BiDi.Storage;
 
-public sealed class StorageModule(Broker broker) : Module(broker)
+public sealed class StorageModule : Module
 {
     public async Task<GetCookiesResult> GetCookiesAsync(GetCookiesOptions? options = null)
     {
@@ -44,4 +44,9 @@ public async Task<SetCookieResult> SetCookieAsync(PartialCookie cookie, SetCooki
 
         return await Broker.ExecuteCommandAsync<SetCookieCommand, SetCookieResult>(new SetCookieCommand(@params), options).ConfigureAwait(false);
     }
+
+    protected internal override void Initialize(Broker broker)
+    {
+        throw new System.NotImplementedException();
+    }
 }