AI-Studio/app/MindWork AI Studio/Tools/MessageBus.cs

using System.Collections.Concurrent;

using Microsoft.AspNetCore.Components;
// ReSharper disable RedundantRecordClassKeyword

namespace AIStudio.Tools;

public sealed class MessageBus
{
    public static readonly MessageBus INSTANCE = new();
    
    private readonly ConcurrentDictionary<IMessageBusReceiver, ComponentBase[]> componentFilters = new();
    private readonly ConcurrentDictionary<IMessageBusReceiver, Event[]> componentEvents = new();
    private readonly ConcurrentDictionary<Event, ConcurrentQueue<Message>> deferredMessages = new();
    private readonly ConcurrentQueue<Message> messageQueue = new();
    private readonly SemaphoreSlim sendingSemaphore = new(1, 1);

    private static ILogger<MessageBus>? LOG; 

    private MessageBus()
    {
    }
    
    public void Initialize(ILogger<MessageBus> logger)
    {
        LOG = logger;
        LOG.LogInformation("Message bus initialized.");
    }

    /// <summary>
    /// Define for which components and events you want to receive messages.
    /// </summary>
    /// <param name="receiver">That's you, the receiver.</param>
    /// <param name="filterComponents">A list of components for which you want to receive messages. Use an empty list to receive messages from all components.</param>
    /// <param name="events">A list of events for which you want to receive messages.</param>
    public void ApplyFilters(IMessageBusReceiver receiver, ComponentBase[] filterComponents, HashSet<Event> events)
    {
        this.componentFilters[receiver] = filterComponents;
        this.componentEvents[receiver] = events.ToArray();
    }
    
    public void RegisterComponent(IMessageBusReceiver receiver)
    {
        this.componentFilters.TryAdd(receiver, []);
        this.componentEvents.TryAdd(receiver, []);
    }
    
    public void Unregister(IMessageBusReceiver receiver)
    {
        this.componentFilters.TryRemove(receiver, out _);
        this.componentEvents.TryRemove(receiver, out _);
    }
    
    private record class Message(ComponentBase? SendingComponent, Event TriggeredEvent, object? Data);
    
    public async Task SendMessage<T>(ComponentBase? sendingComponent, Event triggeredEvent, T? data = default)
    {
        this.messageQueue.Enqueue(new Message(sendingComponent, triggeredEvent, data));

        try
        {
            await this.sendingSemaphore.WaitAsync();
            while (this.messageQueue.TryDequeue(out var message))
            {
                foreach (var (receiver, componentFilter) in this.componentFilters)
                {
                    if (componentFilter.Length > 0 && message.SendingComponent is not null && !componentFilter.Contains(message.SendingComponent))
                        continue;

                    var eventFilter = this.componentEvents[receiver];
                    if (eventFilter.Length == 0 || eventFilter.Contains(message.TriggeredEvent))

                        // We don't await the task here because we don't want to block the message bus:
                        _ = receiver.ProcessMessage(message.SendingComponent, message.TriggeredEvent, message.Data);
                }
            }
        }
        catch (Exception e)
        {
            LOG?.LogError(e, "Error while sending message.");
        }
        finally
        {
            this.sendingSemaphore.Release();
        }
    }

    public Task SendError(DataErrorMessage dataErrorMessage) => this.SendMessage(null, Event.SHOW_ERROR, dataErrorMessage);
    
    public Task SendWarning(DataWarningMessage dataWarningMessage) => this.SendMessage(null, Event.SHOW_WARNING, dataWarningMessage);
    
    public Task SendSuccess(DataSuccessMessage dataSuccessMessage) => this.SendMessage(null, Event.SHOW_SUCCESS, dataSuccessMessage);

    public Task SendInfo(DataInfoMessage dataInfoMessage) => this.SendMessage(null, Event.SHOW_INFO, dataInfoMessage);

    /// <summary>
    /// Stores a message until someone asks for it, cf. TakeDeferredMessages. This is how a
    /// component hands data to a component which does not exist yet, e.g. an assistant which
    /// sends its result to the chat before the user gets there.
    /// </summary>
    /// <param name="sendingComponent">That's you, the sender.</param>
    /// <param name="triggeredEvent">The event this message belongs to.</param>
    /// <param name="data">The data to hand over.</param>
    public void DeferMessage<T>(ComponentBase? sendingComponent, Event triggeredEvent, T? data = default)
    {
        var queue = this.deferredMessages.GetOrAdd(triggeredEvent, _ => new());
        queue.Enqueue(new Message(sendingComponent, triggeredEvent, data));
    }

    /// <summary>
    /// Takes all deferred messages of an event out of the bus.
    /// </summary>
    /// <remarks>
    /// This empties the queue and returns what was in it. It used to be a lazy iterator, which
    /// meant that a caller stopping after the first message left the rest of the queue behind:
    /// those messages were never delivered, and the data they carry — a complete chat thread, for
    /// instance — stayed alive for as long as the app ran. Returning a list makes that impossible.
    /// Callers who expect a single message take the last one, since that is the most recent thing
    /// the user asked for.
    /// </remarks>
    /// <param name="triggeredEvent">The event whose messages you want.</param>
    /// <returns>The deferred messages, oldest first. Empty when there are none.</returns>
    public IReadOnlyList<T?> TakeDeferredMessages<T>(Event triggeredEvent)
    {
        //
        // Removing the queue along with its messages is what keeps the dictionary from growing:
        // otherwise, every event which ever deferred a message would keep an empty queue forever.
        //
        if (!this.deferredMessages.TryRemove(triggeredEvent, out var queue))
            return [];

        var messages = new List<T?>();
        while (queue.TryDequeue(out var message))
            messages.Add(message.Data is T data ? data : default);

        return messages;
    }
    
    public async Task<TResult?> SendMessageUseFirstResult<TPayload, TResult>(ComponentBase? sendingComponent, Event triggeredEvent, TPayload? data = default)
    {
        foreach (var (receiver, componentFilter) in this.componentFilters)
        {
            if (componentFilter.Length > 0 && sendingComponent is not null && !componentFilter.Contains(sendingComponent))
                continue;

            var eventFilter = this.componentEvents[receiver];
            if (eventFilter.Length == 0 || eventFilter.Contains(triggeredEvent))
            {
                var result = await receiver.ProcessMessageWithResult<TPayload, TResult>(sendingComponent, triggeredEvent, data);
                if (result is not null)
                    return (TResult) result;
            }
        }
        
        return default;
    }
}