[dotnet] [bidi] Provide extension points for custom BiDi modules

[RenderMichael]

User description

Description

Enables the right extensibility for users to create their own Modules, and hook them into BiDi.

Example implementation of Permissions module:

public class PermissionsBiDi : BiDi
{
    private readonly Lazy<PermissionsModule> _permissionsModule;

    public PermissionsModule Permissions => _permissionsModule.Value;

    private PermissionsBiDi(string url)
        : base(url)
    {
        Broker.ProvideCustomSerializationContext(PermissionsModule.SerializerContext);
        _permissionsModule = new Lazy<PermissionsModule>(() => new PermissionsModule(Broker));
    }

    public static async Task<PermissionsBiDi> ConnectAsync(string webSocketUrl)
    {
        var bidi = new PermissionsBiDi(webSocketUrl);
        await bidi.Broker.ConnectAsync();
        return bidi;
    }
}

public class PermissionsModule(Broker broker) : Module(broker)
{
    public static JsonSerializerContext SerializerContext => PermissionsBiDiJsonSerializerContext.Default;

    public async Task SetPermissionAsync(string permissionName, PermissionState state, string origin, UserContext? userContext, SetPermissionOptions? options = null)
    {
        var @params = new SetPermissionCommandParameters(new PermissionDescriptor(permissionName), state, origin, userContext);

        await Broker.ExecuteCommandAsync(new SetPermissionCommand(@params), options).ConfigureAwait(false);
    }
}

public record SetPermissionOptions : CommandOptions;

internal class SetPermissionCommand(SetPermissionCommandParameters @params)
    : Command<SetPermissionCommandParameters>(@params, "permissions.setPermission") { }

internal record SetPermissionCommandParameters(PermissionDescriptor Descriptor, PermissionState State, string Origin, UserContext? UserContext) : CommandParameters { }

internal record PermissionDescriptor(string Name);

public enum PermissionState
{
    Granted,
    Denied,
    Prompt
}

[JsonSerializable(typeof(SetPermissionCommand), TypeInfoPropertyName = "Permissions_SetPermissionCommand")]
internal partial class PermissionsBiDiJsonSerializerContext : JsonSerializerContext;

Or, if user does not care about AOT-safe JSON:

Sample Permissions module

public class PermissionsBiDi : BiDi
{
    private readonly Lazy<PermissionsModule> _permissionsModule;

    public PermissionsModule Permissions => _permissionsModule.Value;

    private PermissionsBiDi(string url)
        : base(url)
    {
        Broker.EnableReflectionBasedJson();
        _permissionsModule = new Lazy<PermissionsModule>(() => new PermissionsModule(Broker));
    }

    public static async Task<PermissionsBiDi> ConnectAsync(string webSocketUrl)
    {
        var bidi = new PermissionsBiDi(webSocketUrl);
        await bidi.Broker.ConnectAsync();
        return bidi;
    }
}

public class PermissionsModule(Broker broker) : Module(broker)
{
    public async Task SetPermissionAsync(string permissionName, PermissionState state, string origin, UserContext? userContext, SetPermissionOptions? options = null)
    {
        var @params = new SetPermissionCommandParameters(new PermissionDescriptor(permissionName), state, origin, userContext);

        await Broker.ExecuteCommandAsync(new SetPermissionCommand(@params), options).ConfigureAwait(false);
    }
}

public record SetPermissionOptions : CommandOptions;

internal class SetPermissionCommand(SetPermissionCommandParameters @params)
    : Command<SetPermissionCommandParameters>(@params, "permissions.setPermission") { }

internal record SetPermissionCommandParameters(PermissionDescriptor Descriptor, PermissionState State, string Origin, UserContext? UserContext) : CommandParameters { }

internal record PermissionDescriptor(string Name);

public enum PermissionState
{
    Granted,
    Denied,
    Prompt
}

Motivation and Context

Provides a roadmap for users to implement custom BiDi modules, such as in #15329.

Types of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)

Checklist

• I have read the contributing document.
• My change requires a change to the documentation.
• I have updated the documentation accordingly.
• I have added tests to cover my changes.
• All new and existing tests passed.

---

PR Type

Enhancement

---

Description

• Added extensibility to BiDi for custom modules.

• Introduced ProvideCustomSerializationContext for custom serialization.

• Enhanced DisposeAsync with a virtual core method.

• Improved error handling and serialization in Broker.

---

Changes walkthrough 📝

Relevant files

Enhancement

BiDi.cs Enhanced `BiDi` class for extensibility and disposal          dotnet/src/webdriver/BiDi/BiDi.cs Made Broker property protected for extensibility. Changed constructor to protected internal for subclassing. Updated lazy initialization to use Broker directly. Refactored DisposeAsync to include a virtual core method. +20/-15  Broker.cs Enhanced `Broker` for custom serialization and error handling dotnet/src/webdriver/BiDi/Communication/Broker.cs Added ProvideCustomSerializationContext for custom serializers. Refined JSON serialization with TypeInfoResolverChain. Improved error handling in message deserialization. Fixed typo in MessageError handling. +18/-7

---

Need help?

Type /help how to ... in the comments thread for any questions about Qodo Merge usage.

Check out the documentation for more information.

[qodo-code-review]

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🎫 Ticket compliance analysis ✅ 15329 - PR Code Verified Compliant requirements: Provide a possibility to implement BiDi Permissions module in .NET Implementation does not rely on internal members Exposes minimal required functionality to make it possible to implement new modules Requires further human verification: Verify if the implementation supports all the usage examples mentioned in the ticket (specifically the extension method approach with bidi.AsPermissions())

⏱️ Estimated effort to review: 2 🔵🔵⚪⚪⚪

🧪 No relevant tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Potential Memory Leak The DisposeAsync method now calls DisposeAsyncCore but doesn't handle exceptions. If an exception occurs during disposal of the Broker, resources might not be properly cleaned up. public async ValueTask DisposeAsync() { await DisposeAsyncCore(); GC.SuppressFinalize(this); } Thread Safety The ProvideCustomSerializationContext method modifies the TypeInfoResolverChain collection without any synchronization mechanism, which could lead to thread safety issues if called concurrently. public void ProvideCustomSerializationContext(JsonSerializerContext extensionContext) { _jsonSerializerContext.TypeInfoResolverChain.Add(extensionContext); }

[qodo-code-review]

PR Code Suggestions ✨

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Add exception handling The deserialization code doesn't handle potential exceptions that could occur during message deserialization. If malformed data is received, this could crash the message receiving loop. dotnet/src/webdriver/BiDi/Communication/Broker.cs [127-134] private async Task ReceiveMessagesAsync(CancellationToken cancellationToken) { while (!cancellationToken.IsCancellationRequested) { - var data = await _transport.ReceiveAsync(cancellationToken).ConfigureAwait(false); + try + { + var data = await _transport.ReceiveAsync(cancellationToken).ConfigureAwait(false); - var messageTypeInfo = (JsonTypeInfo<Message>)_jsonSerializerContext.GetTypeInfo(typeof(Message)); - var message = JsonSerializer.Deserialize(new ReadOnlySpan<byte>(data), messageTypeInfo); + var messageTypeInfo = (JsonTypeInfo<Message>)_jsonSerializerContext.GetTypeInfo(typeof(Message)); + var message = JsonSerializer.Deserialize(new ReadOnlySpan<byte>(data), messageTypeInfo); Apply this suggestion Suggestion importance[1-10]: 9 __ Why: This suggestion addresses a critical reliability issue. Without exception handling, any deserialization error would crash the message receiving loop, potentially breaking the entire communication channel. Adding try-catch blocks is essential for robust error handling in network communication.	High
	Ensure proper disposal sequence The _transport field is being disposed after Broker is disposed, but Broker might be using the transport during its disposal. Consider disposing _transport only after ensuring Broker has been fully disposed. dotnet/src/webdriver/BiDi/BiDi.cs [86-90] protected virtual async ValueTask DisposeAsyncCore() { - await Broker.DisposeAsync().ConfigureAwait(false); - _transport?.Dispose(); + try + { + await Broker.DisposeAsync().ConfigureAwait(false); + } + finally + { + _transport?.Dispose(); + } } Apply this suggestion Suggestion importance[1-10]: 8 __ Why: This suggestion addresses a potential resource management issue by ensuring that transport disposal happens even if Broker disposal throws an exception. Using try-finally ensures proper cleanup sequence, which is critical for resource management and preventing resource leaks.	Medium
Learned best practice	Add null validation for method parameters to prevent NullReferenceException and provide clear error messages The extensionContext parameter is not checked for null before being used. If null is passed, it would cause a NullReferenceException when attempting to add it to the TypeInfoResolverChain. Add a null check using ArgumentNullException.ThrowIfNull() to validate the parameter. dotnet/src/webdriver/BiDi/Communication/Broker.cs [113-116] public void ProvideCustomSerializationContext(JsonSerializerContext extensionContext) { + ArgumentNullException.ThrowIfNull(extensionContext, nameof(extensionContext)); _jsonSerializerContext.TypeInfoResolverChain.Add(extensionContext); } Apply this suggestion Suggestion importance[1-10]: 6	Low
Update

[RenderMichael]

BiDi is not a sealed type, we have to implement the IAsyncDisposable pattern for types which users may extend.

[RenderMichael]

BiDiJsonSerializerContext is no longer "king".

We know that the Message JSON type info exists, but we must extract it manually now (AOT-safe).

[RenderMichael]

Guidance for users: do not call this method after Broker.ConnectAsync is called. Json serialization options becomes immutable and this method will throw..

[nvborisenko]

public class PermissionsBiDi : BiDi

It it big mistake, meaning once I get permissionsObject somehow, I will see all members from BiDi like permissionsObject.Script.CallFunctionAsync(...). It is not related to Permissions module.

I would say that PermissionsBiDi class should depends on Broker only.

[RenderMichael]

With this PR, all the core Modules are still internal, so a user cannot see them directly. Even a 3rd party PermissionsBiDi class cannot see them. The only thing this PR does it expose Broker to child classes, for extension (Broker needs to know about additional module, due to JSON).

[RenderMichael]

Ah, I see what you mean. I will try to make that vision a reality (will need to do something about ITransport)

[RenderMichael]

Looking more into it, I'm not sure what having BiDi extensions without the core modules does. I have a follow-up PR up with a Permissions implementation #15421. If you look at the test, it uses core modules a well.

Users should be free to use the modules without core BiDi, but I don't understand the use case there.

[RenderMichael]

BiDi is too closely tied with Broker and everything else. Types like Request carry around the instance of BiDi they were created by. It is not possible to remove the BiDi type from this equation, and I am not sure it is desirable.

[nvborisenko]

Initially it was designed for continuation from EventArgs:

bidi.Network.OnBeforeRequestSent(async args => await args.BiDi.DoSomethingAsync());

We can revisit this design. Most likely it was motivated by network intercepting.

[nvborisenko]

I see the first step: Broker should accept JsonSerializerOptions in ctor, Broker doesn't want BiDi.

[RenderMichael]

Can you help me understand some design:

1. Why are Modules internal? Isn't that the main way users of the BiDi type will use it?
2. Is there a scenario for BiDi interaction without the core Modules?
3. If a user gets something like permissionsObject with a BiDi property on it, what do they need the BiDi for? Maybe users do not need to receive it at all?

[nvborisenko]

I think BiDi can construct his JsonSerializerContext, and then Broker will rely on it. So I am thinking on a way completely new own permissions json context rather than extending default. Why: probably having lighter json context is beneficial for serialization in terms of performance.

And Appium (who is most likely on top of Selenium), will also provide their own json context, they are still able to define it on top of our default.

[nvborisenko]

Seems good idea. Then We can:

- internal Broker(BiDi bidi, Uri url)
+ internal Broker(SessionModule session, Uri url)

Meaning be closer to BiDi -> Broker -> Transport