.NET: Update AgentFeatureCollections with feedback (#2379)

* Update AgentFeatureCollections with feedback

* Address feedback.

* Fix issue with sample.

* Change generic type restriction to notnull

* Remove revision

* Update dotnet/src/Microsoft.Agents.AI.Abstractions/Features/AgentFeatureCollectionExtensions.cs

Co-authored-by: Copilot <[email protected]>

* Add revision back again and improve some formatting.

* Remove virtual from revision.

* Add overloads taking type as param and add unit tests.

---------

Co-authored-by: Copilot <[email protected]>

Author: westey-m

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

@@ -118,9 +118,8 @@ async Task CustomChatMessageStore_UsingFactoryAndExistingExternalId_Async()
 
     // It's possible to create a new thread that uses the same chat message store id by providing
     // the VectorChatMessageStoreThreadDbKeyFeature in the feature collection when creating the new thread.
-    AgentFeatureCollection features = new();
-    features.Set(new VectorChatMessageStoreThreadDbKeyFeature(messageStoreFromFactory.ThreadDbKey!));
-    AgentThread resumedThread = agent.GetNewThread(features);
+    AgentThread resumedThread = agent.GetNewThread(
+        new AgentFeatureCollection().WithFeature(new VectorChatMessageStoreThreadDbKeyFeature(messageStoreFromFactory.ThreadDbKey!)));
 
     // Run the agent with the thread that stores conversation history in the vector store.
     Console.WriteLine(await agent.RunAsync("Now tell the same joke in the voice of a pirate, and add some emojis to the joke.", resumedThread));
@@ -149,9 +148,7 @@ async Task CustomChatMessageStore_PerThread_Async()
     // We can then pass the feature collection when creating a new thread.
     // We also have the opportunity here to pass any id that we want for storing the chat history in the vector store.
     VectorChatMessageStore perThreadMessageStore = new(vectorStore, "chat-history-1");
-    AgentFeatureCollection features = new();
-    features.Set<ChatMessageStore>(perThreadMessageStore);
-    AgentThread thread = agent.GetNewThread(features);
+    AgentThread thread = agent.GetNewThread(new AgentFeatureCollection().WithFeature<ChatMessageStore>(perThreadMessageStore));
 
     Console.WriteLine(await agent.RunAsync("Tell me a joke about a pirate.", thread));
 
@@ -191,10 +188,13 @@ async Task CustomChatMessageStore_PerRun_Async()
     // If the agent doesn't support a message store, it would be ignored.
     // We also have the opportunity here to pass any id that we want for storing the chat history in the vector store.
     VectorChatMessageStore perRunMessageStore = new(vectorStore, "chat-history-1");
-    AgentFeatureCollection features = new();
-    features.Set<ChatMessageStore>(perRunMessageStore);
-
-    Console.WriteLine(await agent.RunAsync("Tell me a joke about a pirate.", thread, options: new AgentRunOptions() { Features = features }));
+    Console.WriteLine(await agent.RunAsync(
+        "Tell me a joke about a pirate.",
+        thread,
+        options: new AgentRunOptions()
+        {
+            Features = new AgentFeatureCollection().WithFeature<ChatMessageStore>(perRunMessageStore)
+        }));
 
     // When serializing this thread, we'll see that it has no messagestore state, since the messagestore was not attached to the thread,
     // but just provided for the single run. Note that, depending on the circumstances, the thread may still contain other state, e.g. Memories,
@@ -237,8 +237,9 @@ public VectorChatMessageStore(VectorStore vectorStore, JsonElement serializedSto
             // or finally we can generate one ourselves.
             this.ThreadDbKey = serializedStoreState.ValueKind is JsonValueKind.String
                 ? serializedStoreState.Deserialize<string>()
-                : features?.Get<VectorChatMessageStoreThreadDbKeyFeature>()?.ThreadDbKey
-                ?? Guid.NewGuid().ToString("N");
+                : features?.TryGet<VectorChatMessageStoreThreadDbKeyFeature>(out var threadDbKeyFeature) is true
+                    ? threadDbKeyFeature.ThreadDbKey
+                    : Guid.NewGuid().ToString("N");
         }
 
         public string? ThreadDbKey { get; }

dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs

@@ -55,7 +55,12 @@ public A2AAgent(A2AClient a2aClient, string? id = null, string? name = null, str
 
     /// <inheritdoc/>
     public sealed override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null)
-        => new A2AAgentThread() { ContextId = featureCollection?.Get<ConversationIdAgentFeature>()?.ConversationId };
+        => new A2AAgentThread()
+        {
+            ContextId = featureCollection?.TryGet<ConversationIdAgentFeature>(out var conversationIdFeature) is true
+                ? conversationIdFeature.ConversationId
+                : null
+        };
 
     /// <summary>
     /// Get a new <see cref="AgentThread"/> instance using an existing context id, to continue that conversation.

dotnet/src/Microsoft.Agents.AI.Abstractions/Features/AgentFeatureCollection.cs

@@ -4,6 +4,7 @@
 using System.Collections;
 using System.Collections.Generic;
 using System.Diagnostics;
+using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using Microsoft.Shared.Diagnostics;
 
@@ -18,9 +19,7 @@ namespace Microsoft.Agents.AI;
 [DebuggerTypeProxy(typeof(FeatureCollectionDebugView))]
 public class AgentFeatureCollection : IAgentFeatureCollection
 {
-    private static readonly KeyComparer s_featureKeyComparer = new();
-    private readonly IAgentFeatureCollection? _defaults;
-    private readonly int _initialCapacity;
+    private readonly IAgentFeatureCollection? _innerCollection;
     private Dictionary<Type, object>? _features;
     private volatile int _containerRevision;
 
@@ -39,124 +38,142 @@ public AgentFeatureCollection()
     public AgentFeatureCollection(int initialCapacity)
     {
         Throw.IfLessThan(initialCapacity, 0);
-
-        this._initialCapacity = initialCapacity;
+        this._features = new(initialCapacity);
     }
 
     /// <summary>
-    /// Initializes a new instance of <see cref="AgentFeatureCollection"/> with the specified defaults.
+    /// Initializes a new instance of <see cref="AgentFeatureCollection"/> with the specified inner collection.
     /// </summary>
-    /// <param name="defaults">The feature defaults.</param>
-    public AgentFeatureCollection(IAgentFeatureCollection defaults)
+    /// <param name="innerCollection">The inner collection.</param>
+    /// <remarks>
+    /// <para>
+    /// When providing an inner collection, and if a feature is not found in this collection,
+    /// an attempt will be made to retrieve it from the inner collection as a fallback.
+    /// </para>
+    /// <para>
+    /// The <see cref="Remove{TFeature}"/> method will only remove features from this collection
+    /// and not from the inner collection. When removing a feature from this collection, and
+    /// it exists in the inner collection, it will still be retrievable from the inner collection.
+    /// </para>
+    /// </remarks>
+    public AgentFeatureCollection(IAgentFeatureCollection innerCollection)
     {
-        this._defaults = defaults;
+        this._innerCollection = Throw.IfNull(innerCollection);
     }
 
     /// <inheritdoc />
-    public virtual int Revision
+    public int Revision
     {
-        get { return this._containerRevision + (this._defaults?.Revision ?? 0); }
+        get { return this._containerRevision + (this._innerCollection?.Revision ?? 0); }
     }
 
     /// <inheritdoc />
     public bool IsReadOnly { get { return false; } }
 
+    IEnumerator IEnumerable.GetEnumerator()
+    {
+        return this.GetEnumerator();
+    }
+
     /// <inheritdoc />
-    public object? this[Type key]
+    public IEnumerator<KeyValuePair<Type, object>> GetEnumerator()
     {
-        get
+        if (this._features is not { Count: > 0 })
         {
-            Throw.IfNull(key);
+            IEnumerable<KeyValuePair<Type, object>> e = ((IEnumerable<KeyValuePair<Type, object>>?)this._innerCollection) ?? [];
+            return e.GetEnumerator();
+        }
 
-            return this._features != null && this._features.TryGetValue(key, out var result) ? result : this._defaults?[key];
+        if (this._innerCollection is null)
+        {
+            return this._features.GetEnumerator();
         }
-        set
+
+        if (this._innerCollection is AgentFeatureCollection innerCollection && innerCollection._features is not { Count: > 0 })
         {
-            Throw.IfNull(key);
+            return this._features.GetEnumerator();
+        }
+
+        return YieldAll();
 
-            if (value == null)
+        IEnumerator<KeyValuePair<Type, object>> YieldAll()
+        {
+            HashSet<Type> set = [];
+
+            foreach (var entry in this._features)
             {
-                if (this._features?.Remove(key) is true)
-                {
-                    this._containerRevision++;
-                }
-                return;
+                set.Add(entry.Key);
+                yield return entry;
             }
 
-            if (this._features == null)
+            foreach (var entry in this._innerCollection.Where(x => !set.Contains(x.Key)))
             {
-                this._features = new Dictionary<Type, object>(this._initialCapacity);
+                yield return entry;
             }
-            this._features[key] = value;
-            this._containerRevision++;
         }
     }
 
-    IEnumerator IEnumerable.GetEnumerator()
-    {
-        return this.GetEnumerator();
-    }
-
     /// <inheritdoc />
-    public IEnumerator<KeyValuePair<Type, object>> GetEnumerator()
+    public bool TryGet<TFeature>([MaybeNullWhen(false)] out TFeature feature)
+        where TFeature : notnull
     {
-        if (this._features != null)
+        if (this.TryGet(typeof(TFeature), out var obj))
         {
-            foreach (var pair in this._features)
-            {
-                yield return pair;
-            }
+            feature = (TFeature)obj;
+            return true;
         }
 
-        if (this._defaults != null)
-        {
-            // Don't return features masked by the wrapper.
-            foreach (var pair in this._features == null ? this._defaults : this._defaults.Except(this._features, s_featureKeyComparer))
-            {
-                yield return pair;
-            }
-        }
+        feature = default;
+        return false;
     }
 
     /// <inheritdoc />
-    public TFeature? Get<TFeature>()
+    public bool TryGet(Type type, [MaybeNullWhen(false)] out object feature)
     {
-        if (typeof(TFeature).IsValueType)
+        if (this._features?.TryGetValue(type, out var obj) is true)
         {
-            var feature = this[typeof(TFeature)];
-            if (feature is null && Nullable.GetUnderlyingType(typeof(TFeature)) is null)
-            {
-                throw new InvalidOperationException(
-                    $"{typeof(TFeature).FullName} does not exist in the feature collection " +
-                    $"and because it is a struct the method can't return null. Use 'AgentFeatureCollection[typeof({typeof(TFeature).FullName})] is not null' to check if the feature exists.");
-            }
-            return (TFeature?)feature;
+            feature = obj;
+            return true;
         }
-        return (TFeature?)this[typeof(TFeature)];
+
+        if (this._innerCollection?.TryGet(type, out var defaultFeature) is true)
+        {
+            feature = defaultFeature;
+            return true;
+        }
+
+        feature = default;
+        return false;
     }
 
     /// <inheritdoc />
-    public void Set<TFeature>(TFeature? instance)
+    public void Set<TFeature>(TFeature instance)
+        where TFeature : notnull
     {
-        this[typeof(TFeature)] = instance;
+        Throw.IfNull(instance);
+
+        this._features ??= new();
+        this._features[typeof(TFeature)] = instance;
+        this._containerRevision++;
     }
 
-    // Used by the debugger. Count over enumerable is required to get the correct value.
-    private int GetCount() => this.Count();
+    /// <inheritdoc />
+    public void Remove<TFeature>()
+        where TFeature : notnull
+        => this.Remove(typeof(TFeature));
 
-    private sealed class KeyComparer : IEqualityComparer<KeyValuePair<Type, object>>
+    /// <inheritdoc />
+    public void Remove(Type type)
     {
-        public bool Equals(KeyValuePair<Type, object> x, KeyValuePair<Type, object> y)
+        if (this._features?.Remove(type) is true)
         {
-            return x.Key.Equals(y.Key);
-        }
-
-        public int GetHashCode(KeyValuePair<Type, object> obj)
-        {
-            return obj.Key.GetHashCode();
+            this._containerRevision++;
         }
     }
 
+    // Used by the debugger. Count over enumerable is required to get the correct value.
+    private int GetCount() => this.Count();
+
     private sealed class FeatureCollectionDebugView(AgentFeatureCollection features)
     {
         private readonly AgentFeatureCollection _features = features;

dotnet/src/Microsoft.Agents.AI.Abstractions/Features/AgentFeatureCollectionExtensions.cs

@@ -0,0 +1,23 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+namespace Microsoft.Agents.AI;
+
+/// <summary>
+/// Extension methods for <see cref="IAgentFeatureCollection"/>.
+/// </summary>
+public static class AgentFeatureCollectionExtensions
+{
+    /// <summary>
+    /// Adds the specified feature to the collection and returns the collection.
+    /// </summary>
+    /// <typeparam name="TFeature">The feature key.</typeparam>
+    /// <param name="features">The feature collection to add the new feature to.</param>
+    /// <param name="feature">The feature to add to the collection.</param>
+    /// <returns>The updated collection.</returns>
+    public static IAgentFeatureCollection WithFeature<TFeature>(this IAgentFeatureCollection features, TFeature feature)
+        where TFeature : notnull
+    {
+        features.Set(feature);
+        return features;
+    }
+}

dotnet/src/Microsoft.Agents.AI.Abstractions/Features/ConversationIdAgentFeature.cs

@@ -8,15 +8,15 @@ namespace Microsoft.Agents.AI;
 /// An agent feature that allows providing a conversation identifier.
 /// </summary>
 /// <remarks>
-/// This feature allows a user to provide a specific identifier for chat history whether stored in the underlying AI service or stored in a 3rd party store.
+/// This feature allows a user to provide a specific identifier for chat history when stored in the underlying AI service.
 /// </remarks>
 public class ConversationIdAgentFeature
 {
     /// <summary>
     /// Initializes a new instance of the <see cref="ConversationIdAgentFeature"/> class with the specified thread
     /// identifier.
     /// </summary>
-    /// <param name="conversationId">The unique identifier of the thread required by the underlying AI service or 3rd party store. Cannot be <see langword="null"/> or empty.</param>
+    /// <param name="conversationId">The unique identifier of the thread required by the underlying AI service. Cannot be <see langword="null"/> or empty.</param>
     public ConversationIdAgentFeature(string conversationId)
     {
         this.ConversationId = Throw.IfNullOrWhitespace(conversationId);