First round of cleanup of runtime abstractions (#156)

dotnet/tests/Microsoft.Extensions.AI.Agents.Runtime.InProcess.UnitTests/InProcessRuntimeTests.cs

@@ -4,6 +4,7 @@ using System;
 using System.Collections.Generic;
 using System.Linq;
 using System.Text.Json;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace Microsoft.Extensions.AI.Agents.Runtime.InProcess.Tests;
@@ -18,7 +19,7 @@ public class InProcessRuntimeTests()
 
         // Assert
         Assert.False(runtime.DeliverToSelf);
-        Assert.Equal(0, runtime.messageQueueCount);
+        Assert.Equal(0, runtime._messageQueueCount);
 
         // Act
         await runtime.StopAsync(); // Already stopped
@@ -29,13 +30,13 @@ public class InProcessRuntimeTests()
         // Assert
         // Invalid to start runtime that is already started
         await Assert.ThrowsAsync<InvalidOperationException>(() => runtime.StartAsync());
-        Assert.Equal(0, runtime.messageQueueCount);
+        Assert.Equal(0, runtime._messageQueueCount);
 
         // Act
         await runtime.StopAsync();
 
         // Assert
-        Assert.Equal(0, runtime.messageQueueCount);
+        Assert.Equal(0, runtime._messageQueueCount);
     }
 
     [Fact]
@@ -43,7 +44,7 @@ public class InProcessRuntimeTests()
     {
         // Arrange
         await using InProcessRuntime runtime = new();
-        TestSubscription subscription = new("TestTopic", "MyAgent");
+        TestSubscription subscription = new("TestTopic", new("MyAgent"));
 
         // Act & Assert
         await Assert.ThrowsAsync<InvalidOperationException>(async () => await runtime.RemoveSubscriptionAsync(subscription.Id));
@@ -68,44 +69,44 @@ public class InProcessRuntimeTests()
         await using InProcessRuntime runtime = new();
 
         // Act & Assert
-        await Assert.ThrowsAsync<InvalidOperationException>(async () => await runtime.GetAgentAsync(AgentType, lazy: false));
+        await Assert.ThrowsAsync<InvalidOperationException>(async () => await runtime.GetActorAsync(AgentType, lazy: false));
 
         // Arrange
-        await runtime.RegisterAgentFactoryAsync(AgentType, factoryFunc);
+        await runtime.RegisterActorFactoryAsync(new(AgentType), factoryFunc);
 
         // Act & Assert
-        await Assert.ThrowsAsync<InvalidOperationException>(async () => await runtime.RegisterAgentFactoryAsync(AgentType, factoryFunc));
+        await Assert.ThrowsAsync<InvalidOperationException>(async () => await runtime.RegisterActorFactoryAsync(new(AgentType), factoryFunc));
 
         // Act: Lookup by type
-        AgentId agentId = await runtime.GetAgentAsync(AgentType, lazy: false);
+        ActorId agentId = await runtime.GetActorAsync(AgentType, lazy: false);
 
         // Assert
         Assert.Single(agents);
-        Assert.Single(runtime.agentInstances);
+        Assert.Single(runtime._actorInstances);
 
         // Act
-        MockAgent agent = await runtime.TryGetUnderlyingAgentInstanceAsync<MockAgent>(agentId);
+        MockAgent agent = await runtime.TryGetUnderlyingActorInstanceAsync<MockAgent>(agentId);
 
         // Assert
         Assert.Equal(agentId, agent.Id);
 
         // Act & Assert
-        await Assert.ThrowsAsync<InvalidOperationException>(async () => await runtime.TryGetUnderlyingAgentInstanceAsync<WrongAgent>(agentId));
+        await Assert.ThrowsAsync<InvalidOperationException>(async () => await runtime.TryGetUnderlyingActorInstanceAsync<WrongAgent>(agentId));
 
         // Act: Lookup by ID
-        AgentId sameId = await runtime.GetAgentAsync(agentId, lazy: false);
+        ActorId sameId = await runtime.GetActorAsync(agentId, lazy: false);
 
         // Assert
         Assert.Equal(agentId, sameId);
 
         // Act: Lookup by Type
-        sameId = await runtime.GetAgentAsync((AgentType)agent.Id.Type, lazy: false);
+        sameId = await runtime.GetActorAsync((ActorType)agent.Id.Type, lazy: false);
 
         // Assert
         Assert.Equal(agentId, sameId);
 
         // Act: Lookup metadata
-        AgentMetadata metadata = await runtime.GetAgentMetadataAsync(agentId);
+        ActorMetadata metadata = await runtime.GetActorMetadataAsync(agentId);
 
         // Assert
         Assert.Equal(agentId.Type, metadata.Type);
@@ -113,15 +114,16 @@ public class InProcessRuntimeTests()
         Assert.Equal(agentId.Key, metadata.Key);
 
         // Act: Access proxy
-        AgentProxy proxy = await runtime.TryGetAgentProxyAsync(agentId);
+        IdProxyActor? proxy = await runtime.TryGetActorProxyAsync(agentId);
 
         // Assert
+        Assert.NotNull(proxy);
         Assert.Equal(agentId, proxy.Id);
         Assert.Equal(metadata.Type, proxy.Metadata.Type);
         Assert.Equal(metadata.Description, proxy.Metadata.Description);
         Assert.Equal(metadata.Key, proxy.Metadata.Key);
 
-        async ValueTask<MockAgent> factoryFunc(AgentId id, IAgentRuntime runtime)
+        async ValueTask<MockAgent> factoryFunc(ActorId id, IAgentRuntime runtime)
         {
             MockAgent agent = new(id, runtime, AgentDescription);
             agents.Add(agent);
@@ -137,35 +139,35 @@ public class InProcessRuntimeTests()
         const string TestMessage = "test message";
 
         await using InProcessRuntime firstRuntime = new();
-        await firstRuntime.RegisterAgentFactoryAsync(AgentType, factoryFunc);
+        await firstRuntime.RegisterActorFactoryAsync(new(AgentType), factoryFunc);
 
         // Act
-        AgentId agentId = await firstRuntime.GetAgentAsync(AgentType, lazy: false);
+        ActorId agentId = await firstRuntime.GetActorAsync(AgentType, lazy: false);
 
         // Assert
-        Assert.Single(firstRuntime.agentInstances);
+        Assert.Single(firstRuntime._actorInstances);
 
         // Arrange
-        MockAgent agent = (MockAgent)firstRuntime.agentInstances[agentId];
+        MockAgent agent = (MockAgent)firstRuntime._actorInstances[agentId];
         agent.ReceivedMessages.Add(TestMessage);
 
         // Act
-        JsonElement agentState = await firstRuntime.SaveAgentStateAsync(agentId);
+        JsonElement agentState = await firstRuntime.SaveActorStateAsync(agentId);
 
         // Arrange
         await using InProcessRuntime secondRuntime = new();
-        await secondRuntime.RegisterAgentFactoryAsync(AgentType, factoryFunc);
+        await secondRuntime.RegisterActorFactoryAsync(new(AgentType), factoryFunc);
 
         // Act
-        await secondRuntime.LoadAgentStateAsync(agentId, agentState);
+        await secondRuntime.LoadActorStateAsync(agentId, agentState);
 
         // Assert
-        Assert.Single(secondRuntime.agentInstances);
-        MockAgent copy = (MockAgent)secondRuntime.agentInstances[agentId];
+        Assert.Single(secondRuntime._actorInstances);
+        MockAgent copy = (MockAgent)secondRuntime._actorInstances[agentId];
         Assert.Single(copy.ReceivedMessages);
         Assert.Equal(TestMessage, copy.ReceivedMessages.Single().ToString());
 
-        static async ValueTask<MockAgent> factoryFunc(AgentId id, IAgentRuntime runtime)
+        static async ValueTask<MockAgent> factoryFunc(ActorId id, IAgentRuntime runtime)
         {
             MockAgent agent = new(id, runtime, "A test agent");
             return agent;
@@ -178,14 +180,14 @@ public class InProcessRuntimeTests()
         // Arrange
         await using InProcessRuntime runtime = new();
         MockAgent? agent = null;
-        await runtime.RegisterAgentFactoryAsync("MyAgent", async (id, runtime) =>
+        await runtime.RegisterActorFactoryAsync(new("MyAgent"), async (id, runtime) =>
         {
             agent = new MockAgent(id, runtime, "A test agent");
             return agent;
         });
 
         // Act: Ensure the agent is actually created
-        AgentId agentId = await runtime.GetAgentAsync("MyAgent", lazy: false);
+        ActorId agentId = await runtime.GetActorAsync("MyAgent", lazy: false);
 
         // Assert
         Assert.NotNull(agent);
@@ -197,7 +199,7 @@ public class InProcessRuntimeTests()
         await runtime.RunUntilIdleAsync();
 
         // Assert
-        Assert.Equal(0, runtime.messageQueueCount);
+        Assert.Equal(0, runtime._messageQueueCount);
         Assert.Single(agent.ReceivedMessages);
     }
 
@@ -214,21 +216,21 @@ public class InProcessRuntimeTests()
         };
 
         MockAgent? agent = null;
-        await runtime.RegisterAgentFactoryAsync("MyAgent", async (id, runtime) =>
+        await runtime.RegisterActorFactoryAsync(new("MyAgent"), async (id, runtime) =>
         {
             agent = new MockAgent(id, runtime, "A test agent");
             return agent;
         });
 
         // Assert
-        Assert.Empty(runtime.agentInstances);
+        Assert.Empty(runtime._actorInstances);
 
         // Act: Ensure the agent is actually created
-        AgentId agentId = await runtime.GetAgentAsync("MyAgent", lazy: false);
+        ActorId agentId = await runtime.GetActorAsync("MyAgent", lazy: false);
 
         // Assert
         Assert.NotNull(agent);
-        Assert.Single(runtime.agentInstances);
+        Assert.Single(runtime._actorInstances);
 
         const string TopicType = "TestTopic";
 
@@ -250,14 +252,14 @@ public class InProcessRuntimeTests()
         // Arrange: Create a runtime and register an agent
         await using InProcessRuntime runtime = new();
         MockAgent? agent = null;
-        await runtime.RegisterAgentFactoryAsync("MyAgent", async (id, runtime) =>
+        await runtime.RegisterActorFactoryAsync(new("MyAgent"), async (id, runtime) =>
         {
             agent = new MockAgent(id, runtime, "test agent");
             return agent;
         });
 
         // Get agent ID and instantiate agent by publishing
-        AgentId agentId = await runtime.GetAgentAsync("MyAgent", lazy: false);
+        ActorId agentId = await runtime.GetActorAsync("MyAgent", lazy: false);
         const string TopicType = "TestTopic";
         await runtime.AddSubscriptionAsync(new TestSubscription(TopicType, agentId.Type));
 
@@ -286,46 +288,44 @@ public class InProcessRuntimeTests()
         agent = null;
         await using InProcessRuntime newRuntime = new();
         await newRuntime.StartAsync();
-        await newRuntime.RegisterAgentFactoryAsync("MyAgent", async (id, runtime) =>
+        await newRuntime.RegisterActorFactoryAsync(new("MyAgent"), async (id, runtime) =>
         {
             agent = new MockAgent(id, runtime, "another agent");
             return agent;
         });
 
         // Assert: Show that no agent instances exist in the new runtime
-        Assert.Empty(newRuntime.agentInstances);
+        Assert.Empty(newRuntime._actorInstances);
 
         // Act: Load the state into the new runtime and show that agent is now instantiated
         await newRuntime.LoadStateAsync(savedState);
 
         // Assert
         Assert.NotNull(agent);
-        Assert.Single(newRuntime.agentInstances);
-        Assert.True(newRuntime.agentInstances.ContainsKey(agentId));
+        Assert.Single(newRuntime._actorInstances);
+        Assert.True(newRuntime._actorInstances.ContainsKey(agentId));
         Assert.Single(agent.ReceivedMessages);
     }
 
 #pragma warning disable CA1812 // Avoid uninstantiated internal classes
-    private sealed class WrongAgent : IHostableAgent
+    private sealed class WrongAgent : IRuntimeActor
 #pragma warning restore CA1812
     {
-        public AgentId Id => throw new NotImplementedException();
+        public ActorId Id => throw new NotImplementedException();
 
-        public AgentMetadata Metadata => throw new NotImplementedException();
+        public ActorMetadata Metadata => throw new NotImplementedException();
 
-        public ValueTask CloseAsync() => default;
-
-        public ValueTask LoadStateAsync(JsonElement state)
+        public ValueTask LoadStateAsync(JsonElement state, CancellationToken cancellationToken = default)
         {
             throw new NotImplementedException();
         }
 
-        public ValueTask<object?> OnMessageAsync(object message, MessageContext messageContext)
+        public ValueTask<object?> OnMessageAsync(object message, MessageContext messageContext, CancellationToken cancellationToken = default)
         {
             throw new NotImplementedException();
         }
 
-        public ValueTask<JsonElement> SaveStateAsync()
+        public ValueTask<JsonElement> SaveStateAsync(CancellationToken cancellationToken = default)
         {
             throw new NotImplementedException();
         }

dotnet/tests/Microsoft.Extensions.AI.Agents.Runtime.InProcess.UnitTests/MessageEnvelopeTests.cs

@@ -46,7 +46,7 @@ public class MessageEnvelopeTests
     {
         // Arrange
         MessageEnvelope envelope = new("test");
-        AgentId sender = new("testtype", "testkey");
+        ActorId sender = new("testtype", "testkey");
 
         // Act
         MessageEnvelope result = envelope.WithSender(sender);
@@ -61,7 +61,7 @@ public class MessageEnvelopeTests
     {
         // Arrange
         MessageEnvelope envelope = new("test");
-        AgentId receiver = new("receivertype", "receiverkey");
+        ActorId receiver = new("receivertype", "receiverkey");
         object expectedResult = new { Response = "Success" };
 
         ValueTask<object?> servicer(MessageEnvelope env, CancellationToken ct) => new(expectedResult);
@@ -76,8 +76,8 @@ public class MessageEnvelopeTests
 
         // Invoke the servicer to verify result sink works
         await delivery.InvokeAsync(CancellationToken.None);
-        Assert.True(delivery.ResultSink.Future.IsCompleted);
-        object? result = await delivery.ResultSink.Future;
+        Assert.True(delivery.ResultTask.IsCompleted);
+        object? result = await delivery.ResultTask;
         Assert.Same(expectedResult, result);
     }
 

dotnet/tests/Microsoft.Extensions.AI.Agents.Runtime.InProcess.UnitTests/MessagingTestFixture.cs

@@ -16,114 +16,87 @@ public sealed class BasicMessage
 public sealed class TestException : Exception;
 #pragma warning restore RCS1194 // Implement exception constructors
 
-public sealed class PublisherAgent : TestAgent, IHandle<BasicMessage>
+public sealed class PublisherAgent : TestAgent
 {
-    private readonly IList<TopicId> _targetTopics;
-
-    public PublisherAgent(AgentId id, IAgentRuntime runtime, string description, IList<TopicId> targetTopics)
-        : base(id, runtime, description)
+    public PublisherAgent(ActorId id, IAgentRuntime runtime, string description, IList<TopicId> targetTopics) : base(id, runtime, description)
     {
-        this._targetTopics = targetTopics;
-    }
-
-    public async ValueTask HandleAsync(BasicMessage item, MessageContext messageContext)
-    {
-        this.ReceivedMessages.Add(item);
-        foreach (TopicId targetTopic in this._targetTopics)
+        this.RegisterMessageHandler<BasicMessage>(async (item, messageContext, cancellationToken) =>
         {
-            await this.PublishMessageAsync(
-                new BasicMessage { Content = $"@{targetTopic}: {item.Content}" },
-                targetTopic);
-        }
+            this.ReceivedMessages.Add(item);
+            foreach (TopicId targetTopic in targetTopics)
+            {
+                await this.PublishMessageAsync(
+                    new BasicMessage { Content = $"@{targetTopic}: {item.Content}" },
+                    targetTopic,
+                    cancellationToken: cancellationToken);
+            }
+        });
     }
 }
 
-public sealed class SendOnAgent : TestAgent, IHandle<BasicMessage>
+public sealed class SendOnAgent : TestAgent
 {
-    private readonly IList<Guid> _targetKeys;
-
-    public SendOnAgent(AgentId id, IAgentRuntime runtime, string description, IList<Guid> targetKeys)
-        : base(id, runtime, description)
+    public SendOnAgent(ActorId id, IAgentRuntime runtime, string description, IList<Guid> targetKeys) : base(id, runtime, description)
     {
-        this._targetKeys = targetKeys;
-    }
-
-    public async ValueTask HandleAsync(BasicMessage item, MessageContext messageContext)
-    {
-        foreach (Guid targetKey in this._targetKeys)
+        this.RegisterMessageHandler<BasicMessage>(async (item, messageContext, cancellationToken) =>
         {
-            AgentId targetId = new(nameof(ReceiverAgent), targetKey.ToString());
-            BasicMessage response = new() { Content = $"@{targetKey}: {item.Content}" };
-            await this.SendMessageAsync(response, targetId);
-        }
+            foreach (Guid targetKey in targetKeys)
+            {
+                ActorId targetId = new(nameof(ReceiverAgent), targetKey.ToString());
+                BasicMessage response = new() { Content = $"@{targetKey}: {item.Content}" };
+                await this.SendMessageAsync(response, targetId, cancellationToken: cancellationToken);
+            }
+        });
     }
 }
 
-public sealed class ReceiverAgent : TestAgent, IHandle<BasicMessage>
+public sealed class ReceiverAgent : TestAgent
 {
     public List<BasicMessage> Messages { get; } = [];
 
-    public ReceiverAgent(AgentId id, IAgentRuntime runtime, string description)
-        : base(id, runtime, description)
+    public ReceiverAgent(ActorId id, IAgentRuntime runtime, string description) : base(id, runtime, description)
     {
-    }
-
-    public ValueTask HandleAsync(BasicMessage item, MessageContext messageContext)
-    {
-        this.Messages.Add(item);
-        return default;
+        this.RegisterMessageHandler<BasicMessage>(async (item, messageContext, cancellationToken) =>
+        {
+            this.Messages.Add(item);
+        });
     }
 }
 
-public sealed class ProcessorAgent : TestAgent, IHandle<BasicMessage, BasicMessage>
+public sealed class ProcessorAgent : TestAgent
 {
-    private Func<string, string> ProcessFunc { get; }
-
-    public ProcessorAgent(AgentId id, IAgentRuntime runtime, Func<string, string> processFunc, string description)
-        : base(id, runtime, description)
+    public ProcessorAgent(ActorId id, IAgentRuntime runtime, Func<string, string> processFunc, string description) : base(id, runtime, description)
     {
-        this.ProcessFunc = processFunc;
-    }
-
-    public ValueTask<BasicMessage> HandleAsync(BasicMessage item, MessageContext messageContext)
-    {
-        BasicMessage result = new() { Content = this.ProcessFunc.Invoke(((BasicMessage)item).Content) };
-
-        return new(result);
+        this.RegisterMessageHandler<BasicMessage, BasicMessage>(async (item, messageContext, cancellationtoken) =>
+        {
+            return new BasicMessage() { Content = processFunc.Invoke(((BasicMessage)item).Content) };
+        });
     }
 }
 
-public sealed class CancelAgent : TestAgent, IHandle<BasicMessage>
+public sealed class CancelAgent : TestAgent
 {
-    public CancelAgent(AgentId id, IAgentRuntime runtime, string description)
-        : base(id, runtime, description)
+    public CancelAgent(ActorId id, IAgentRuntime runtime, string description) : base(id, runtime, description)
     {
-    }
-
-    public ValueTask HandleAsync(BasicMessage item, MessageContext messageContext)
-    {
-        CancellationToken cancelledToken = new(canceled: true);
-        cancelledToken.ThrowIfCancellationRequested();
-
-        return default;
+        this.RegisterMessageHandler<BasicMessage>(async (item, messageContext, cancellationToken) =>
+        {
+            CancellationToken cancelledToken = new(canceled: true);
+            cancelledToken.ThrowIfCancellationRequested();
+        });
     }
 }
 
-public sealed class ErrorAgent : TestAgent, IHandle<BasicMessage>
+public sealed class ErrorAgent : TestAgent
 {
-    public ErrorAgent(AgentId id, IAgentRuntime runtime, string description)
-        : base(id, runtime, description)
+    public ErrorAgent(ActorId id, IAgentRuntime runtime, string description) : base(id, runtime, description)
     {
+        this.RegisterMessageHandler<BasicMessage>(async (item, messageContext, cancellationToken) =>
+        {
+            this.DidThrow = true;
+            throw new TestException();
+        });
     }
-
     public bool DidThrow { get; private set; }
-
-    public ValueTask HandleAsync(BasicMessage item, MessageContext messageContext)
-    {
-        this.DidThrow = true;
-
-        throw new TestException();
-    }
 }
 
 public sealed class MessagingTestFixture
@@ -131,23 +104,23 @@ public sealed class MessagingTestFixture
     private Dictionary<Type, object> AgentsTypeMap { get; } = [];
     public InProcessRuntime Runtime { get; } = new();
 
-    public ValueTask<AgentType> RegisterFactoryMapInstances<TAgent>(AgentType type, Func<AgentId, IAgentRuntime, ValueTask<TAgent>> factory)
-        where TAgent : IHostableAgent
+    public ValueTask<ActorType> RegisterFactoryMapInstances<TAgent>(ActorType type, Func<ActorId, IAgentRuntime, ValueTask<TAgent>> factory)
+        where TAgent : IRuntimeActor
     {
-        async ValueTask<TAgent> WrappedFactory(AgentId id, IAgentRuntime runtime)
+        async ValueTask<TAgent> WrappedFactory(ActorId id, IAgentRuntime runtime)
         {
             TAgent agent = await factory(id, runtime);
             this.GetAgentInstances<TAgent>()[id] = agent;
             return agent;
         }
 
-        return this.Runtime.RegisterAgentFactoryAsync(type, WrappedFactory);
+        return this.Runtime.RegisterActorFactoryAsync(type, WrappedFactory);
     }
 
-    public Dictionary<AgentId, TAgent> GetAgentInstances<TAgent>() where TAgent : IHostableAgent
+    public Dictionary<ActorId, TAgent> GetAgentInstances<TAgent>() where TAgent : IRuntimeActor
     {
         if (!this.AgentsTypeMap.TryGetValue(typeof(TAgent), out object? maybeAgentMap) ||
-            maybeAgentMap is not Dictionary<AgentId, TAgent> result)
+            maybeAgentMap is not Dictionary<ActorId, TAgent> result)
         {
             this.AgentsTypeMap[typeof(TAgent)] = result = [];
         }
@@ -157,24 +130,24 @@ public sealed class MessagingTestFixture
     public async ValueTask RegisterReceiverAgentAsync(string? agentNameSuffix = null, params string[] topicTypes)
     {
         await this.RegisterFactoryMapInstances(
-            $"{nameof(ReceiverAgent)}{agentNameSuffix ?? string.Empty}",
+            new($"{nameof(ReceiverAgent)}{agentNameSuffix ?? string.Empty}"),
             (id, runtime) => new ValueTask<ReceiverAgent>(new ReceiverAgent(id, runtime, string.Empty)));
 
         foreach (string topicType in topicTypes)
         {
-            await this.Runtime.AddSubscriptionAsync(new TestSubscription(topicType, $"{nameof(ReceiverAgent)}{agentNameSuffix ?? string.Empty}"));
+            await this.Runtime.AddSubscriptionAsync(new TestSubscription(topicType, new($"{nameof(ReceiverAgent)}{agentNameSuffix ?? string.Empty}")));
         }
     }
 
     public async ValueTask RegisterErrorAgentAsync(string? agentNameSuffix = null, params string[] topicTypes)
     {
         await this.RegisterFactoryMapInstances(
-            $"{nameof(ErrorAgent)}{agentNameSuffix ?? string.Empty}",
+            new($"{nameof(ErrorAgent)}{agentNameSuffix ?? string.Empty}"),
             (id, runtime) => new ValueTask<ErrorAgent>(new ErrorAgent(id, runtime, string.Empty)));
 
         foreach (string topicType in topicTypes)
         {
-            await this.Runtime.AddSubscriptionAsync(new TestSubscription(topicType, $"{nameof(ErrorAgent)}{agentNameSuffix ?? string.Empty}"));
+            await this.Runtime.AddSubscriptionAsync(new TestSubscription(topicType, new($"{nameof(ErrorAgent)}{agentNameSuffix ?? string.Empty}")));
         }
     }
 
@@ -187,7 +160,7 @@ public sealed class MessagingTestFixture
         await this.Runtime.RunUntilIdleAsync();
     }
 
-    public async ValueTask<object?> RunSendTestAsync(AgentId sendTarget, object message, string? messageId = null)
+    public async ValueTask<object?> RunSendTestAsync(ActorId sendTarget, object message, string? messageId = null)
     {
         messageId ??= Guid.NewGuid().ToString();
 

dotnet/tests/Microsoft.Extensions.AI.Agents.Runtime.InProcess.UnitTests/PublishMessageTests.cs

@@ -93,10 +93,10 @@ public class PublishMessageTests
         MessagingTestFixture fixture = new();
 
         await fixture.RegisterFactoryMapInstances(
-            nameof(PublisherAgent),
+            new(nameof(PublisherAgent)),
             (id, runtime) => new ValueTask<PublisherAgent>(new PublisherAgent(id, runtime, string.Empty, [new TopicId("TestTopic")])));
 
-        await fixture.Runtime.AddSubscriptionAsync(new TestSubscription("RunTest", nameof(PublisherAgent)));
+        await fixture.Runtime.AddSubscriptionAsync(new TestSubscription("RunTest", new(nameof(PublisherAgent))));
 
         await fixture.RegisterReceiverAgentAsync(topicTypes: "TestTopic");
         await fixture.RegisterReceiverAgentAsync("2", topicTypes: "TestTopic");

dotnet/tests/Microsoft.Extensions.AI.Agents.Runtime.InProcess.UnitTests/ResultSinkTests.cs

@@ -1,103 +0,0 @@
-// Copyright (c) Microsoft. All rights reserved.
-
-using System;
-using System.Threading.Tasks;
-using System.Threading.Tasks.Sources;
-
-namespace Microsoft.Extensions.AI.Agents.Runtime.InProcess.Tests;
-
-public class ResultSinkTests
-{
-    [Fact]
-    public void GetResultTest()
-    {
-        // Arrange
-        ResultSink<int> sink = new();
-        const int ExpectedResult = 42;
-
-        // Act
-        sink.SetResult(ExpectedResult);
-        int result = sink.GetResult(0);
-
-        // Assert
-        Assert.Equal(ExpectedResult, result);
-        Assert.Equal(ValueTaskSourceStatus.Succeeded, sink.GetStatus(0));
-    }
-
-    [Fact]
-    public async Task FutureResultTestAsync()
-    {
-        // Arrange
-        ResultSink<string> sink = new();
-        const string ExpectedResult = "test";
-
-        // Act
-        sink.SetResult(ExpectedResult);
-        string result = await sink.Future;
-
-        // Assert
-        Assert.Equal(ExpectedResult, result);
-        Assert.Equal(ValueTaskSourceStatus.Succeeded, sink.GetStatus(0));
-    }
-
-    [Fact]
-    public async Task SetExceptionTestAsync()
-    {
-        // Arrange
-        ResultSink<int> sink = new();
-        InvalidOperationException expectedException = new("Test exception");
-
-        // Act
-        sink.SetException(expectedException);
-
-        // Assert
-        Exception exception = await Assert.ThrowsAsync<InvalidOperationException>(async () => await sink.Future);
-        Assert.Equal(expectedException.Message, exception.Message);
-        exception = Assert.Throws<InvalidOperationException>(() => sink.GetResult(0));
-        Assert.Equal(expectedException.Message, exception.Message);
-        Assert.Equal(ValueTaskSourceStatus.Faulted, sink.GetStatus(0));
-    }
-
-    [Fact]
-    public async Task SetCancelledTestAsync()
-    {
-        // Arrange
-        ResultSink<int> sink = new();
-
-        // Act
-        sink.SetCancelled();
-
-        // Assert
-        Assert.True(sink.IsCancelled);
-        Assert.Throws<OperationCanceledException>(() => sink.GetResult(0));
-        await Assert.ThrowsAsync<OperationCanceledException>(async () => await sink.Future);
-        Assert.Equal(ValueTaskSourceStatus.Canceled, sink.GetStatus(0));
-    }
-
-    [Fact]
-    public void OnCompletedTest()
-    {
-        // Arrange
-        ResultSink<int> sink = new();
-        bool continuationCalled = false;
-        const int ExpectedResult = 42;
-
-        // Register the continuation
-        sink.OnCompleted(
-            state => continuationCalled = true,
-            state: null,
-            token: 0,
-            ValueTaskSourceOnCompletedFlags.None);
-
-        // Assert
-        Assert.False(continuationCalled, "Continuation should have been called");
-
-        // Act
-        sink.SetResult(ExpectedResult);
-
-        // Assert
-        Assert.Equal(ExpectedResult, sink.GetResult(0));
-        Assert.Equal(ValueTaskSourceStatus.Succeeded, sink.GetStatus(0));
-        Assert.True(continuationCalled, "Continuation should have been called");
-    }
-}

dotnet/tests/Microsoft.Extensions.AI.Agents.Runtime.InProcess.UnitTests/SendMessageTests.cs

@@ -16,10 +16,10 @@ public class SendMessageTests
 
         MessagingTestFixture fixture = new();
 
-        await fixture.RegisterFactoryMapInstances(nameof(ProcessorAgent),
+        await fixture.RegisterFactoryMapInstances(new(nameof(ProcessorAgent)),
             (id, runtime) => new ValueTask<ProcessorAgent>(new ProcessorAgent(id, runtime, ProcessFunc, string.Empty)));
 
-        AgentId targetAgent = new(nameof(ProcessorAgent), Guid.NewGuid().ToString());
+        ActorId targetAgent = new(nameof(ProcessorAgent), Guid.NewGuid().ToString());
         object? maybeResult = await fixture.RunSendTestAsync(targetAgent, new BasicMessage { Content = "1" });
 
         Assert.Equal("Processed(1)", Assert.IsType<BasicMessage>(maybeResult).Content);
@@ -30,12 +30,12 @@ public class SendMessageTests
     {
         MessagingTestFixture fixture = new();
 
-        await fixture.RegisterFactoryMapInstances(nameof(CancelAgent),
+        await fixture.RegisterFactoryMapInstances(new(nameof(CancelAgent)),
             (id, runtime) => new ValueTask<CancelAgent>(new CancelAgent(id, runtime, string.Empty)));
 
-        AgentId targetAgent = new(nameof(CancelAgent), Guid.NewGuid().ToString());
+        ActorId targetAgent = new(nameof(CancelAgent), Guid.NewGuid().ToString());
 
-        await Assert.ThrowsAsync<OperationCanceledException>(() => fixture.RunSendTestAsync(targetAgent, new BasicMessage { Content = "1" }).AsTask());
+        await Assert.ThrowsAnyAsync<OperationCanceledException>(() => fixture.RunSendTestAsync(targetAgent, new BasicMessage { Content = "1" }).AsTask());
     }
 
     [Fact]
@@ -43,10 +43,10 @@ public class SendMessageTests
     {
         MessagingTestFixture fixture = new();
 
-        await fixture.RegisterFactoryMapInstances(nameof(ErrorAgent),
+        await fixture.RegisterFactoryMapInstances(new(nameof(ErrorAgent)),
             (id, runtime) => new ValueTask<ErrorAgent>(new ErrorAgent(id, runtime, string.Empty)));
 
-        AgentId targetAgent = new(nameof(ErrorAgent), Guid.NewGuid().ToString());
+        ActorId targetAgent = new(nameof(ErrorAgent), Guid.NewGuid().ToString());
 
         await Assert.ThrowsAsync<TestException>(() => fixture.RunSendTestAsync(targetAgent, new BasicMessage { Content = "1" }).AsTask());
     }
@@ -58,16 +58,16 @@ public class SendMessageTests
 
         MessagingTestFixture fixture = new();
 
-        Dictionary<AgentId, SendOnAgent> sendAgents = fixture.GetAgentInstances<SendOnAgent>();
-        Dictionary<AgentId, ReceiverAgent> receiverAgents = fixture.GetAgentInstances<ReceiverAgent>();
+        Dictionary<ActorId, SendOnAgent> sendAgents = fixture.GetAgentInstances<SendOnAgent>();
+        Dictionary<ActorId, ReceiverAgent> receiverAgents = fixture.GetAgentInstances<ReceiverAgent>();
 
-        await fixture.RegisterFactoryMapInstances(nameof(SendOnAgent),
+        await fixture.RegisterFactoryMapInstances(new(nameof(SendOnAgent)),
             (id, runtime) => new ValueTask<SendOnAgent>(new SendOnAgent(id, runtime, string.Empty, targetGuids)));
 
-        await fixture.RegisterFactoryMapInstances(nameof(ReceiverAgent),
+        await fixture.RegisterFactoryMapInstances(new(nameof(ReceiverAgent)),
             (id, runtime) => new ValueTask<ReceiverAgent>(new ReceiverAgent(id, runtime, string.Empty)));
 
-        AgentId targetAgent = new(nameof(SendOnAgent), Guid.NewGuid().ToString());
+        ActorId targetAgent = new(nameof(SendOnAgent), Guid.NewGuid().ToString());
         BasicMessage input = new() { Content = "Hello" };
         Task testTask = fixture.RunSendTestAsync(targetAgent, input).AsTask();
 
@@ -82,7 +82,7 @@ public class SendMessageTests
         // Check that each of the target agents received the message
         foreach (Guid targetKey in targetGuids)
         {
-            AgentId targetId = new(nameof(ReceiverAgent), targetKey.ToString());
+            ActorId targetId = new(nameof(ReceiverAgent), targetKey.ToString());
             Assert.Single(receiverAgents[targetId].Messages);
             Assert.Contains(receiverAgents[targetId].Messages, m => m.Content == $"@{targetKey}: {input.Content}");
         }

dotnet/tests/Microsoft.Extensions.AI.Agents.Runtime.InProcess.UnitTests/TestAgents.cs

@@ -3,15 +3,16 @@
 using System;
 using System.Collections.Generic;
 using System.Text.Json;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace Microsoft.Extensions.AI.Agents.Runtime.InProcess.Tests;
 
-public abstract class TestAgent : BaseAgent
+public abstract class TestAgent : RuntimeActor
 {
     internal List<object> ReceivedMessages = [];
 
-    protected TestAgent(AgentId id, IAgentRuntime runtime, string description)
+    protected TestAgent(ActorId id, IAgentRuntime runtime, string description)
         : base(id, runtime, description)
     {
     }
@@ -21,24 +22,26 @@ public abstract class TestAgent : BaseAgent
 /// A test agent that captures the messages it receives and
 /// is able to save and load its state.
 /// </summary>
-public sealed class MockAgent : TestAgent, IHandle<string>
+public sealed class MockAgent : TestAgent
 {
-    public MockAgent(AgentId id, IAgentRuntime runtime, string description)
-        : base(id, runtime, description) { }
+    public MockAgent(ActorId id, IAgentRuntime runtime, string description) : base(id, runtime, description)
+    {
+        this.RegisterMessageHandler<string>(this.HandleAsync);
+    }
 
-    public ValueTask HandleAsync(string item, MessageContext messageContext)
+    public ValueTask HandleAsync(string item, MessageContext messageContext, CancellationToken cancellationToken)
     {
         this.ReceivedMessages.Add(item);
         return default;
     }
 
-    public override async ValueTask<JsonElement> SaveStateAsync()
+    public override async ValueTask<JsonElement> SaveStateAsync(CancellationToken cancellationToken = default)
     {
         JsonElement json = JsonSerializer.SerializeToElement(this.ReceivedMessages);
         return json;
     }
 
-    public override ValueTask LoadStateAsync(JsonElement state)
+    public override ValueTask LoadStateAsync(JsonElement state, CancellationToken cancellationToken = default)
     {
         this.ReceivedMessages = JsonSerializer.Deserialize<List<object>>(state) ?? throw new InvalidOperationException("Failed to deserialize state");
         return default;

dotnet/tests/Microsoft.Extensions.AI.Agents.Runtime.InProcess.UnitTests/TestSubscription.cs

@@ -5,20 +5,20 @@ using System.Diagnostics.CodeAnalysis;
 
 namespace Microsoft.Extensions.AI.Agents.Runtime.InProcess.Tests;
 
-public class TestSubscription(string topicType, string agentType, string? id = null) : ISubscriptionDefinition
+public class TestSubscription(string topicType, ActorType agentType, string? id = null) : ISubscriptionDefinition
 {
     public string Id { get; } = id ?? Guid.NewGuid().ToString();
 
     public string TopicType { get; } = topicType;
 
-    public AgentId MapToAgent(TopicId topic)
+    public ActorId MapToActor(TopicId topic)
     {
         if (!this.Matches(topic))
         {
             throw new InvalidOperationException("TopicId does not match the subscription.");
         }
 
-        return new AgentId(agentType, topic.Source);
+        return new ActorId(agentType, topic.Source);
     }
 
     public bool Equals(ISubscriptionDefinition? other) => this.Id == other?.Id;