From 80c1e2ee0a3c972afdf4784b40b942bd21dad12d Mon Sep 17 00:00:00 2001 From: Stephen Toub Date: Fri, 11 Jul 2025 21:50:51 -0400 Subject: [PATCH] Allow InProcessRuntime subscriptions to be processed concurrently (#179) * Allow InProcessRuntime subscriptions to be processed concurrently * Fix tests --- .../InProcessRuntime.cs | 19 ++++++++----------- .../PublishMessageTests.cs | 12 +++--------- 2 files changed, 11 insertions(+), 20 deletions(-) diff --git a/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Abstractions/InProcessRuntime.cs b/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Abstractions/InProcessRuntime.cs index ffa59c48e5..a947860759 100644 --- a/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Abstractions/InProcessRuntime.cs +++ b/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Abstractions/InProcessRuntime.cs @@ -262,21 +262,22 @@ public sealed partial class InProcessRuntime : IAgentRuntime, IAsyncDisposable { Debug.Assert(message.Topic.HasValue); - List? exceptions = null; + List? tasks = null; TopicId topic = message.Topic!.Value; foreach (KeyValuePair subscription in message.Runtime._subscriptions) { - if (!subscription.Value.Matches(topic)) + if (subscription.Value.Matches(topic)) { - continue; + (tasks ??= []).Add(ProcessSubscriptionAsync(message, subscription.Value, topic, cancellationToken)); } - try + static async Task ProcessSubscriptionAsync( + MessageToProcess message, ISubscriptionDefinition subscription, TopicId topic, CancellationToken cancellationToken) { using CancellationTokenSource combinedSource = CancellationTokenSource.CreateLinkedTokenSource(message.Cancellation, cancellationToken); combinedSource.Token.ThrowIfCancellationRequested(); - ActorId actorId = subscription.Value.MapToActor(topic); + ActorId actorId = subscription.MapToActor(topic); ActorId? sender = message.Sender; if (sender is null || sender != actorId) { @@ -289,15 +290,11 @@ public sealed partial class InProcessRuntime : IAgentRuntime, IAsyncDisposable }, combinedSource.Token).ConfigureAwait(false); } } - catch (Exception ex) - { - (exceptions ??= []).Add(ex); - } } - if (exceptions is not null) + if (tasks is not null) { - throw new AggregateException("One or more exceptions occurred while processing the message.", exceptions); + await Task.WhenAll(tasks).ConfigureAwait(false); } // This method is effectively void, with the result never being used. But it's typed the same as SendMessageServicerAsync diff --git a/dotnet/tests/Microsoft.Extensions.AI.Agents.Runtime.Abstractions.UnitTests/PublishMessageTests.cs b/dotnet/tests/Microsoft.Extensions.AI.Agents.Runtime.Abstractions.UnitTests/PublishMessageTests.cs index 4f03c1e962..75a369de35 100644 --- a/dotnet/tests/Microsoft.Extensions.AI.Agents.Runtime.Abstractions.UnitTests/PublishMessageTests.cs +++ b/dotnet/tests/Microsoft.Extensions.AI.Agents.Runtime.Abstractions.UnitTests/PublishMessageTests.cs @@ -1,6 +1,5 @@ // Copyright (c) Microsoft. All rights reserved. -using System; using System.Threading.Tasks; namespace Microsoft.Extensions.AI.Agents.Runtime.InProcess.Tests; @@ -35,8 +34,7 @@ public class PublishMessageTests await fixture.RegisterErrorAgentAsync(topicTypes: "TestTopic"); // Test that we wrap single errors appropriately - var e = await Assert.ThrowsAsync(async () => await fixture.RunPublishTestAsync(new TopicId("TestTopic"), new BasicMessage { Content = "1" })); - Assert.IsType(Assert.Single(e.InnerExceptions)); + await Assert.ThrowsAsync(async () => await fixture.RunPublishTestAsync(new TopicId("TestTopic"), new BasicMessage { Content = "1" })); var values = fixture.GetAgentInstances().Values; } @@ -50,9 +48,7 @@ public class PublishMessageTests await fixture.RegisterErrorAgentAsync("2", topicTypes: "TestTopic"); // What we are really testing here is that a single exception does not prevent sending to the remaining agents - var e = await Assert.ThrowsAsync(async () => await fixture.RunPublishTestAsync(new TopicId("TestTopic"), new BasicMessage { Content = "1" })); - Assert.Equal(2, e.InnerExceptions.Count); - Assert.All(e.InnerExceptions, innerException => Assert.IsType(innerException)); + await Assert.ThrowsAsync(async () => await fixture.RunPublishTestAsync(new TopicId("TestTopic"), new BasicMessage { Content = "1" })); var values = fixture.GetAgentInstances().Values; Assert.Equal(2, values.Count); @@ -70,9 +66,7 @@ public class PublishMessageTests await fixture.RegisterErrorAgentAsync("2", topicTypes: "TestTopic"); // What we are really testing here is that raising exceptions does not prevent sending to the remaining agents - var e = await Assert.ThrowsAsync(async () => await fixture.RunPublishTestAsync(new TopicId("TestTopic"), new BasicMessage { Content = "1" })); - Assert.Equal(2, e.InnerExceptions.Count); - Assert.All(e.InnerExceptions, innerException => Assert.IsType(innerException)); + await Assert.ThrowsAsync(async () => await fixture.RunPublishTestAsync(new TopicId("TestTopic"), new BasicMessage { Content = "1" })); var agents = fixture.GetAgentInstances().Values; Assert.Equal(2, agents.Count);