feat: [dotnet] open control channel for grpc save/load

dotnet/src/Microsoft.AutoGen/Core.Grpc/GrpcAgentRuntime.cs

@@ -483,6 +483,7 @@ public async ValueTask OnMessageAsync(Message message, CancellationToken cancell
                 var cloudEvent = message.CloudEvent ?? throw new InvalidOperationException("CloudEvent is null.");
                 await HandlePublish(cloudEvent);
                 break;
+                throw new InvalidOperationException("ControlMessage is not supported.");
             default:
                 throw new InvalidOperationException($"Unexpected message '{message}'.");
         }

dotnet/src/Microsoft.AutoGen/RuntimeGateway.Grpc/Services/Grpc/GrpcGateway.cs

@@ -31,10 +31,19 @@
     /// </summary>
     private readonly IMessageRegistryGrain _messageRegistry;
     private readonly IGateway _reference;
-    private readonly ConcurrentDictionary<string, List<GrpcWorkerConnection>> _supportedAgentTypes = [];
-    public readonly ConcurrentDictionary<string, GrpcWorkerConnection> _workers = new();
-    private readonly ConcurrentDictionary<(string Type, string Key), GrpcWorkerConnection> _agentDirectory = new();
-    private readonly ConcurrentDictionary<(GrpcWorkerConnection, string), TaskCompletionSource<RpcResponse>> _pendingRequests = new();
+    private readonly ConcurrentDictionary<string, List<GrpcWorkerConnection<Message>>> _supportedAgentTypes = [];
+    /// <summary>
+    /// a map of the clientids form the grpc connection headers to the worker connections that service that id
+    /// for the Message Channel
+    /// </summary>
+    public readonly ConcurrentDictionary<string, GrpcWorkerConnection<Message>> _workers = new();
+    /// <summary>
+    /// a map of the clientids form the grpc connection headers to the worker connections for the Control Channel
+    /// </summary>
+    public readonly ConcurrentDictionary<string, GrpcWorkerConnection<ControlMessage>> _controlWorkers = new();
+    private readonly ConcurrentDictionary<(string Type, string Key), GrpcWorkerConnection<Message>> _agentDirectory = new();
+    private readonly ConcurrentDictionary<(GrpcWorkerConnection<Message>, string), TaskCompletionSource<RpcResponse>> _pendingRequests = new();
+    private readonly ConcurrentDictionary<(GrpcWorkerConnection<ControlMessage>, string), TaskCompletionSource<RpcResponse>> _pendingControlRequests = new();
 
     /// <summary>
     /// Initializes a new instance of the <see cref="GrpcGateway"/> class.
@@ -255,16 +264,25 @@
     /// <param name="responseStream">The response stream.</param>
     /// <param name="context">The server call context.</param>
     /// <returns>A task that represents the asynchronous operation.</returns>
-    internal async Task ConnectToWorkerProcess(IAsyncStreamReader<Message> requestStream, IServerStreamWriter<Message> responseStream, ServerCallContext context)
+    internal async Task ConnectToWorkerProcess<TMessage>(IAsyncStreamReader<TMessage> requestStream, IServerStreamWriter<TMessage> responseStream, ServerCallContext context)
+    where TMessage : class
     {
         _logger.LogInformation("Received new connection from {Peer}.", context.Peer);
-        var clientId = (context.RequestHeaders.Get("client-id")?.Value) ??
-            throw new RpcException(new Status(StatusCode.InvalidArgument, "Client ID is required."));
-        var workerProcess = new GrpcWorkerConnection(this, requestStream, responseStream, context);
-        _workers.GetOrAdd(clientId, workerProcess);
-
-        await this.AttachDanglingRegistrations(clientId).ConfigureAwait(false);
-
+        var clientId = context.RequestHeaders.Get("client-id")?.Value
+            ?? throw new RpcException(new Status(StatusCode.InvalidArgument, "Client ID is required."));
+        var workerProcess = new GrpcWorkerConnection<TMessage>(this, requestStream, responseStream, context);
+        if (typeof(TMessage) == typeof(Message))
+        {
+            _workers.GetOrAdd(clientId, _ => (GrpcWorkerConnection<Message>)(object)workerProcess);
+        }
+        else if (typeof(TMessage) == typeof(ControlMessage))
+        {
+            _controlWorkers.GetOrAdd(clientId, _ => (GrpcWorkerConnection<ControlMessage>)(object)workerProcess);
+        }
+        else
+        {
+            throw new InvalidOperationException($"Unsupported message type: {typeof(TMessage).Name}");
+        }
         await workerProcess.Connect().ConfigureAwait(false);
     }
 
@@ -301,20 +319,35 @@
     /// <param name="message">The received message.</param>
     /// <param name="cancellationToken">The cancellation token.</param>
     /// <returns>A task that represents the asynchronous operation.</returns>
-    internal async Task OnReceivedMessageAsync(GrpcWorkerConnection connection, Message message, CancellationToken cancellationToken = default)
+    internal async Task OnReceivedMessageAsync<TMessage>(GrpcWorkerConnection<TMessage> connection, TMessage message, CancellationToken cancellationToken = default)
+    where TMessage : class
     {
         _logger.LogInformation("Received message {Message} from connection {Connection}.", message, connection);
-        switch (message.MessageCase)
+
+        switch (message)
         {
-            case Message.MessageOneofCase.Request:
-                await DispatchRequestAsync(connection, message.Request);
-                break;
-            case Message.MessageOneofCase.Response:
-                DispatchResponse(connection, message.Response);
-                break;
-            case Message.MessageOneofCase.CloudEvent:
-                await DispatchEventAsync(message.CloudEvent, cancellationToken);
+            case Message msg:
+                switch (msg.MessageCase)
+                {
+                    case Message.MessageOneofCase.Request:
+                        await DispatchRequestAsync(connection, msg.Request);
+                        break;
+                    case Message.MessageOneofCase.Response:
+                        DispatchResponse(connection, msg.Response);
+                        break;
+                    case Message.MessageOneofCase.CloudEvent:
+                        await DispatchEventAsync(msg.CloudEvent, cancellationToken);
+                        break;
+                    default:
+                        await RespondBadRequestAsync(connection, $"Unknown message type for message '{msg}'.");
+                        break;
+                }
                 break;
+
+            //case ControlMessage controlMsg:
+            //    //await HandleControlMessageAsync(connection, controlMsg, cancellationToken);
+            //    break;
+
             default:
                 await RespondBadRequestAsync(connection, $"Unknown message type for message '{message}'.");
                 break;
@@ -326,14 +359,25 @@
     /// </summary>
     /// <param name="connection">The worker connection.</param>
     /// <param name="response">The RPC response.</param>
-    private void DispatchResponse(GrpcWorkerConnection connection, RpcResponse response)
+    private void DispatchResponse<TMessage>(GrpcWorkerConnection<TMessage> connection, RpcResponse response)
+    where TMessage : class
     {
-        if (!_pendingRequests.TryRemove((connection, response.RequestId), out var completion))
+        if (connection is GrpcWorkerConnection<Message> messageConnection)
         {
-            _logger.LogWarning("Received response for unknown request id: {RequestId}.", response.RequestId);
-            return;
+            if (_pendingRequests.TryRemove((messageConnection, response.RequestId), out var completion))
+            {
+                completion.SetResult(response);
+                return;
+            }
+        }
+        else if (connection is GrpcWorkerConnection<ControlMessage> controlConnection)
+        {
+            if (_pendingControlRequests.TryRemove((controlConnection, response.RequestId), out var completion))
+            {
+                completion.SetResult(response);
+                return;
+            }
         }
-        completion.SetResult(response);
     }
 
     /// <summary>
@@ -392,7 +436,8 @@
     /// <param name="connection">The worker connection.</param>
     /// <param name="request">The RPC request.</param>
     /// <returns>A task that represents the asynchronous operation.</returns>
-    private async ValueTask DispatchRequestAsync(GrpcWorkerConnection connection, RpcRequest request)
+    private async ValueTask DispatchRequestAsync<TMessage>(GrpcWorkerConnection<TMessage> connection, RpcRequest request)
+    where TMessage : class
     {
         var requestId = request.RequestId;
         if (request.Target is null)
@@ -422,42 +467,60 @@
     /// <param name="request">The RPC request.</param>
     /// <param name="func">The function to invoke.</param>
     /// <returns>A task that represents the asynchronous operation.</returns>
-    private static async Task InvokeRequestDelegate(GrpcWorkerConnection connection, RpcRequest request, Func<RpcRequest, Task<RpcResponse>> func)
+    private static async Task InvokeRequestDelegate<TMessage>(GrpcWorkerConnection<TMessage> connection, RpcRequest request, Func<RpcRequest, Task<RpcResponse>> func)
+    where TMessage : class
     {
         try
         {
             var response = await func(request);
             response.RequestId = request.RequestId;
-            await connection.ResponseStream.WriteAsync(new Message { Response = response }).ConfigureAwait(false);
+
+            if (connection is GrpcWorkerConnection<Message> messageConnection)
+            {
+                await messageConnection.ResponseStream.WriteAsync(new Message { Response = response }).ConfigureAwait(false);
+            }
         }
         catch (Exception ex)
         {
-            await connection.ResponseStream.WriteAsync(new Message { Response = new RpcResponse { RequestId = request.RequestId, Error = ex.Message } }).ConfigureAwait(false);
+            if (connection is GrpcWorkerConnection<Message> messageConnection)
+            {
+                await messageConnection.ResponseStream.WriteAsync(
+                    new Message { Response = new RpcResponse { RequestId = request.RequestId, Error = ex.Message } }
+                ).ConfigureAwait(false);
+            }
         }
     }
 
     /// <summary>
     /// Handles the removal of a worker process.
     /// </summary>
     /// <param name="workerProcess">The worker process.</param>
-    internal void OnRemoveWorkerProcess(GrpcWorkerConnection workerProcess)
+    internal void OnRemoveWorkerProcess<TMessage>(GrpcWorkerConnection<TMessage> workerProcess)
+    where TMessage : class
     {
-        var clientId = workerProcess.ServerCallContext.RequestHeaders.Get("client-id")?.Value ??
-            throw new RpcException(new Status(StatusCode.InvalidArgument, "Grpc Client ID is required."));
+        var clientId = workerProcess.ServerCallContext.RequestHeaders.Get("client-id")?.Value
+            ?? throw new RpcException(new Status(StatusCode.InvalidArgument, "Grpc Client ID is required."));
+
         _workers.TryRemove(clientId, out _);
+        _controlWorkers.TryRemove(clientId, out _);
+
         var types = workerProcess.GetSupportedTypes();
         foreach (var type in types)
         {
-            if (_supportedAgentTypes.TryGetValue(type, out var supported))
+            if (_supportedAgentTypes.TryGetValue(type, out var supported) && workerProcess is GrpcWorkerConnection<Message> messageWorker)
             {
-                supported.Remove(workerProcess);
+                supported.Remove(messageWorker);
             }
         }
-        foreach (var pair in _agentDirectory)
+
+        if (workerProcess is GrpcWorkerConnection<Message> messageWorkerInstance)
         {
-            if (pair.Value == workerProcess)
+            foreach (var pair in _agentDirectory.ToList())
             {
-                ((IDictionary<(string Type, string Key), GrpcWorkerConnection>)_agentDirectory).Remove(pair);
+                if (ReferenceEquals(pair.Value, messageWorkerInstance)) // Ensures exact instance match
+                {
+                    _agentDirectory.TryRemove(pair.Key, out _);
+                }
             }
         }
     }
@@ -468,7 +531,8 @@
     /// <param name="connection">The worker connection.</param>
     /// <param name="error">The error message.</param>
     /// <returns>A task that represents the asynchronous operation.</returns>
-    private static async ValueTask RespondBadRequestAsync(GrpcWorkerConnection connection, string error)
+    private static async ValueTask RespondBadRequestAsync<TMessage>(GrpcWorkerConnection<TMessage> connection, string error)
+    where TMessage : class
     {
         throw new RpcException(new Status(StatusCode.InvalidArgument, error));
     }
@@ -480,7 +544,7 @@
     /// <param name="cloudEvent">The cloud event.</param>
     /// <param name="cancellationToken">The cancellation token.</param>
     /// <returns>A task that represents the asynchronous operation.</returns>
-    private async Task WriteResponseAsync(GrpcWorkerConnection connection, CloudEvent cloudEvent, CancellationToken cancellationToken = default)
+    private async Task WriteResponseAsync(GrpcWorkerConnection<Message> connection, CloudEvent cloudEvent, CancellationToken cancellationToken = default)
     {
         await connection.ResponseStream.WriteAsync(new Message { CloudEvent = cloudEvent }, cancellationToken).ConfigureAwait(false);
     }

dotnet/src/Microsoft.AutoGen/RuntimeGateway.Grpc/Services/Grpc/GrpcGatewayService.cs

@@ -35,6 +35,31 @@
         }
     }
 
+    /// <summary>
+    /// Open channel for the Control Channel (defined in the proto file).
+    /// </summary>
+    /// <param name="requestStream">The request stream.</param>
+    /// <param name="responseStream">The response stream.</param>
+    /// <param name="context">The server call context.</param>
+    /// <returns>A task that represents the asynchronous operation.</returns>
+    /// <remarks>Control channel is used for control messages between the agent worker and the cluster.</remarks>
+    ///       public virtual global::System.Threading.Tasks.Task OpenControlChannel(grpc::IAsyncStreamReader<global::Microsoft.AutoGen.Protobuf.ControlMessage> requestStream, grpc::IServerStreamWriter<global::Microsoft.AutoGen.Protobuf.ControlMessage> responseStream, grpc::ServerCallContext context)
+    public override async Task OpenControlChannel(IAsyncStreamReader<ControlMessage> requestStream, IServerStreamWriter<ControlMessage> responseStream, ServerCallContext context)
+    {
+        try
+        {
+            await Gateway.ConnectToWorkerProcess(requestStream, responseStream, context).ConfigureAwait(true);
+        }
+        catch
+        {
+            if (context.CancellationToken.IsCancellationRequested)
+            {
+                return;
+            }
+            throw;
+        }
+    }
+
     /// <summary>
     /// Adds a subscription.
     /// </summary>

dotnet/src/Microsoft.AutoGen/RuntimeGateway.Grpc/Services/Grpc/GrpcWorkerConnection.cs

@@ -2,11 +2,11 @@
 // GrpcWorkerConnection.cs
 using System.Threading.Channels;
 using Grpc.Core;
-using Microsoft.AutoGen.Protobuf;
 
 namespace Microsoft.AutoGen.RuntimeGateway.Grpc;
 
-public sealed class GrpcWorkerConnection : IAsyncDisposable
+public sealed class GrpcWorkerConnection<TMessage> : IAsyncDisposable
+where TMessage : class
 {
     private static long s_nextConnectionId;
     private Task _readTask = Task.CompletedTask;
@@ -17,13 +17,13 @@
     private readonly GrpcGateway _gateway;
     private readonly CancellationTokenSource _shutdownCancellationToken = new();
     public Task Completion { get; private set; } = Task.CompletedTask;
-    public GrpcWorkerConnection(GrpcGateway agentWorker, IAsyncStreamReader<Message> requestStream, IServerStreamWriter<Message> responseStream, ServerCallContext context)
+    public GrpcWorkerConnection(GrpcGateway agentWorker, IAsyncStreamReader<TMessage> requestStream, IServerStreamWriter<TMessage> responseStream, ServerCallContext context)
     {
         _gateway = agentWorker;
         RequestStream = requestStream;
         ResponseStream = responseStream;
         ServerCallContext = context;
-        _outboundMessages = Channel.CreateUnbounded<Message>(new UnboundedChannelOptions { AllowSynchronousContinuations = true, SingleReader = true, SingleWriter = false });
+        _outboundMessages = Channel.CreateUnbounded<TMessage>(new UnboundedChannelOptions { AllowSynchronousContinuations = true, SingleReader = true, SingleWriter = false });
     }
     public Task Connect()
     {
@@ -50,11 +50,11 @@
         return Completion = Task.WhenAll(_readTask, _writeTask);
     }
 
-    public IAsyncStreamReader<Message> RequestStream { get; }
-    public IServerStreamWriter<Message> ResponseStream { get; }
+    public IAsyncStreamReader<TMessage> RequestStream { get; }
+    public IServerStreamWriter<TMessage> ResponseStream { get; }
     public ServerCallContext ServerCallContext { get; }
 
-    private readonly Channel<Message> _outboundMessages;
+    private readonly Channel<TMessage> _outboundMessages;
 
     /// <inheritdoc/>
     public void AddSupportedType(string type)
@@ -75,7 +75,7 @@
     }
 
     /// <inheritdoc/>
-    public async Task SendMessage(Message message)
+    public async Task SendMessage(TMessage message)
     {
         await _outboundMessages.Writer.WriteAsync(message).ConfigureAwait(false);
     }