Allow InProcessRuntime subscriptions to be processed concurrently (#179)

* Allow InProcessRuntime subscriptions to be processed concurrently

* Fix tests
This commit is contained in:
Stephen Toub
2025-07-11 21:50:51 -04:00
committed by GitHub
Unverified
parent 18299bc701
commit 80c1e2ee0a
2 changed files with 11 additions and 20 deletions
@@ -262,21 +262,22 @@ public sealed partial class InProcessRuntime : IAgentRuntime, IAsyncDisposable
{
Debug.Assert(message.Topic.HasValue);
List<Exception>? exceptions = null;
List<Task>? tasks = null;
TopicId topic = message.Topic!.Value;
foreach (KeyValuePair<string, ISubscriptionDefinition> 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
@@ -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<AggregateException>(async () => await fixture.RunPublishTestAsync(new TopicId("TestTopic"), new BasicMessage { Content = "1" }));
Assert.IsType<TestException>(Assert.Single(e.InnerExceptions));
await Assert.ThrowsAsync<TestException>(async () => await fixture.RunPublishTestAsync(new TopicId("TestTopic"), new BasicMessage { Content = "1" }));
var values = fixture.GetAgentInstances<ReceiverAgent>().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<AggregateException>(async () => await fixture.RunPublishTestAsync(new TopicId("TestTopic"), new BasicMessage { Content = "1" }));
Assert.Equal(2, e.InnerExceptions.Count);
Assert.All(e.InnerExceptions, innerException => Assert.IsType<TestException>(innerException));
await Assert.ThrowsAsync<TestException>(async () => await fixture.RunPublishTestAsync(new TopicId("TestTopic"), new BasicMessage { Content = "1" }));
var values = fixture.GetAgentInstances<ErrorAgent>().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<AggregateException>(async () => await fixture.RunPublishTestAsync(new TopicId("TestTopic"), new BasicMessage { Content = "1" }));
Assert.Equal(2, e.InnerExceptions.Count);
Assert.All(e.InnerExceptions, innerException => Assert.IsType<TestException>(innerException));
await Assert.ThrowsAsync<TestException>(async () => await fixture.RunPublishTestAsync(new TopicId("TestTopic"), new BasicMessage { Content = "1" }));
var agents = fixture.GetAgentInstances<ReceiverAgent>().Values;
Assert.Equal(2, agents.Count);