Repository: microsoft/autogen
Pull request: fix: Recurrent SendMessageAsync deadlock in message handler
In the .NET InProcessRuntime we tried to minimize the amount of tasks created. Unfortunately, this results in a deadlock when a send message handler is attempting to SendMessage during the handling of the incoming message.

The fix is to create a new task for delivering the message in the message processing loop, which creates a new logical stack to run the delivery, freeing the loop to process the next delivery request.

* Also fixes passing exceptions and cancellations back to the waiting task.
* Also fixes a build slowdown on Windows due to uv and llama-cpp

Closes #5915

Problem statement:
Deadlock in logically recurrent calls to SendMessageAsync
### Confirmation

- [x] I confirm that I am a maintainer and so can use this template. If I am not, I understand this issue will be closed and I will be asked to use a different template.

### Issue body

If a message handler invoked during SendMessage then attempts to send a message and await its result, the InProcessRuntime will deadlock.

Hints from issue discussion:

Commit message: fix: Recurrent SendMessageAsync deadlock in message handler

In the .NET InProcessRuntime we tried to minimize the amount of tasks created. Unfortunately, this results in a deadlock when a send message handler is attempting to SendMessage during the handling of the incoming message.

The fix is to create a new task for delivering the message in the message processing loop, which creates a new logical stack to run the delivery, freeing the loop to process the next delivery request.

* Also fixes passing exceptions and cancellations back to the waiting task.
* Also fixes a build slowdown on Windows due to uv and llama-cpp

Patch to review:
diff --git a/dotnet/src/Microsoft.AutoGen/Core/InProcessRuntime.cs b/dotnet/src/Microsoft.AutoGen/Core/InProcessRuntime.cs
index bf18b1da2c35..3b2c472a555a 100644
--- a/dotnet/src/Microsoft.AutoGen/Core/InProcessRuntime.cs
+++ b/dotnet/src/Microsoft.AutoGen/Core/InProcessRuntime.cs
@@ -3,6 +3,7 @@
 
 using System.Collections.Concurrent;
 using System.Diagnostics;
+using System.Reflection;
 using System.Text.Json;
 using Microsoft.AutoGen.Contracts;
 using Microsoft.Extensions.Hosting;
@@ -104,14 +105,29 @@ public ValueTask PublishMessageAsync(object message, TopicId topic, AgentId? sen
 
     public ValueTask<object?> SendMessageAsync(object message, AgentId recepient, AgentId? sender = null, string? messageId = null, CancellationToken cancellationToken = default)
     {
-        return this.ExecuteTracedAsync(() =>
+        return this.ExecuteTracedAsync(async () =>
         {
             MessageDelivery delivery = new MessageEnvelope(message, messageId, cancellationToken)
                                             .WithSender(sender)
                                             .ForSend(recepient, this.SendMessageServicer);
 
             this.messageDeliveryQueue.Enqueue(delivery);
-            return delivery.Future;
+
+            try
+            {
+                return await delivery.Future;
+            }
+            catch (TargetInvocationException ex) when (ex.InnerException is OperationCanceledException innerOCEx)
+            {
+                if (!cancellationToken.IsCancellationRequested)
+                {
+                    // If the parent token was not cancelled, we do not have the true cause of the cancellation.
+                    // Pass an arbitrary cancelled token through the exception.
+                    cancellationToken = new CancellationToken(canceled: true);
+                }
+
+                throw new OperationCanceledException($"Delivery of message {messageId} was cancelled.", innerOCEx, cancellationToken);
+            }
         });
     }
 
@@ -242,7 +258,7 @@ public ValueTask<AgentProxy> TryGetAgentProxyAsync(AgentId agentId)
         return ValueTask.FromResult(new AgentProxy(agentId, this));
     }
 
-    public ValueTask ProcessNextMessage(CancellationToken cancellation = default)
+    public ValueTask ProcessNextMessageAsync(CancellationToken cancellation = default)
     {
         Debug.WriteLine("Processing next message...");
         if (this.messageDeliveryQueue.TryDequeue(out MessageDelivery? delivery))
@@ -260,7 +276,7 @@ private async Task RunAsync(CancellationToken cancellation)
     {
         while (!cancellation.IsCancellationRequested && shouldContinue())
         {
-            await this.ProcessNextMessage(cancellation);
+            _ = Task.Run(async () => await this.ProcessNextMessageAsync(cancellation));
         }
 
         await this.FinishAsync(this.finishSource?.Token ?? CancellationToken.None);
diff --git a/dotnet/src/Microsoft.AutoGen/Core/MessageDelivery.cs b/dotnet/src/Microsoft.AutoGen/Core/MessageDelivery.cs
index ab97dac84e79..4ed9f0b0df67 100644
--- a/dotnet/src/Microsoft.AutoGen/Core/MessageDelivery.cs
+++ b/dotnet/src/Microsoft.AutoGen/Core/MessageDelivery.cs
@@ -48,8 +48,19 @@ public MessageDelivery ForSend(AgentId receiver, Func<MessageEnvelope, Cancellat
         ResultSink<object?> resultSink = new ResultSink<object?>();
         Func<MessageEnvelope, CancellationToken, ValueTask> boundServicer = async (envelope, cancellation) =>
         {
-            object? result = await servicer(envelope, cancellation);
-            resultSink.SetResult(result);
+            try
+            {
+                object? result = await servicer(envelope, cancellation);
+                resultSink.SetResult(result);
+            }
+            catch (OperationCanceledException)
+            {
+                resultSink.SetCancelled();
+            }
+            catch (Exception ex)
+            {
+                resultSink.SetException(ex);
+            }
         };
 
         return new MessageDelivery(this, boundServicer, resultSink);
diff --git a/dotnet/test/Microsoft.AutoGen.Core.Tests/SendMessageTests.cs b/dotnet/test/Microsoft.AutoGen.Core.Tests/SendMessageTests.cs
new file mode 100644
index 000000000000..a921164b196c
--- /dev/null
+++ b/dotnet/test/Microsoft.AutoGen.Core.Tests/SendMessageTests.cs
@@ -0,0 +1,232 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// SendMessageTests.cs
+
+using System.Diagnostics;
+using System.Reflection;
+using FluentAssertions;
+using Microsoft.AutoGen.Contracts;
+using Microsoft.Extensions.Logging;
+using Xunit;
+
+namespace Microsoft.AutoGen.Core.Tests;
+
+[Trait("Category", "UnitV2")]
+public class SendMessageTests
+{
+    private sealed class BasicMessage
+    {
+        public string Content { get; set; } = string.Empty;
+    }
+
+    private sealed class SendOnAgent : BaseAgent, IHandle<BasicMessage>
+    {
+        private IList<Guid> targetKeys;
+
+        public SendOnAgent(AgentId id, IAgentRuntime runtime, string description, IList<Guid> targetKeys, ILogger<BaseAgent>? logger = null)
+            : base(id, runtime, description, logger)
+        {
+            this.targetKeys = targetKeys;
+        }
+
+        public async ValueTask HandleAsync(BasicMessage item, MessageContext messageContext)
+        {
+            foreach (Guid targetKey in targetKeys)
+            {
+                AgentId targetId = new(nameof(ReceiverAgent), targetKey.ToString());
+                BasicMessage message = new BasicMessage { Content = $"@{targetKey}: item.Content" };
+                await this.Runtime.SendMessageAsync(message, targetId);
+            }
+        }
+    }
+
+    private sealed class ReceiverAgent : BaseAgent, IHandle<BasicMessage>
+    {
+        public List<BasicMessage> Messages { get; } = new();
+
+        public ReceiverAgent(AgentId id, IAgentRuntime runtime, string description, ILogger<BaseAgent>? logger = null)
+            : base(id, runtime, description, logger)
+        {
+        }
+
+        public ValueTask HandleAsync(BasicMessage item, MessageContext messageContext)
+        {
+            this.Messages.Add(item);
+
+            return ValueTask.CompletedTask;
+        }
+    }
+
+    private sealed class ProcessorAgent : BaseAgent, IHandle<BasicMessage, BasicMessage>
+    {
+        private Func<string, string> ProcessFunc { get; }
+
+        public ProcessorAgent(AgentId id, IAgentRuntime runtime, Func<string, string> processFunc, string description, ILogger<BaseAgent>? logger = null)
+            : base(id, runtime, description, logger)
+        {
+            this.ProcessFunc = processFunc;
+        }
+
+        public ValueTask<BasicMessage> HandleAsync(BasicMessage item, MessageContext messageContext)
+        {
+            BasicMessage result = new() { Content = this.ProcessFunc(item.Content) };
+
+            return ValueTask.FromResult(result);
+        }
+    }
+
+    private sealed class TestException : Exception
+    { }
+
+    private sealed class CancelAgent : BaseAgent, IHandle<BasicMessage, BasicMessage>
+    {
+        public CancelAgent(AgentId id, IAgentRuntime runtime, string description, ILogger<BaseAgent>? logger = null)
+            : base(id, runtime, description, logger)
+        {
+        }
+
+        public ValueTask<BasicMessage> HandleAsync(BasicMessage item, MessageContext messageContext)
+        {
+            CancellationToken cancelledToken = new CancellationToken(canceled: true);
+            cancelledToken.ThrowIfCancellationRequested();
+
+            return ValueTask.FromResult(item);
+        }
+    }
+
+    private sealed class ErrorAgent : BaseAgent, IHandle<BasicMessage, BasicMessage>
+    {
+        public ErrorAgent(AgentId id, IAgentRuntime runtime, string description, ILogger<BaseAgent>? logger = null)
+            : base(id, runtime, description, logger)
+        {
+        }
+
+        public ValueTask<BasicMessage> HandleAsync(BasicMessage item, MessageContext messageContext)
+        {
+            throw new TestException();
+        }
+    }
+
+    private sealed class SendTestFixture
+    {
+        private Dictionary<Type, object> AgentsTypeMap { get; } = new();
+        public InProcessRuntime Runtime { get; private set; } = new();
+
+        public ValueTask<AgentType> RegisterFactoryMapInstances<TAgent>(AgentType type, Func<AgentId, IAgentRuntime, ValueTask<TAgent>> factory)
+            where TAgent : IHostableAgent
+        {
+            Func<AgentId, IAgentRuntime, ValueTask<TAgent>> wrappedFactory = async (id, runtime) =>
+            {
+                TAgent agent = await factory(id, runtime);
+                this.GetAgentInstances<TAgent>()[id] = agent;
+                return agent;
+            };
+
+            return this.Runtime.RegisterAgentFactoryAsync(type, wrappedFactory);
+        }
+
+        public Dictionary<AgentId, TAgent> GetAgentInstances<TAgent>() where TAgent : IHostableAgent
+        {
+            if (!AgentsTypeMap.TryGetValue(typeof(TAgent), out object? maybeAgentMap) ||
+                maybeAgentMap is not Dictionary<AgentId, TAgent> result)
+            {
+                this.AgentsTypeMap[typeof(TAgent)] = result = new Dictionary<AgentId, TAgent>();
+            }
+
+            return result;
+        }
+
+        public async ValueTask<object?> RunTestAsync(AgentId sendTarget, object message, string? messageId = null)
+        {
+            messageId ??= Guid.NewGuid().ToString();
+
+            await this.Runtime.StartAsync();
+
+            object? result = await this.Runtime.SendMessageAsync(message, sendTarget, messageId: messageId);
+
+            await this.Runtime.RunUntilIdleAsync();
+
+            return result;
+        }
+    }
+
+    [Fact]
+    public async Task Test_SendMessage_ReturnsValue()
+    {
+        Func<string, string> ProcessFunc = (s) => $"Processed({s})";
+
+        SendTestFixture fixture = new SendTestFixture();
+
+        await fixture.RegisterFactoryMapInstances(nameof(ProcessorAgent),
+            (id, runtime) => ValueTask.FromResult(new ProcessorAgent(id, runtime, ProcessFunc, string.Empty)));
+
+        AgentId targetAgent = new AgentId(nameof(ProcessorAgent), Guid.NewGuid().ToString());
+        object? maybeResult = await fixture.RunTestAsync(targetAgent, new BasicMessage { Content = "1" });
+
+        maybeResult.Should().NotBeNull()
+                        .And.BeOfType<BasicMessage>()
+                        .And.Match<BasicMessage>(m => m.Content == "Processed(1)");
+    }
+
+    [Fact]
+    public async Task Test_SendMessage_Cancellation()
+    {
+        SendTestFixture fixture = new SendTestFixture();
+
+        await fixture.RegisterFactoryMapInstances(nameof(CancelAgent),
+            (id, runtime) => ValueTask.FromResult(new CancelAgent(id, runtime, string.Empty)));
+
+        AgentId targetAgent = new AgentId(nameof(CancelAgent), Guid.NewGuid().ToString());
+        Func<Task> testAction = () => fixture.RunTestAsync(targetAgent, new BasicMessage { Content = "1" }).AsTask();
+
+        // TODO: Do we want to do the unwrap in this case?
+        await testAction.Should().ThrowAsync<OperationCanceledException>();
+    }
+
+    [Fact]
+    public async Task Test_SendMessage_Error()
+    {
+        SendTestFixture fixture = new SendTestFixture();
+
+        await fixture.RegisterFactoryMapInstances(nameof(ErrorAgent),
+            (id, runtime) => ValueTask.FromResult(new ErrorAgent(id, runtime, string.Empty)));
+
+        AgentId targetAgent = new AgentId(nameof(ErrorAgent), Guid.NewGuid().ToString());
+        Func<Task> testAction = () => fixture.RunTestAsync(targetAgent, new BasicMessage { Content = "1" }).AsTask();
+
+        (await testAction.Should().ThrowAsync<TargetInvocationException>())
+                                  .WithInnerException<TestException>();
+    }
+
+    [Fact]
+    public async Task TesT_SendMessage_FromSendMessageHandler()
+    {
+        Guid[] targetGuids = [Guid.NewGuid(), Guid.NewGuid()];
+
+        SendTestFixture fixture = new SendTestFixture();
+
+        Dictionary<AgentId, SendOnAgent> sendAgents = fixture.GetAgentInstances<SendOnAgent>();
+        Dictionary<AgentId, ReceiverAgent> receiverAgents = fixture.GetAgentInstances<ReceiverAgent>();
+
+        await fixture.RegisterFactoryMapInstances(nameof(SendOnAgent),
+            (id, runtime) => ValueTask.FromResult(new SendOnAgent(id, runtime, string.Empty, targetGuids)));
+
+        await fixture.RegisterFactoryMapInstances(nameof(ReceiverAgent),
+            (id, runtime) => ValueTask.FromResult(new ReceiverAgent(id, runtime, string.Empty)));
+
+        AgentId targetAgent = new AgentId(nameof(SendOnAgent), Guid.NewGuid().ToString());
+        Task testTask = fixture.RunTestAsync(targetAgent, new BasicMessage { Content = "Hello" }).AsTask();
+
+        TimeSpan timeout = Debugger.IsAttached ? TimeSpan.FromSeconds(10) : TimeSpan.FromSeconds(1);
+        Task timeoutTask = Task.Delay(timeout);
+
+        Task completedTask = await Task.WhenAny([testTask, timeoutTask]);
+        completedTask.Should().Be(testTask, "SendOnAgent should complete before timeout");
+
+        // Check that each of the target agents received the message
+        foreach (var targetKey in targetGuids)
+        {
+            var targetId = new AgentId(nameof(ReceiverAgent), targetKey.ToString());
+            receiverAgents[targetId].Messages.Should().ContainSingle(m => m.Content == $"@{targetKey}: item.Content");
+        }
+    }
+}
diff --git a/dotnet/test/Microsoft.AutoGen.Integration.Tests/Microsoft.AutoGen.Integration.Tests.csproj b/dotnet/test/Microsoft.AutoGen.Integration.Tests/Microsoft.AutoGen.Integration.Tests.csproj
index 661c49d3ba47..9ddc467f460b 100644
--- a/dotnet/test/Microsoft.AutoGen.Integration.Tests/Microsoft.AutoGen.Integration.Tests.csproj
+++ b/dotnet/test/Microsoft.AutoGen.Integration.Tests/Microsoft.AutoGen.Integration.Tests.csproj
@@ -56,7 +56,7 @@
     </PropertyGroup>  
     <Message Importance="Normal" Text="Initializing virtual environment for $(PythonVenvRoot)" />
 
-    <Exec Command="$(UvPath) sync --all-extras" WorkingDirectory="$(PythonRoot)" Condition=" '$(CreateVenv)' == 'True' " />
+    <Exec Command="$(UvPath) sync --all-extras --no-install-package llama-cpp-python" WorkingDirectory="$(PythonRoot)" Condition=" '$(CreateVenv)' == 'True' " />
       <!--Output TaskParameter="ConsoleOutput" PropertyName="OutputOfExec" />
     </Exec>
     <Message Importance="Normal" Text="$(OutputOfExec)" />-->
@@ -64,3 +64,4 @@
 
 
 </Project>
+  
\ No newline at end of file


Code review comments:
path: dotnet/src/Microsoft.AutoGen/Core/InProcessRuntime.cs
text: @user1:
With the change to multiple concurrent tasks which may send messages resulting in new agent instances, it's probably a good idea to make `agentInstances` a `ConcurrentDictionary` and switch to `TryAdd(...)` (line 144).

@author:
This is a good point, but fixing it right is a bit more involved than making the Dictionaries concurrent, due to interaction between the `agentFactories` and `agentInstances` dictionaries in `EnsureAgentAsync()`, so I will pull it out into a separate change