From c8e4ad82423ff8f5d8560e16b495c79b0989b94d Mon Sep 17 00:00:00 2001 From: Griffin Bassman Date: Thu, 6 Feb 2025 17:09:26 -0500 Subject: [PATCH] feat: save/load test for dotnet agents (#5284) --- .../Microsoft.AutoGen/Contracts/AgentProxy.cs | 6 +- .../Contracts/IAgentRuntime.cs | 2 +- .../Microsoft.AutoGen/Contracts/ISaveState.cs | 2 +- .../Core.Grpc/GrpcAgentRuntime.cs | 39 ++--- .../src/Microsoft.AutoGen/Core/BaseAgent.cs | 7 +- .../Core/InProcessRuntime.cs | 26 ++-- .../Core/Properties/AssemblyInfo.cs | 6 + .../AgentRuntimeTests.cs | 83 ----------- .../AgentTests.cs | 23 +-- .../InProcessRuntimeTests.cs | 141 ++++++++++++++++++ .../Microsoft.AutoGen.Core.Tests/TestAgent.cs | 35 ++++- 11 files changed, 230 insertions(+), 140 deletions(-) create mode 100644 dotnet/src/Microsoft.AutoGen/Core/Properties/AssemblyInfo.cs delete mode 100644 dotnet/test/Microsoft.AutoGen.Core.Tests/AgentRuntimeTests.cs create mode 100644 dotnet/test/Microsoft.AutoGen.Core.Tests/InProcessRuntimeTests.cs diff --git a/dotnet/src/Microsoft.AutoGen/Contracts/AgentProxy.cs b/dotnet/src/Microsoft.AutoGen/Contracts/AgentProxy.cs index 44ad9b0e1..d37d6284b 100644 --- a/dotnet/src/Microsoft.AutoGen/Contracts/AgentProxy.cs +++ b/dotnet/src/Microsoft.AutoGen/Contracts/AgentProxy.cs @@ -1,6 +1,8 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // AgentProxy.cs +using System.Text.Json; + namespace Microsoft.AutoGen.Contracts; /// @@ -55,7 +57,7 @@ public class AgentProxy(AgentId agentId, IAgentRuntime runtime) /// /// A dictionary representing the state of the agent. Must be JSON serializable. /// A task representing the asynchronous operation. - public ValueTask LoadStateAsync(IDictionary state) + public ValueTask LoadStateAsync(IDictionary state) { return this.runtime.LoadAgentStateAsync(this.Id, state); } @@ -64,7 +66,7 @@ public class AgentProxy(AgentId agentId, IAgentRuntime runtime) /// Saves the state of the agent. The result must be JSON serializable. /// /// A task representing the asynchronous operation, returning a dictionary containing the saved state. - public ValueTask> SaveStateAsync() + public ValueTask> SaveStateAsync() { return this.runtime.SaveAgentStateAsync(this.Id); } diff --git a/dotnet/src/Microsoft.AutoGen/Contracts/IAgentRuntime.cs b/dotnet/src/Microsoft.AutoGen/Contracts/IAgentRuntime.cs index 0d84fbe72..c4b2e998f 100644 --- a/dotnet/src/Microsoft.AutoGen/Contracts/IAgentRuntime.cs +++ b/dotnet/src/Microsoft.AutoGen/Contracts/IAgentRuntime.cs @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // IAgentRuntime.cs -using StateDict = System.Collections.Generic.IDictionary; +using StateDict = System.Collections.Generic.IDictionary; namespace Microsoft.AutoGen.Contracts; diff --git a/dotnet/src/Microsoft.AutoGen/Contracts/ISaveState.cs b/dotnet/src/Microsoft.AutoGen/Contracts/ISaveState.cs index ed6d15d1d..4f98f1fc4 100644 --- a/dotnet/src/Microsoft.AutoGen/Contracts/ISaveState.cs +++ b/dotnet/src/Microsoft.AutoGen/Contracts/ISaveState.cs @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ISaveState.cs -using StateDict = System.Collections.Generic.IDictionary; +using StateDict = System.Collections.Generic.IDictionary; namespace Microsoft.AutoGen.Contracts; diff --git a/dotnet/src/Microsoft.AutoGen/Core.Grpc/GrpcAgentRuntime.cs b/dotnet/src/Microsoft.AutoGen/Core.Grpc/GrpcAgentRuntime.cs index 1ff103601..461148843 100644 --- a/dotnet/src/Microsoft.AutoGen/Core.Grpc/GrpcAgentRuntime.cs +++ b/dotnet/src/Microsoft.AutoGen/Core.Grpc/GrpcAgentRuntime.cs @@ -2,6 +2,7 @@ // GrpcAgentRuntime.cs using System.Collections.Concurrent; +using System.Text.Json; using Grpc.Core; using Microsoft.AutoGen.Contracts; using Microsoft.AutoGen.Protobuf; @@ -319,13 +320,13 @@ public sealed class GrpcAgentRuntime : IHostedService, IAgentRuntime, IMessageSi public ValueTask GetAgentAsync(string agent, string key = "default", bool lazy = true) => this.GetAgentAsync(new Contracts.AgentId(agent, key), lazy); - public async ValueTask> SaveAgentStateAsync(Contracts.AgentId agentId) + public async ValueTask> SaveAgentStateAsync(Contracts.AgentId agentId) { IHostableAgent agent = await this._agentsContainer.EnsureAgentAsync(agentId); return await agent.SaveStateAsync(); } - public async ValueTask LoadAgentStateAsync(Contracts.AgentId agentId, IDictionary state) + public async ValueTask LoadAgentStateAsync(Contracts.AgentId agentId, IDictionary state) { IHostableAgent agent = await this._agentsContainer.EnsureAgentAsync(agentId); await agent.LoadStateAsync(state); @@ -375,37 +376,41 @@ public sealed class GrpcAgentRuntime : IHostedService, IAgentRuntime, IMessageSi return ValueTask.FromResult(new AgentProxy(agentId, this)); } - public async ValueTask> SaveStateAsync() - { - Dictionary state = new(); - foreach (var agent in this._agentsContainer.LiveAgents) - { - state[agent.Id.ToString()] = await agent.SaveStateAsync(); - } - - return state; - } - - public async ValueTask LoadStateAsync(IDictionary state) + public async ValueTask LoadStateAsync(IDictionary state) { HashSet registeredTypes = this._agentsContainer.RegisteredAgentTypes; foreach (var agentIdStr in state.Keys) { Contracts.AgentId agentId = Contracts.AgentId.FromStr(agentIdStr); - if (state[agentIdStr] is not IDictionary agentStateDict) + + if (state[agentIdStr].ValueKind != JsonValueKind.Object) { - throw new Exception($"Agent state for {agentId} is not a {typeof(IDictionary)}: {state[agentIdStr].GetType()}"); + throw new Exception($"Agent state for {agentId} is not a valid JSON object."); } + var agentState = JsonSerializer.Deserialize>(state[agentIdStr].GetRawText()) + ?? throw new Exception($"Failed to deserialize state for {agentId}."); + if (registeredTypes.Contains(agentId.Type)) { IHostableAgent agent = await this._agentsContainer.EnsureAgentAsync(agentId); - await agent.LoadStateAsync(agentStateDict); + await agent.LoadStateAsync(agentState); } } } + public async ValueTask> SaveStateAsync() + { + Dictionary state = new(); + foreach (var agent in this._agentsContainer.LiveAgents) + { + var agentState = await agent.SaveStateAsync(); + state[agent.Id.ToString()] = JsonSerializer.SerializeToElement(agentState); + } + return state; + } + public async ValueTask OnMessageAsync(Message message, CancellationToken cancellation = default) { switch (message.MessageCase) diff --git a/dotnet/src/Microsoft.AutoGen/Core/BaseAgent.cs b/dotnet/src/Microsoft.AutoGen/Core/BaseAgent.cs index 99ff001ba..a3899280f 100644 --- a/dotnet/src/Microsoft.AutoGen/Core/BaseAgent.cs +++ b/dotnet/src/Microsoft.AutoGen/Core/BaseAgent.cs @@ -4,6 +4,7 @@ using System.Diagnostics; using System.Reflection; +using System.Text.Json; using Microsoft.AutoGen.Contracts; using Microsoft.Extensions.Logging; @@ -92,11 +93,11 @@ public abstract class BaseAgent : IAgent, IHostableAgent return null; } - public virtual ValueTask> SaveStateAsync() + public virtual ValueTask> SaveStateAsync() { - return ValueTask.FromResult>(new Dictionary()); + return ValueTask.FromResult>(new Dictionary()); } - public virtual ValueTask LoadStateAsync(IDictionary state) + public virtual ValueTask LoadStateAsync(IDictionary state) { return ValueTask.CompletedTask; } diff --git a/dotnet/src/Microsoft.AutoGen/Core/InProcessRuntime.cs b/dotnet/src/Microsoft.AutoGen/Core/InProcessRuntime.cs index 69b2d314e..9acf96e64 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.Text.Json; using Microsoft.AutoGen.Contracts; using Microsoft.Extensions.Hosting; @@ -12,7 +13,7 @@ public sealed class InProcessRuntime : IAgentRuntime, IHostedService { public bool DeliverToSelf { get; set; } //= false; - Dictionary agentInstances = new(); + internal Dictionary agentInstances = new(); Dictionary subscriptions = new(); Dictionary>> agentFactories = new(); @@ -152,13 +153,13 @@ public sealed class InProcessRuntime : IAgentRuntime, IHostedService return agent.Metadata; } - public async ValueTask LoadAgentStateAsync(AgentId agentId, IDictionary state) + public async ValueTask LoadAgentStateAsync(AgentId agentId, IDictionary state) { IHostableAgent agent = await this.EnsureAgentAsync(agentId); await agent.LoadStateAsync(state); } - public async ValueTask> SaveAgentStateAsync(AgentId agentId) + public async ValueTask> SaveAgentStateAsync(AgentId agentId) { IHostableAgent agent = await this.EnsureAgentAsync(agentId); return await agent.SaveStateAsync(); @@ -187,16 +188,21 @@ public sealed class InProcessRuntime : IAgentRuntime, IHostedService return ValueTask.CompletedTask; } - public async ValueTask LoadStateAsync(IDictionary state) + public async ValueTask LoadStateAsync(IDictionary state) { foreach (var agentIdStr in state.Keys) { AgentId agentId = AgentId.FromStr(agentIdStr); - if (state[agentIdStr] is not IDictionary agentState) + + if (state[agentIdStr].ValueKind != JsonValueKind.Object) { - throw new Exception($"Agent state for {agentId} is not a {typeof(IDictionary)}: {state[agentIdStr].GetType()}"); + throw new Exception($"Agent state for {agentId} is not a valid JSON object."); } + // Deserialize before using + var agentState = JsonSerializer.Deserialize>(state[agentIdStr].GetRawText()) + ?? throw new Exception($"Failed to deserialize state for {agentId}."); + if (this.agentFactories.ContainsKey(agentId.Type)) { IHostableAgent agent = await this.EnsureAgentAsync(agentId); @@ -205,14 +211,14 @@ public sealed class InProcessRuntime : IAgentRuntime, IHostedService } } - public async ValueTask> SaveStateAsync() + public async ValueTask> SaveStateAsync() { - Dictionary state = new(); + Dictionary state = new(); foreach (var agentId in this.agentInstances.Keys) { - state[agentId.ToString()] = await this.agentInstances[agentId].SaveStateAsync(); + var agentState = await this.agentInstances[agentId].SaveStateAsync(); + state[agentId.ToString()] = JsonSerializer.SerializeToElement(agentState); } - return state; } diff --git a/dotnet/src/Microsoft.AutoGen/Core/Properties/AssemblyInfo.cs b/dotnet/src/Microsoft.AutoGen/Core/Properties/AssemblyInfo.cs new file mode 100644 index 000000000..8ff444817 --- /dev/null +++ b/dotnet/src/Microsoft.AutoGen/Core/Properties/AssemblyInfo.cs @@ -0,0 +1,6 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// AssemblyInfo.cs + +using System.Runtime.CompilerServices; + +[assembly: InternalsVisibleTo("Microsoft.AutoGen.Core.Tests, PublicKey=0024000004800000940000000602000000240000525341310004000001000100f1d038d0b85ae392ad72011df91e9343b0b5df1bb8080aa21b9424362d696919e0e9ac3a8bca24e283e10f7a569c6f443e1d4e3ebc84377c87ca5caa562e80f9932bf5ea91b7862b538e13b8ba91c7565cf0e8dfeccfea9c805ae3bda044170ecc7fc6f147aeeac422dd96aeb9eb1f5a5882aa650efe2958f2f8107d2038f2ab")] diff --git a/dotnet/test/Microsoft.AutoGen.Core.Tests/AgentRuntimeTests.cs b/dotnet/test/Microsoft.AutoGen.Core.Tests/AgentRuntimeTests.cs deleted file mode 100644 index 812d47c2d..000000000 --- a/dotnet/test/Microsoft.AutoGen.Core.Tests/AgentRuntimeTests.cs +++ /dev/null @@ -1,83 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// AgentRuntimeTests.cs -using FluentAssertions; -using Microsoft.AutoGen.Contracts; -using Microsoft.Extensions.Logging; -using Xunit; - -namespace Microsoft.AutoGen.Core.Tests; - -[Trait("Category", "UnitV2")] -public class AgentRuntimeTests() -{ - // Agent will not deliver to self will success when runtime.DeliverToSelf is false (default) - [Fact] - public async Task RuntimeAgentPublishToSelfDefaultNoSendTest() - { - var runtime = new InProcessRuntime(); - await runtime.StartAsync(); - - Logger logger = new(new LoggerFactory()); - SubscribedSelfPublishAgent agent = null!; - - await runtime.RegisterAgentFactoryAsync("MyAgent", (id, runtime) => - { - agent = new SubscribedSelfPublishAgent(id, runtime, logger); - return ValueTask.FromResult(agent); - }); - - // Ensure the agent is actually created - AgentId agentId = await runtime.GetAgentAsync("MyAgent", lazy: false); - - // Validate agent ID - agentId.Should().Be(agent.Id, "Agent ID should match the registered agent"); - - await runtime.RegisterImplicitAgentSubscriptionsAsync("MyAgent"); - - var topicType = "TestTopic"; - - await runtime.PublishMessageAsync("SelfMessage", new TopicId(topicType)).ConfigureAwait(true); - - await runtime.RunUntilIdleAsync(); - - // Agent has default messages and could not publish to self - agent.Text.Source.Should().Be("DefaultTopic"); - agent.Text.Content.Should().Be("DefaultContent"); - } - - // Agent delivery to self will success when runtime.DeliverToSelf is true - [Fact] - public async Task RuntimeAgentPublishToSelfDeliverToSelfTrueTest() - { - var runtime = new InProcessRuntime(); - runtime.DeliverToSelf = true; - await runtime.StartAsync(); - - Logger logger = new(new LoggerFactory()); - SubscribedSelfPublishAgent agent = null!; - - await runtime.RegisterAgentFactoryAsync("MyAgent", (id, runtime) => - { - agent = new SubscribedSelfPublishAgent(id, runtime, logger); - return ValueTask.FromResult(agent); - }); - - // Ensure the agent is actually created - AgentId agentId = await runtime.GetAgentAsync("MyAgent", lazy: false); - - // Validate agent ID - agentId.Should().Be(agent.Id, "Agent ID should match the registered agent"); - - await runtime.RegisterImplicitAgentSubscriptionsAsync("MyAgent"); - - var topicType = "TestTopic"; - - await runtime.PublishMessageAsync("SelfMessage", new TopicId(topicType)).ConfigureAwait(true); - - await runtime.RunUntilIdleAsync(); - - // Agent sucessfully published to self - agent.Text.Source.Should().Be("TestTopic"); - agent.Text.Content.Should().Be("SelfMessage"); - } -} diff --git a/dotnet/test/Microsoft.AutoGen.Core.Tests/AgentTests.cs b/dotnet/test/Microsoft.AutoGen.Core.Tests/AgentTests.cs index c091f9eb7..805fbc871 100644 --- a/dotnet/test/Microsoft.AutoGen.Core.Tests/AgentTests.cs +++ b/dotnet/test/Microsoft.AutoGen.Core.Tests/AgentTests.cs @@ -54,7 +54,7 @@ public class AgentTests() return ValueTask.FromResult(agent); }); - // Ensure the agent is actually created + // Ensure the agent id is registered AgentId agentId = await runtime.GetAgentAsync("MyAgent", lazy: false); // Validate agent ID @@ -146,25 +146,4 @@ public class AgentTests() Assert.True(agent.ReceivedItems.Count == 1); } - - [Fact] - public async Task AgentShouldSaveStateCorrectlyTest() - { - var runtime = new InProcessRuntime(); - await runtime.StartAsync(); - - Logger logger = new(new LoggerFactory()); - TestAgent agent = new TestAgent(new AgentId("TestType", "TestKey"), runtime, logger); - - var state = await agent.SaveStateAsync(); - - // Ensure state is a dictionary - state.Should().NotBeNull(); - state.Should().BeOfType>(); - state.Should().BeEmpty("Default SaveStateAsync should return an empty dictionary."); - - // Add a sample value and verify it updates correctly - state["testKey"] = "testValue"; - state.Should().ContainKey("testKey").WhoseValue.Should().Be("testValue"); - } } diff --git a/dotnet/test/Microsoft.AutoGen.Core.Tests/InProcessRuntimeTests.cs b/dotnet/test/Microsoft.AutoGen.Core.Tests/InProcessRuntimeTests.cs new file mode 100644 index 000000000..174f8b781 --- /dev/null +++ b/dotnet/test/Microsoft.AutoGen.Core.Tests/InProcessRuntimeTests.cs @@ -0,0 +1,141 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// InProcessRuntimeTests.cs +using System.Text.Json; +using FluentAssertions; +using Microsoft.AutoGen.Contracts; +using Microsoft.Extensions.Logging; +using Xunit; + +namespace Microsoft.AutoGen.Core.Tests; + +[Trait("Category", "UnitV2")] +public class InProcessRuntimeTests() +{ + // Agent will not deliver to self will success when runtime.DeliverToSelf is false (default) + [Fact] + public async Task RuntimeAgentPublishToSelfDefaultNoSendTest() + { + var runtime = new InProcessRuntime(); + await runtime.StartAsync(); + + Logger logger = new(new LoggerFactory()); + SubscribedSelfPublishAgent agent = null!; + + await runtime.RegisterAgentFactoryAsync("MyAgent", (id, runtime) => + { + agent = new SubscribedSelfPublishAgent(id, runtime, logger); + return ValueTask.FromResult(agent); + }); + + // Ensure the agent is actually created + AgentId agentId = await runtime.GetAgentAsync("MyAgent", lazy: false); + + // Validate agent ID + agentId.Should().Be(agent.Id, "Agent ID should match the registered agent"); + + await runtime.RegisterImplicitAgentSubscriptionsAsync("MyAgent"); + + var topicType = "TestTopic"; + + await runtime.PublishMessageAsync("SelfMessage", new TopicId(topicType)).ConfigureAwait(true); + + await runtime.RunUntilIdleAsync(); + + // Agent has default messages and could not publish to self + agent.Text.Source.Should().Be("DefaultTopic"); + agent.Text.Content.Should().Be("DefaultContent"); + } + + // Agent delivery to self will success when runtime.DeliverToSelf is true + [Fact] + public async Task RuntimeAgentPublishToSelfDeliverToSelfTrueTest() + { + var runtime = new InProcessRuntime(); + runtime.DeliverToSelf = true; + await runtime.StartAsync(); + + Logger logger = new(new LoggerFactory()); + SubscribedSelfPublishAgent agent = null!; + + await runtime.RegisterAgentFactoryAsync("MyAgent", (id, runtime) => + { + agent = new SubscribedSelfPublishAgent(id, runtime, logger); + return ValueTask.FromResult(agent); + }); + + // Ensure the agent is actually created + AgentId agentId = await runtime.GetAgentAsync("MyAgent", lazy: false); + + // Validate agent ID + agentId.Should().Be(agent.Id, "Agent ID should match the registered agent"); + + await runtime.RegisterImplicitAgentSubscriptionsAsync("MyAgent"); + + var topicType = "TestTopic"; + + await runtime.PublishMessageAsync("SelfMessage", new TopicId(topicType)).ConfigureAwait(true); + + await runtime.RunUntilIdleAsync(); + + // Agent sucessfully published to self + agent.Text.Source.Should().Be("TestTopic"); + agent.Text.Content.Should().Be("SelfMessage"); + } + + [Fact] + public async Task RuntimeShouldSaveLoadStateCorrectlyTest() + { + // Create a runtime and register an agent + var runtime = new InProcessRuntime(); + await runtime.StartAsync(); + Logger logger = new(new LoggerFactory()); + SubscribedSaveLoadAgent agent = null!; + await runtime.RegisterAgentFactoryAsync("MyAgent", (id, runtime) => + { + agent = new SubscribedSaveLoadAgent(id, runtime, logger); + return ValueTask.FromResult(agent); + }); + + // Get agent ID and instantiate agent by publishing + AgentId agentId = await runtime.GetAgentAsync("MyAgent", lazy: true); + await runtime.RegisterImplicitAgentSubscriptionsAsync("MyAgent"); + var topicType = "TestTopic"; + await runtime.PublishMessageAsync(new TextMessage { Source = topicType, Content = "test" }, new TopicId(topicType)).ConfigureAwait(true); + await runtime.RunUntilIdleAsync(); + agent.ReceivedMessages.Any().Should().BeTrue("Agent should receive messages when subscribed."); + + // Save the state + var savedState = await runtime.SaveStateAsync(); + + // Ensure saved state contains the agent's state + savedState.Should().ContainKey(agentId.ToString()); + + // Ensure the agent's state is stored as a valid JSON object + savedState[agentId.ToString()].ValueKind.Should().Be(JsonValueKind.Object, "Agent state should be stored as a JSON object"); + + // Serialize and Deserialize the state to simulate persistence + string json = JsonSerializer.Serialize(savedState); + json.Should().NotBeNullOrEmpty("Serialized state should not be empty"); + var deserializedState = JsonSerializer.Deserialize>(json) + ?? throw new Exception("Deserialized state is unexpectedly null"); + deserializedState.Should().ContainKey(agentId.ToString()); + + // Start new runtime and restore the state + var newRuntime = new InProcessRuntime(); + await newRuntime.StartAsync(); + await newRuntime.RegisterAgentFactoryAsync("MyAgent", (id, runtime) => + { + agent = new SubscribedSaveLoadAgent(id, runtime, logger); + return ValueTask.FromResult(agent); + }); + await newRuntime.RegisterImplicitAgentSubscriptionsAsync("MyAgent"); + + // Show that no agent instances exist in the new runtime + newRuntime.agentInstances.Count.Should().Be(0, "Agent should be registered in the new runtime"); + + // Load the state into the new runtime and show that agent is now instantiated + await newRuntime.LoadStateAsync(deserializedState); + newRuntime.agentInstances.Count.Should().Be(1, "Agent should be registered in the new runtime"); + newRuntime.agentInstances.Should().ContainKey(agentId, "Agent should be loaded into the new runtime"); + } +} diff --git a/dotnet/test/Microsoft.AutoGen.Core.Tests/TestAgent.cs b/dotnet/test/Microsoft.AutoGen.Core.Tests/TestAgent.cs index b6dadc833..ed87a7105 100644 --- a/dotnet/test/Microsoft.AutoGen.Core.Tests/TestAgent.cs +++ b/dotnet/test/Microsoft.AutoGen.Core.Tests/TestAgent.cs @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // TestAgent.cs +using System.Text.Json; using Microsoft.AutoGen.Contracts; using Microsoft.Extensions.Logging; @@ -59,7 +60,7 @@ public class TestAgent(AgentId id, /// Key: source /// Value: message /// - private readonly Dictionary _receivedMessages = new(); + protected Dictionary _receivedMessages = new(); public Dictionary ReceivedMessages => _receivedMessages; } @@ -73,6 +74,38 @@ public class SubscribedAgent : TestAgent } } +[TypeSubscription("TestTopic")] +public class SubscribedSaveLoadAgent : TestAgent +{ + public SubscribedSaveLoadAgent(AgentId id, + IAgentRuntime runtime, + Logger? logger = null) : base(id, runtime, logger) + { + } + + public override ValueTask> SaveStateAsync() + { + var jsonSafeDictionary = _receivedMessages.ToDictionary( + kvp => kvp.Key, + kvp => JsonSerializer.SerializeToElement(kvp.Value) // Convert each object to JsonElement + ); + + return ValueTask.FromResult>(jsonSafeDictionary); + } + + public override ValueTask LoadStateAsync(IDictionary state) + { + _receivedMessages.Clear(); + + foreach (var kvp in state) + { + _receivedMessages[kvp.Key] = kvp.Value.Deserialize() ?? throw new Exception($"Failed to deserialize key: {kvp.Key}"); + } + + return ValueTask.CompletedTask; + } +} + /// /// The test agent showing an agent that subscribes to itself. ///