add SendMessageAsync

microsoft · Nov 12, 2024 · 55147a9 · 55147a9
1 parent 3530a99
commit 55147a9
Show file tree

Hide file tree

Showing 6 changed files with 21 additions and 24 deletions.
diff --git a/dotnet/src/Microsoft.AutoGen/Abstractions/IAgentRuntime.cs b/dotnet/src/Microsoft.AutoGen/Abstractions/IAgentRuntime.cs
@@ -13,6 +13,7 @@ public interface IAgentRuntime
     ValueTask<AgentState> ReadAsync(AgentId agentId, CancellationToken cancellationToken = default);
     ValueTask SendResponseAsync(RpcRequest request, RpcResponse response, CancellationToken cancellationToken = default);
     ValueTask SendRequestAsync(IAgentBase agent, RpcRequest request, CancellationToken cancellationToken = default);
+    ValueTask SendMessageAsync(Message message, CancellationToken cancellationToken = default);
     ValueTask PublishEventAsync(CloudEvent @event, CancellationToken cancellationToken = default);
     void Update(Activity? activity, RpcRequest request);
     void Update(Activity? activity, CloudEvent cloudEvent);

diff --git a/dotnet/src/Microsoft.AutoGen/Abstractions/IAgentWorker.cs b/dotnet/src/Microsoft.AutoGen/Abstractions/IAgentWorker.cs
@@ -8,6 +8,7 @@ public interface IAgentWorker
     ValueTask PublishEventAsync(CloudEvent evt, CancellationToken cancellationToken = default);
     ValueTask SendRequestAsync(IAgentBase agent, RpcRequest request, CancellationToken cancellationToken = default);
     ValueTask SendResponseAsync(RpcResponse response, CancellationToken cancellationToken = default);
+    ValueTask SendMessageAsync(Message message, CancellationToken cancellationToken = default);
     ValueTask StoreAsync(AgentState value, CancellationToken cancellationToken = default);
     ValueTask<AgentState> ReadAsync(AgentId agentId, CancellationToken cancellationToken = default);
 }
diff --git a/dotnet/src/Microsoft.AutoGen/Agents/AgentRuntime.cs b/dotnet/src/Microsoft.AutoGen/Agents/AgentRuntime.cs
@@ -45,6 +45,10 @@ public async ValueTask SendRequestAsync(IAgentBase agent, RpcRequest request, Ca
     {
         await worker.SendRequestAsync(agent, request, cancellationToken).ConfigureAwait(false);
     }
+    public async ValueTask SendMessageAsync(Message message, CancellationToken cancellationToken = default)
+    {
+        await worker.SendMessageAsync(message, cancellationToken).ConfigureAwait(false);
+    }
     public async ValueTask PublishEventAsync(CloudEvent @event, CancellationToken cancellationToken = default)
     {
         await worker.PublishEventAsync(@event, cancellationToken).ConfigureAwait(false);

diff --git a/dotnet/src/Microsoft.AutoGen/Agents/Services/AgentWorker.cs b/dotnet/src/Microsoft.AutoGen/Agents/Services/AgentWorker.cs
@@ -63,6 +63,10 @@ public ValueTask SendResponseAsync(RpcResponse response, CancellationToken cance
     {
         return _mailbox.Writer.WriteAsync(new Message { Response = response }, cancellationToken);
     }
+    public ValueTask SendMessageAsync(Message message, CancellationToken cancellationToken = default)
+    {
+        return _mailbox.Writer.WriteAsync(message, cancellationToken);
+    }
     public ValueTask StoreAsync(AgentState value, CancellationToken cancellationToken = default)
     {
         var agentId = value.AgentId ?? throw new InvalidOperationException("AgentId is required when saving AgentState.");

diff --git a/dotnet/src/Microsoft.AutoGen/Agents/Services/Grpc/GrpcAgentWorker.cs b/dotnet/src/Microsoft.AutoGen/Agents/Services/Grpc/GrpcAgentWorker.cs
@@ -237,6 +237,11 @@ await WriteChannelAsync(new Message
         await WriteChannelAsync(new Message { Request = request }, cancellationToken).ConfigureAwait(false);
     }
     // new is intentional
+    public new async ValueTask SendMessageAsync(Message message, CancellationToken cancellationToken = default)
+    {
+        await WriteChannelAsync(message, cancellationToken).ConfigureAwait(false);
+    }
+    // new is intentional
     public new async ValueTask PublishEventAsync(CloudEvent @event, CancellationToken cancellationToken = default)
     {
         await WriteChannelAsync(new Message { CloudEvent = @event }, cancellationToken).ConfigureAwait(false);

diff --git a/dotnet/test/Microsoft.AutoGen.Agents.Tests/GrpcGatewayTests.cs b/dotnet/test/Microsoft.AutoGen.Agents.Tests/GrpcGatewayTests.cs
@@ -28,26 +28,7 @@ private async ValueTask Setup()
     [Fact]
     public async Task AddSubscriptionAsync_SendsAddSubscriptionRequest_AndChecksAddSubscriptionResponse()
     {
-        // Arrange
-        var request = new AddSubscriptionRequest
-        {
-            RequestId = "test-request-id",
-            Subscription = new Subscription
-            {
-                TypeSubscription = new TypeSubscription
-                {
-                    TopicType = "test-topic",
-                    AgentType = "test-agent-type"
-                }
-            }
-        };
-
-
-        var responseStream = new Mock<IServerStreamWriter<Message>>();
-        _mockConnection.SetupGet(c => c.ResponseStream).Returns(responseStream.Object);
 
-        // Act
-        await _grpcGateway.AddSubscriptionAsync(_mockConnection.Object, request);
 
         // Assert
         responseStream.Verify(stream => stream.WriteAsync(It.Is<Message>(msg =>
@@ -64,13 +45,14 @@ public class TestAgent(
     {
         public async Task Handle(NewMessageReceived item)
         {
-
+            // update our subscription requests
+            await SendSubscriptionRequestAsync().ConfigureAwait(false);
         }
-        private async SendSubscriptionRequestAsync()
+        private async ValueTask SendSubscriptionRequestAsync(CancellationToken cancellationToken = default)
         {
-            RpcRequest request = new()
+            Message request = new()
             {
-                Payload = new AddSubscriptionRequest()
+                AddSubscriptionRequest = new()
                 {
                     RequestId = "test-request-id",
                     Subscription = new Subscription
@@ -83,7 +65,7 @@ private async SendSubscriptionRequestAsync()
                     }
                 }
             };
-            return this.Context.SendRequestAsync(this,request);
+            await Context.SendMessageAsync(request, cancellationToken).ConfigureAwait(false);
         }
     }
 }