fixing how IGateway works

dotnet/src/Microsoft.AutoGen/Abstractions/IGateway.cs

@@ -9,4 +9,6 @@ public interface IGateway : IGrainObserver
     ValueTask BroadcastEvent(CloudEvent evt, CancellationToken cancellationToken = default);
     ValueTask StoreAsync(AgentState value, CancellationToken cancellationToken = default);
     ValueTask<AgentState> ReadAsync(AgentId agentId, CancellationToken cancellationToken = default);
+    Task SendMessageAsync(IConnection connection, CloudEvent cloudEvent, CancellationToken cancellationToken = default);
+
 }

dotnet/src/Microsoft.AutoGen/Agents/Services/Gateway.cs

@@ -19,7 +19,7 @@ internal class Gateway : BackgroundService, IGateway, IGrainWithIntegerKey
 
     private readonly ConcurrentDictionary<string, List<InMemoryQueue<CloudEvent>>> _supportedAgentTypes = [];
     private readonly ConcurrentDictionary<(string Type, string Key), InMemoryQueue<CloudEvent>> _agentDirectory = new();
-    private readonly ConcurrentDictionary<InMemoryQueue<CloudEvent>, InMemoryQueue<CloudEvent>> _workers = new();
+    public readonly ConcurrentDictionary<IConnection, IConnection> _workers = new();
     private readonly ConcurrentDictionary<(InMemoryQueue<Message>, string), TaskCompletionSource<RpcResponse>> _pendingRequests = new();
     private readonly InMemoryQueue<Message> _messageQueue = new();
 
@@ -44,7 +44,7 @@ public async ValueTask BroadcastEvent(CloudEvent evt, CancellationToken cancella
     }
 
     // intentionally not static
-    private async Task SendMessageAsync(IConnection connection, CloudEvent cloudEvent, CancellationToken cancellationToken = default)
+    public virtual async Task SendMessageAsync(IConnection connection, CloudEvent cloudEvent, CancellationToken cancellationToken = default)
     {
         var queue = (InMemoryQueue<CloudEvent>)connection;
         await queue.Writer.WriteAsync(cloudEvent, cancellationToken).AsTask().ConfigureAwait(false);

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

@@ -15,8 +15,6 @@ internal sealed class GrpcGateway : Gateway, IGateway
     private readonly IAgentRegistry _gatewayRegistry;
     private readonly IGateway _reference;
 
-    // The local mapping of agents to worker processes.
-    private readonly ConcurrentDictionary<GrpcWorkerConnection, GrpcWorkerConnection> _workers = new();
     // The agents supported by each worker process.
     private readonly ConcurrentDictionary<string, List<GrpcWorkerConnection>> _supportedAgentTypes = [];
     // The mapping from agent id to worker process.
@@ -33,7 +31,7 @@ public GrpcGateway(IClusterClient clusterClient, ILogger<Gateway> logger) : base
         _gatewayRegistry = clusterClient.GetGrain<IAgentRegistry>(0);
     }
     //intetionally not static so can be called by some methods implemented in base class
-    private async ValueTask SendMessageAsync(IConnection connection, CloudEvent cloudEvent, CancellationToken cancellationToken = default)
+    public override async Task SendMessageAsync(IConnection connection, CloudEvent cloudEvent, CancellationToken cancellationToken = default)
     {
         var queue = (GrpcWorkerConnection)connection;
         await queue.ResponseStream.WriteAsync(new Message { CloudEvent = cloudEvent }, cancellationToken).ConfigureAwait(false);