.NET: [BREAKING] Add GetService for AIContextProviders and ChatMessageStore (#874)

* Add GetService for AIContextProviders and ChatMessageStore

* Change styling and fix format issues

* Update sample code to allow for missing memory component.

dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyThreadStorage/Program.cs

@@ -70,14 +70,17 @@ AgentThread resumedThread = agent.DeserializeThread(serializedThread);
 // Run the agent with the thread that stores conversation history in the vector store a second time.
 Console.WriteLine(await agent.RunAsync("Now tell the same joke in the voice of a pirate, and add some emojis to the joke.", resumedThread));
 
+// We can access the VectorChatMessageStore via the thread's GetService method if we need to read the key under which threads are stored.
+var messageStore = resumedThread.GetService<VectorChatMessageStore>()!;
+Console.WriteLine($"\nThread is stored in vector store under key: {messageStore.ThreadDbKey}");
+
 namespace SampleApp
 {
     /// <summary>
-    /// A sample implementation of <see cref="IChatMessageStore"/> that stores chat messages in a vector store.
+    /// A sample implementation of <see cref="ChatMessageStore"/> that stores chat messages in a vector store.
     /// </summary>
-    internal sealed class VectorChatMessageStore : IChatMessageStore
+    internal sealed class VectorChatMessageStore : ChatMessageStore
     {
-        private string? _threadId;
         private readonly VectorStore _vectorStore;
 
         public VectorChatMessageStore(VectorStore vectorStore, JsonElement serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null)
@@ -87,35 +90,37 @@ namespace SampleApp
             if (serializedStoreState.ValueKind is JsonValueKind.String)
             {
                 // Here we can deserialize the thread id so that we can access the same messages as before the suspension.
-                this._threadId = serializedStoreState.Deserialize<string>();
+                this.ThreadDbKey = serializedStoreState.Deserialize<string>();
             }
         }
 
-        public async Task AddMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken)
+        public string? ThreadDbKey { get; private set; }
+
+        public override async Task AddMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken)
         {
-            this._threadId ??= Guid.NewGuid().ToString();
+            this.ThreadDbKey ??= Guid.NewGuid().ToString();
 
             var collection = this._vectorStore.GetCollection<string, ChatHistoryItem>("ChatHistory");
             await collection.EnsureCollectionExistsAsync(cancellationToken);
 
             await collection.UpsertAsync(messages.Select(x => new ChatHistoryItem()
             {
-                Key = this._threadId + x.MessageId,
+                Key = this.ThreadDbKey + x.MessageId,
                 Timestamp = DateTimeOffset.UtcNow,
-                ThreadId = this._threadId,
+                ThreadId = this.ThreadDbKey,
                 SerializedMessage = JsonSerializer.Serialize(x),
                 MessageText = x.Text
             }), cancellationToken);
         }
 
-        public async Task<IEnumerable<ChatMessage>> GetMessagesAsync(CancellationToken cancellationToken)
+        public override async Task<IEnumerable<ChatMessage>> GetMessagesAsync(CancellationToken cancellationToken)
         {
             var collection = this._vectorStore.GetCollection<string, ChatHistoryItem>("ChatHistory");
             await collection.EnsureCollectionExistsAsync(cancellationToken);
 
             var records = await collection
                 .GetAsync(
-                    x => x.ThreadId == this._threadId, 10,
+                    x => x.ThreadId == this.ThreadDbKey, 10,
                     new() { OrderBy = x => x.Descending(y => y.Timestamp) },
                     cancellationToken)
                 .ToListAsync(cancellationToken);
@@ -126,9 +131,9 @@ namespace SampleApp
             return messages;
         }
 
-        public ValueTask<JsonElement?> SerializeStateAsync(JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) =>
+        public override ValueTask<JsonElement?> SerializeStateAsync(JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) =>
             // We have to serialize the thread id, so that on deserialization we can retrieve the messages using the same thread id.
-            new(JsonSerializer.SerializeToElement(this._threadId));
+            new(JsonSerializer.SerializeToElement(this.ThreadDbKey));
 
         /// <summary>
         /// The data structure used to store chat history items in the vector store.

dotnet/samples/GettingStarted/Agents/Agent_Step13_Memory/Program.cs

@@ -62,20 +62,22 @@ Console.WriteLine(await agent.RunAsync("What is my name and age?", deserializedT
 
 Console.WriteLine("\n>> Read memories from memory component\n");
 
-// It's possible to access the memory component via the thread's AIContextProvider property.
-var userInfo = ((deserializedThread as ChatClientAgentThread)!.AIContextProvider as UserInfoMemory)!.UserInfo;
+// It's possible to access the memory component via the thread's GetService method.
+var userInfo = deserializedThread.GetService<UserInfoMemory>()?.UserInfo;
 
 // Output the user info that was captured by the memory component.
-Console.WriteLine($"MEMORY - User Name: {userInfo.UserName}");
-Console.WriteLine($"MEMORY - User Age: {userInfo.UserAge}");
+Console.WriteLine($"MEMORY - User Name: {userInfo?.UserName}");
+Console.WriteLine($"MEMORY - User Age: {userInfo?.UserAge}");
 
 Console.WriteLine("\n>> Use new thread with previously created memories\n");
 
 // It is also possible to set the memories in a memory component on an individual thread.
 // This is useful if we want to start a new thread, but have it share the same memories as a previous thread.
-// For this scenario, we have to know the underlying agent thread type and the memory component type.
 var newThread = agent.GetNewThread();
-((newThread as ChatClientAgentThread)!.AIContextProvider as UserInfoMemory)!.UserInfo = userInfo;
+if (userInfo is not null && newThread.GetService<UserInfoMemory>() is UserInfoMemory newThreadMemory)
+{
+    newThreadMemory.UserInfo = userInfo;
+}
 
 // Invoke the agent and output the text result.
 // This time the agent should remember the user's name and use it in the response.
@@ -150,12 +152,6 @@ namespace SampleApp
         {
             return new ValueTask<JsonElement?>(JsonSerializer.SerializeToElement(this.UserInfo, jsonSerializerOptions));
         }
-
-        public override ValueTask DeserializeAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
-        {
-            this.UserInfo = serializedState.Deserialize<UserInfo>(jsonSerializerOptions) ?? new UserInfo();
-            return default;
-        }
     }
 
     internal sealed class UserInfo

dotnet/src/Microsoft.Agents.Workflows/WorkflowMessageStore.cs

@@ -10,7 +10,7 @@ using Microsoft.Extensions.AI.Agents;
 
 namespace Microsoft.Agents.Workflows;
 
-internal sealed class WorkflowMessageStore : IChatMessageStore
+internal sealed class WorkflowMessageStore : ChatMessageStore
 {
     private int _bookmark;
     private readonly List<ChatMessage> _chatMessages = [];
@@ -46,14 +46,14 @@ internal sealed class WorkflowMessageStore : IChatMessageStore
 
     internal void AddMessages(params ChatMessage[] messages) => this._chatMessages.AddRange(messages);
 
-    public Task AddMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken)
+    public override Task AddMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken)
     {
         this._chatMessages.AddRange(messages);
 
         return Task.CompletedTask;
     }
 
-    public Task<IEnumerable<ChatMessage>> GetMessagesAsync(CancellationToken cancellationToken) => Task.FromResult<IEnumerable<ChatMessage>>(this._chatMessages.AsReadOnly());
+    public override Task<IEnumerable<ChatMessage>> GetMessagesAsync(CancellationToken cancellationToken) => Task.FromResult<IEnumerable<ChatMessage>>(this._chatMessages.AsReadOnly());
 
     public IEnumerable<ChatMessage> GetFromBookmark()
     {
@@ -65,7 +65,7 @@ internal sealed class WorkflowMessageStore : IChatMessageStore
 
     public void UpdateBookmark() => this._bookmark = this._chatMessages.Count;
 
-    public ValueTask<JsonElement?> SerializeStateAsync(JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+    public override ValueTask<JsonElement?> SerializeStateAsync(JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
     {
         StoreState state = new()
         {

dotnet/src/Microsoft.Extensions.AI.Agents.Abstractions/AIContextProvider.cs

@@ -5,6 +5,7 @@ using System.Collections.Generic;
 using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.Shared.Diagnostics;
 
 namespace Microsoft.Extensions.AI.Agents;
 
@@ -52,18 +53,35 @@ public abstract class AIContextProvider
         return default;
     }
 
-    /// <summary>
-    /// Deserializes the state contained in the provided <see cref="JsonElement"/> into the properties on this object.
-    /// </summary>
-    /// <param name="serializedState">A <see cref="JsonElement"/> representing the state of the object.</param>
-    /// <param name="jsonSerializerOptions">Optional settings for customizing the JSON deserialization process.</param>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns>A <see cref="ValueTask"/> that completes when the state has been deserialized.</returns>
-    public virtual ValueTask DeserializeAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+    /// <summary>Asks the <see cref="AIContextProvider"/> for an object of the specified type <paramref name="serviceType"/>.</summary>
+    /// <param name="serviceType">The type of object being requested.</param>
+    /// <param name="serviceKey">An optional key that can be used to help identify the target service.</param>
+    /// <returns>The found object, otherwise <see langword="null"/>.</returns>
+    /// <exception cref="ArgumentNullException"><paramref name="serviceType"/> is <see langword="null"/>.</exception>
+    /// <remarks>
+    /// The purpose of this method is to allow for the retrieval of strongly-typed services that might be provided by the <see cref="AIContextProvider"/>,
+    /// including itself or any services it might be wrapping.
+    /// </remarks>
+    public virtual object? GetService(Type serviceType, object? serviceKey = null)
     {
-        return default;
+        _ = Throw.IfNull(serviceType);
+
+        return serviceKey is null && serviceType.IsInstanceOfType(this)
+            ? this
+            : null;
     }
 
+    /// <summary>Asks the <see cref="AIContextProvider"/> for an object of type <typeparamref name="TService"/>.</summary>
+    /// <typeparam name="TService">The type of the object to be retrieved.</typeparam>
+    /// <param name="serviceKey">An optional key that can be used to help identify the target service.</param>
+    /// <returns>The found object, otherwise <see langword="null"/>.</returns>
+    /// <remarks>
+    /// The purpose of this method is to allow for the retrieval of strongly typed services that may be provided by the <see cref="AIContextProvider"/>,
+    /// including itself or any services it might be wrapping.
+    /// </remarks>
+    public TService? GetService<TService>(object? serviceKey = null)
+        => this.GetService(typeof(TService), serviceKey) is TService service ? service : default;
+
     /// <summary>
     /// Contains the event context provided to <see cref="InvokingAsync(InvokingContext, CancellationToken)"/>.
     /// </summary>

dotnet/src/Microsoft.Extensions.AI.Agents.Abstractions/ChatMessageStore.cs

@@ -0,0 +1,85 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Text.Json;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Shared.Diagnostics;
+
+namespace Microsoft.Extensions.AI.Agents;
+
+/// <summary>
+/// Defines methods for storing and retrieving chat messages associated with a specific thread.
+/// </summary>
+/// <remarks>
+/// Implementations of this interface are responsible for managing the storage of chat messages,
+/// including handling large volumes of data by truncating or summarizing messages as necessary.
+/// </remarks>
+public abstract class ChatMessageStore
+{
+    /// <summary>
+    /// Gets all the messages from the store that should be used for the next agent invocation.
+    /// </summary>
+    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
+    /// <returns>A collection of chat messages.</returns>
+    /// <remarks>
+    /// <para>
+    /// Messages are returned in ascending chronological order, with the oldest message first.
+    /// </para>
+    /// <para>
+    /// If the messages stored in the store become very large, it is up to the store to
+    /// truncate, summarize or otherwise limit the number of messages returned.
+    /// </para>
+    /// <para>
+    /// When using implementations of <see cref="ChatMessageStore"/>, a new one should be created for each thread
+    /// since they may contain state that is specific to a thread.
+    /// </para>
+    /// </remarks>
+    public abstract Task<IEnumerable<ChatMessage>> GetMessagesAsync(CancellationToken cancellationToken = default);
+
+    /// <summary>
+    /// Adds messages to the store.
+    /// </summary>
+    /// <param name="messages">The messages to add.</param>
+    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
+    /// <returns>An async task.</returns>
+    public abstract Task AddMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken = default);
+
+    /// <summary>
+    /// Serializes the current object's state to a <see cref="JsonElement"/> using the specified serialization options.
+    /// </summary>
+    /// <param name="jsonSerializerOptions">The JSON serialization options to use.</param>
+    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
+    /// <returns>A <see cref="JsonElement"/> representation of the object's state.</returns>
+    public abstract ValueTask<JsonElement?> SerializeStateAsync(JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default);
+
+    /// <summary>Asks the <see cref="ChatMessageStore"/> for an object of the specified type <paramref name="serviceType"/>.</summary>
+    /// <param name="serviceType">The type of object being requested.</param>
+    /// <param name="serviceKey">An optional key that can be used to help identify the target service.</param>
+    /// <returns>The found object, otherwise <see langword="null"/>.</returns>
+    /// <exception cref="ArgumentNullException"><paramref name="serviceType"/> is <see langword="null"/>.</exception>
+    /// <remarks>
+    /// The purpose of this method is to allow for the retrieval of strongly-typed services that might be provided by the <see cref="ChatMessageStore"/>,
+    /// including itself or any services it might be wrapping.
+    /// </remarks>
+    public virtual object? GetService(Type serviceType, object? serviceKey = null)
+    {
+        _ = Throw.IfNull(serviceType);
+
+        return serviceKey is null && serviceType.IsInstanceOfType(this)
+            ? this
+            : null;
+    }
+
+    /// <summary>Asks the <see cref="ChatMessageStore"/> for an object of type <typeparamref name="TService"/>.</summary>
+    /// <typeparam name="TService">The type of the object to be retrieved.</typeparam>
+    /// <param name="serviceKey">An optional key that can be used to help identify the target service.</param>
+    /// <returns>The found object, otherwise <see langword="null"/>.</returns>
+    /// <remarks>
+    /// The purpose of this method is to allow for the retrieval of strongly typed services that may be provided by the <see cref="ChatMessageStore"/>,
+    /// including itself or any services it might be wrapping.
+    /// </remarks>
+    public TService? GetService<TService>(object? serviceKey = null)
+        => this.GetService(typeof(TService), serviceKey) is TService service ? service : default;
+}

dotnet/src/Microsoft.Extensions.AI.Agents.Abstractions/IChatMessageStore.cs

@@ -1,54 +0,0 @@
-// Copyright (c) Microsoft. All rights reserved.
-
-using System.Collections.Generic;
-using System.Text.Json;
-using System.Threading;
-using System.Threading.Tasks;
-
-namespace Microsoft.Extensions.AI.Agents;
-
-/// <summary>
-/// Defines methods for storing and retrieving chat messages associated with a specific thread.
-/// </summary>
-/// <remarks>
-/// Implementations of this interface are responsible for managing the storage of chat messages,
-/// including handling large volumes of data by truncating or summarizing messages as necessary.
-/// </remarks>
-public interface IChatMessageStore
-{
-    /// <summary>
-    /// Gets all the messages from the store that should be used for the next agent invocation.
-    /// </summary>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns>A collection of chat messages.</returns>
-    /// <remarks>
-    /// <para>
-    /// Messages are returned in ascending chronological order, with the oldest message first.
-    /// </para>
-    /// <para>
-    /// If the messages stored in the store become very large, it is up to the store to
-    /// truncate, summarize or otherwise limit the number of messages returned.
-    /// </para>
-    /// <para>
-    /// When using implementations of <see cref="IChatMessageStore"/>, a new one should be created for each thread
-    /// since they may contain state that is specific to a thread.
-    /// </para>
-    /// </remarks>
-    Task<IEnumerable<ChatMessage>> GetMessagesAsync(CancellationToken cancellationToken = default);
-
-    /// <summary>
-    /// Adds messages to the store.
-    /// </summary>
-    /// <param name="messages">The messages to add.</param>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns>An async task.</returns>
-    Task AddMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken = default);
-
-    /// <summary>
-    /// Serializes the current object's state to a <see cref="JsonElement"/> using the specified serialization options.
-    /// </summary>
-    /// <param name="jsonSerializerOptions">The JSON serialization options to use.</param>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns>A <see cref="JsonElement"/> representation of the object's state.</returns>
-    ValueTask<JsonElement?> SerializeStateAsync(JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default);
-}

dotnet/src/Microsoft.Extensions.AI.Agents.Abstractions/InMemoryAgentThread.cs

@@ -85,6 +85,10 @@ public abstract class InMemoryAgentThread : AgentThread
         return JsonSerializer.SerializeToElement(state, AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(InMemoryAgentThreadState)));
     }
 
+    /// <inheritdoc/>
+    public override object? GetService(Type serviceType, object? serviceKey = null) =>
+        base.GetService(serviceType, serviceKey) ?? this.MessageStore?.GetService(serviceType, serviceKey);
+
     /// <inheritdoc />
     protected internal override Task MessagesReceivedAsync(IEnumerable<ChatMessage> newMessages, CancellationToken cancellationToken = default)
         => this.MessageStore.AddMessagesAsync(newMessages, cancellationToken);

dotnet/src/Microsoft.Extensions.AI.Agents.Abstractions/InMemoryChatMessageStore.cs

@@ -13,7 +13,7 @@ namespace Microsoft.Extensions.AI.Agents;
 /// <summary>
 /// Represents an in-memory store for chat messages associated with a specific thread.
 /// </summary>
-public sealed class InMemoryChatMessageStore : IList<ChatMessage>, IChatMessageStore
+public sealed class InMemoryChatMessageStore : ChatMessageStore, IList<ChatMessage>
 {
     private List<ChatMessage> _messages;
 
@@ -96,7 +96,7 @@ public sealed class InMemoryChatMessageStore : IList<ChatMessage>, IChatMessageS
     }
 
     /// <inheritdoc />
-    public async Task AddMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken)
+    public override async Task AddMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken)
     {
         _ = Throw.IfNull(messages);
 
@@ -109,7 +109,7 @@ public sealed class InMemoryChatMessageStore : IList<ChatMessage>, IChatMessageS
     }
 
     /// <inheritdoc />
-    public async Task<IEnumerable<ChatMessage>> GetMessagesAsync(CancellationToken cancellationToken)
+    public override async Task<IEnumerable<ChatMessage>> GetMessagesAsync(CancellationToken cancellationToken)
     {
         if (this.ReducerTriggerEvent is ChatReducerTriggerEvent.BeforeMessagesRetrieval && this.ChatReducer is not null)
         {
@@ -120,7 +120,7 @@ public sealed class InMemoryChatMessageStore : IList<ChatMessage>, IChatMessageS
     }
 
     /// <inheritdoc />
-    public ValueTask<JsonElement?> SerializeStateAsync(JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+    public override ValueTask<JsonElement?> SerializeStateAsync(JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
     {
         StoreState state = new()
         {

dotnet/src/Microsoft.Extensions.AI.Agents/ChatCompletion/ChatClientAgent.cs

@@ -262,7 +262,7 @@ public sealed class ChatClientAgent : AIAgent
     /// <inheritdoc/>
     public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
     {
-        Func<JsonElement, JsonSerializerOptions?, IChatMessageStore>? chatMessageStoreFactory = this._agentOptions?.ChatMessageStoreFactory is null ?
+        Func<JsonElement, JsonSerializerOptions?, ChatMessageStore>? chatMessageStoreFactory = this._agentOptions?.ChatMessageStoreFactory is null ?
             null :
             (jse, jso) => this._agentOptions.ChatMessageStoreFactory.Invoke(new() { SerializedState = jse, JsonSerializerOptions = jso });
 

dotnet/src/Microsoft.Extensions.AI.Agents/ChatCompletion/ChatClientAgentOptions.cs

@@ -76,10 +76,10 @@ public class ChatClientAgentOptions
     public ChatOptions? ChatOptions { get; set; }
 
     /// <summary>
-    /// Gets or sets a factory function to create an instance of <see cref="IChatMessageStore"/>
+    /// Gets or sets a factory function to create an instance of <see cref="ChatMessageStore"/>
     /// which will be used to store chat messages for this agent.
     /// </summary>
-    public Func<ChatMessageStoreFactoryContext, IChatMessageStore>? ChatMessageStoreFactory { get; set; }
+    public Func<ChatMessageStoreFactoryContext, ChatMessageStore>? ChatMessageStoreFactory { get; set; }
 
     /// <summary>
     /// Gets or sets a factory function to create an instance of <see cref="AIContextProvider"/>
@@ -135,14 +135,14 @@ public class ChatClientAgentOptions
     }
 
     /// <summary>
-    /// Context object passed to the <see cref="ChatMessageStoreFactory"/> to create a new instance of <see cref="IChatMessageStore"/>.
+    /// Context object passed to the <see cref="ChatMessageStoreFactory"/> to create a new instance of <see cref="ChatMessageStore"/>.
     /// </summary>
     public class ChatMessageStoreFactoryContext
     {
         /// <summary>
         /// Gets or sets the serialized state of the chat message store, if any.
         /// </summary>
-        /// <value><see langword="default"/> if there is no state, e.g. when the <see cref="IChatMessageStore"/> is first created.</value>
+        /// <value><see langword="default"/> if there is no state, e.g. when the <see cref="ChatMessageStore"/> is first created.</value>
         public JsonElement SerializedState { get; set; }
 
         /// <summary>

dotnet/src/Microsoft.Extensions.AI.Agents/ChatCompletion/ChatClientAgentThread.cs

@@ -16,7 +16,7 @@ namespace Microsoft.Extensions.AI.Agents;
 public class ChatClientAgentThread : AgentThread
 {
     private string? _conversationId;
-    private IChatMessageStore? _messageStore;
+    private ChatMessageStore? _messageStore;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="ChatClientAgentThread"/> class.
@@ -30,12 +30,12 @@ public class ChatClientAgentThread : AgentThread
     /// </summary>
     /// <param name="serializedThreadState">A <see cref="JsonElement"/> representing the serialized state of the thread.</param>
     /// <param name="jsonSerializerOptions">Optional settings for customizing the JSON deserialization process.</param>
-    /// <param name="chatMessageStoreFactory">An optional factory function to create a custom <see cref="IChatMessageStore"/>.</param>
+    /// <param name="chatMessageStoreFactory">An optional factory function to create a custom <see cref="ChatMessageStore"/>.</param>
     /// <param name="aiContextProviderFactory">An optional factory function to create a custom <see cref="AIContextProvider"/>.</param>
     internal ChatClientAgentThread(
         JsonElement serializedThreadState,
         JsonSerializerOptions? jsonSerializerOptions = null,
-        Func<JsonElement, JsonSerializerOptions?, IChatMessageStore>? chatMessageStoreFactory = null,
+        Func<JsonElement, JsonSerializerOptions?, ChatMessageStore>? chatMessageStoreFactory = null,
         Func<JsonElement, JsonSerializerOptions?, AIContextProvider>? aiContextProviderFactory = null)
     {
         if (serializedThreadState.ValueKind != JsonValueKind.Object)
@@ -77,7 +77,7 @@ public class ChatClientAgentThread : AgentThread
     /// <para>
     /// This property may be null in the following cases:
     /// <list type="bullet">
-    /// <item>The thread stores messages via the <see cref="IChatMessageStore"/> and not in the agent service.</item>
+    /// <item>The thread stores messages via the <see cref="ChatMessageStore"/> and not in the agent service.</item>
     /// <item>This thread object is new and a server managed thread has not yet been created in the agent service.</item>
     /// </list>
     /// </para>
@@ -110,7 +110,7 @@ public class ChatClientAgentThread : AgentThread
     }
 
     /// <summary>
-    /// Gets or sets the <see cref="IChatMessageStore"/> used by this thread, for cases where messages should be stored in a custom location.
+    /// Gets or sets the <see cref="ChatMessageStore"/> used by this thread, for cases where messages should be stored in a custom location.
     /// </summary>
     /// <remarks>
     /// <para>
@@ -121,12 +121,12 @@ public class ChatClientAgentThread : AgentThread
     /// <para>
     /// This property may be null in the following cases:
     /// <list type="bullet">
-    /// <item>The thread stores messages in the agent service and just has an id to the remove thread, instead of in an <see cref="IChatMessageStore"/>.</item>
-    /// <item>This thread object is new it is not yet clear whether it will be backed by a server managed thread or an <see cref="IChatMessageStore"/>.</item>
+    /// <item>The thread stores messages in the agent service and just has an id to the remove thread, instead of in an <see cref="ChatMessageStore"/>.</item>
+    /// <item>This thread object is new it is not yet clear whether it will be backed by a server managed thread or an <see cref="ChatMessageStore"/>.</item>
     /// </list>
     /// </para>
     /// </remarks>
-    public IChatMessageStore? MessageStore
+    public ChatMessageStore? MessageStore
     {
         get => this._messageStore;
         internal set
@@ -179,12 +179,12 @@ public class ChatClientAgentThread : AgentThread
     }
 
     /// <inheritdoc/>
-    public override object? GetService(Type serviceType, object? serviceKey = null)
-    {
-        return serviceType == typeof(AgentThreadMetadata) ?
-            new AgentThreadMetadata(this.ConversationId) :
-            base.GetService(serviceType, serviceKey);
-    }
+    public override object? GetService(Type serviceType, object? serviceKey = null) =>
+        serviceType == typeof(AgentThreadMetadata)
+            ? new AgentThreadMetadata(this.ConversationId)
+            : base.GetService(serviceType, serviceKey)
+            ?? this.AIContextProvider?.GetService(serviceType, serviceKey)
+            ?? this.MessageStore?.GetService(serviceType, serviceKey);
 
     /// <inheritdoc />
     protected override async Task MessagesReceivedAsync(IEnumerable<ChatMessage> newMessages, CancellationToken cancellationToken = default)

dotnet/tests/Microsoft.Extensions.AI.Agents.Abstractions.UnitTests/AIContextProviderTests.cs

@@ -27,15 +27,6 @@ public class AIContextProviderTests
         Assert.Equal(default, actual);
     }
 
-    [Fact]
-    public async Task DeserializeAsync_ReturnsCompletedTaskAsync()
-    {
-        var provider = new TestAIContextProvider();
-        var element = default(JsonElement);
-        var actual = provider.DeserializeAsync(element);
-        Assert.Equal(default, actual);
-    }
-
     [Fact]
     public void InvokingContext_Constructor_ThrowsForNullMessages()
     {
@@ -48,6 +39,122 @@ public class AIContextProviderTests
         Assert.Throws<ArgumentNullException>(() => new AIContextProvider.InvokedContext(null!));
     }
 
+    #region GetService Method Tests
+
+    /// <summary>
+    /// Verify that GetService returns the context provider itself when requesting the exact context provider type.
+    /// </summary>
+    [Fact]
+    public void GetService_RequestingExactContextProviderType_ReturnsContextProvider()
+    {
+        // Arrange
+        var contextProvider = new TestAIContextProvider();
+
+        // Act
+        var result = contextProvider.GetService(typeof(TestAIContextProvider));
+
+        // Assert
+        Assert.NotNull(result);
+        Assert.Same(contextProvider, result);
+    }
+
+    /// <summary>
+    /// Verify that GetService returns the context provider itself when requesting the base AIContextProvider type.
+    /// </summary>
+    [Fact]
+    public void GetService_RequestingAIContextProviderType_ReturnsContextProvider()
+    {
+        // Arrange
+        var contextProvider = new TestAIContextProvider();
+
+        // Act
+        var result = contextProvider.GetService(typeof(AIContextProvider));
+
+        // Assert
+        Assert.NotNull(result);
+        Assert.Same(contextProvider, result);
+    }
+
+    /// <summary>
+    /// Verify that GetService returns null when requesting an unrelated type.
+    /// </summary>
+    [Fact]
+    public void GetService_RequestingUnrelatedType_ReturnsNull()
+    {
+        // Arrange
+        var contextProvider = new TestAIContextProvider();
+
+        // Act
+        var result = contextProvider.GetService(typeof(string));
+
+        // Assert
+        Assert.Null(result);
+    }
+
+    /// <summary>
+    /// Verify that GetService returns null when a service key is provided, even for matching types.
+    /// </summary>
+    [Fact]
+    public void GetService_WithServiceKey_ReturnsNull()
+    {
+        // Arrange
+        var contextProvider = new TestAIContextProvider();
+
+        // Act
+        var result = contextProvider.GetService(typeof(TestAIContextProvider), "some-key");
+
+        // Assert
+        Assert.Null(result);
+    }
+
+    /// <summary>
+    /// Verify that GetService throws ArgumentNullException when serviceType is null.
+    /// </summary>
+    [Fact]
+    public void GetService_WithNullServiceType_ThrowsArgumentNullException()
+    {
+        // Arrange
+        var contextProvider = new TestAIContextProvider();
+
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>(() => contextProvider.GetService(null!));
+    }
+
+    /// <summary>
+    /// Verify that GetService generic method works correctly.
+    /// </summary>
+    [Fact]
+    public void GetService_Generic_ReturnsCorrectType()
+    {
+        // Arrange
+        var contextProvider = new TestAIContextProvider();
+
+        // Act
+        var result = contextProvider.GetService<TestAIContextProvider>();
+
+        // Assert
+        Assert.NotNull(result);
+        Assert.Same(contextProvider, result);
+    }
+
+    /// <summary>
+    /// Verify that GetService generic method returns null for unrelated types.
+    /// </summary>
+    [Fact]
+    public void GetService_Generic_ReturnsNullForUnrelatedType()
+    {
+        // Arrange
+        var contextProvider = new TestAIContextProvider();
+
+        // Act
+        var result = contextProvider.GetService<string>();
+
+        // Assert
+        Assert.Null(result);
+    }
+
+    #endregion
+
     private sealed class TestAIContextProvider : AIContextProvider
     {
         public override ValueTask<AIContext> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default)
@@ -59,10 +166,5 @@ public class AIContextProviderTests
         {
             return await base.SerializeAsync(jsonSerializerOptions, cancellationToken);
         }
-
-        public override async ValueTask DeserializeAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
-        {
-            await base.DeserializeAsync(serializedState, jsonSerializerOptions, cancellationToken);
-        }
     }
 }

dotnet/tests/Microsoft.Extensions.AI.Agents.Abstractions.UnitTests/ChatMessageStoreTests.cs

@@ -0,0 +1,89 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Text.Json;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.Extensions.AI.Agents.Abstractions.UnitTests;
+
+/// <summary>
+/// Contains tests for the <see cref="ChatMessageStore"/> class.
+/// </summary>
+public class ChatMessageStoreTests
+{
+    #region GetService Method Tests
+
+    [Fact]
+    public void GetService_RequestingExactStoreType_ReturnsStore()
+    {
+        var store = new TestChatMessageStore();
+        var result = store.GetService(typeof(TestChatMessageStore));
+        Assert.NotNull(result);
+        Assert.Same(store, result);
+    }
+
+    [Fact]
+    public void GetService_RequestingBaseStoreType_ReturnsStore()
+    {
+        var store = new TestChatMessageStore();
+        var result = store.GetService(typeof(ChatMessageStore));
+        Assert.NotNull(result);
+        Assert.Same(store, result);
+    }
+
+    [Fact]
+    public void GetService_RequestingUnrelatedType_ReturnsNull()
+    {
+        var store = new TestChatMessageStore();
+        var result = store.GetService(typeof(string));
+        Assert.Null(result);
+    }
+
+    [Fact]
+    public void GetService_WithServiceKey_ReturnsNull()
+    {
+        var store = new TestChatMessageStore();
+        var result = store.GetService(typeof(TestChatMessageStore), "some-key");
+        Assert.Null(result);
+    }
+
+    [Fact]
+    public void GetService_WithNullServiceType_ThrowsArgumentNullException()
+    {
+        var store = new TestChatMessageStore();
+        Assert.Throws<ArgumentNullException>(() => store.GetService(null!));
+    }
+
+    [Fact]
+    public void GetService_Generic_ReturnsCorrectType()
+    {
+        var store = new TestChatMessageStore();
+        var result = store.GetService<TestChatMessageStore>();
+        Assert.NotNull(result);
+        Assert.Same(store, result);
+    }
+
+    [Fact]
+    public void GetService_Generic_ReturnsNullForUnrelatedType()
+    {
+        var store = new TestChatMessageStore();
+        var result = store.GetService<string>();
+        Assert.Null(result);
+    }
+
+    #endregion
+
+    private sealed class TestChatMessageStore : ChatMessageStore
+    {
+        public override Task<IEnumerable<ChatMessage>> GetMessagesAsync(CancellationToken cancellationToken = default)
+            => Task.FromResult<IEnumerable<ChatMessage>>(Array.Empty<ChatMessage>());
+
+        public override Task AddMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken = default)
+            => Task.CompletedTask;
+
+        public override ValueTask<JsonElement?> SerializeStateAsync(JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+            => default;
+    }
+}

dotnet/tests/Microsoft.Extensions.AI.Agents.Abstractions.UnitTests/InMemoryAgentThreadTests.cs

@@ -131,6 +131,22 @@ public class InMemoryAgentThreadTests
 
     #endregion
 
+    #region GetService Tests
+
+    [Fact]
+    public void GetService_RequestingChatMessageStore_ReturnsChatMessageStore()
+    {
+        // Arrange
+        var thread = new TestInMemoryAgentThread();
+
+        // Act & Assert
+        Assert.NotNull(thread.GetService(typeof(ChatMessageStore)));
+        Assert.Same(thread.GetMessageStore(), thread.GetService(typeof(ChatMessageStore)));
+        Assert.Same(thread.GetMessageStore(), thread.GetService(typeof(InMemoryChatMessageStore)));
+    }
+
+    #endregion
+
     // Sealed test subclass to expose protected members for testing
     private sealed class TestInMemoryAgentThread : InMemoryAgentThread
     {

dotnet/tests/Microsoft.Extensions.AI.Agents.UnitTests/ChatCompletion/ChatClientAgentOptionsTests.cs

@@ -164,7 +164,7 @@ public class ChatClientAgentOptionsTests
         const string Name = "Test name";
         const string Description = "Test description";
         var tools = new List<AITool> { AIFunctionFactory.Create(() => "test") };
-        static IChatMessageStore ChatMessageStoreFactory(ChatClientAgentOptions.ChatMessageStoreFactoryContext ctx) => new Mock<IChatMessageStore>().Object;
+        static ChatMessageStore ChatMessageStoreFactory(ChatClientAgentOptions.ChatMessageStoreFactoryContext ctx) => new Mock<ChatMessageStore>().Object;
         static AIContextProvider AIContextProviderFactory(ChatClientAgentOptions.AIContextProviderFactoryContext ctx) => new Mock<AIContextProvider>().Object;
 
         var original = new ChatClientAgentOptions(Instructions, Name, Description, tools)

dotnet/tests/Microsoft.Extensions.AI.Agents.UnitTests/ChatCompletion/ChatClientAgentTests.cs

@@ -1707,7 +1707,7 @@ public class ChatClientAgentTests
     {
         // Arrange
         var mockChatClient = new Mock<IChatClient>();
-        var mockStore = new Mock<IChatMessageStore>();
+        var mockStore = new Mock<ChatMessageStore>();
         var factoryCalled = false;
 
         var agent = new ChatClientAgent(mockChatClient.Object, new ChatClientAgentOptions

dotnet/tests/Microsoft.Extensions.AI.Agents.UnitTests/ChatCompletion/ChatClientAgentThreadTests.cs

@@ -308,7 +308,7 @@ public class ChatClientAgentThreadTests
             new Dictionary<string, object> { ["Key"] = "TestValue" },
             TestJsonSerializerContext.Default.DictionaryStringObject);
 
-        var messageStoreMock = new Mock<IChatMessageStore>();
+        var messageStoreMock = new Mock<ChatMessageStore>();
         messageStoreMock
             .Setup(m => m.SerializeStateAsync(options, It.IsAny<CancellationToken>()))
             .ReturnsAsync(storeStateElement);
@@ -333,6 +333,45 @@ public class ChatClientAgentThreadTests
 
     #endregion Serialize Tests
 
+    #region GetService Tests
+
+    [Fact]
+    public void GetService_RequestingAIContextProvider_ReturnsAIContextProvider()
+    {
+        // Arrange
+        var thread = new ChatClientAgentThread();
+        var mockProvider = new Mock<AIContextProvider>();
+        mockProvider
+            .Setup(m => m.GetService(It.Is<Type>(x => x == typeof(AIContextProvider)), null))
+            .Returns(mockProvider.Object);
+        thread.AIContextProvider = mockProvider.Object;
+
+        // Act
+        var result = thread.GetService(typeof(AIContextProvider));
+
+        // Assert
+        Assert.NotNull(result);
+        Assert.Same(mockProvider.Object, result);
+    }
+
+    [Fact]
+    public void GetService_RequestingChatMessageStore_ReturnsChatMessageStore()
+    {
+        // Arrange
+        var thread = new ChatClientAgentThread();
+        var messageStore = new InMemoryChatMessageStore();
+        thread.MessageStore = messageStore;
+
+        // Act
+        var result = thread.GetService(typeof(ChatMessageStore));
+
+        // Assert
+        Assert.NotNull(result);
+        Assert.Same(messageStore, result);
+    }
+
+    #endregion
+
     private sealed class MessageSendingAgent : AIAgent
     {
         public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)