Add GetService<>() capabilities to interrogate AgentSessionStore delegation chain

Co-authored-by: lokitoth <6936551+lokitoth@users.noreply.github.com>

dotnet/src/Microsoft.Agents.AI.Hosting/AgentSessionStore.cs

@@ -1,7 +1,9 @@
 // Copyright (c) Microsoft. All rights reserved.
 
+using System;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.Shared.Diagnostics;
 
 namespace Microsoft.Agents.AI.Hosting;
 
@@ -43,4 +45,35 @@ public abstract class AgentSessionStore
         AIAgent agent,
         string conversationId,
         CancellationToken cancellationToken = default);
+
+    /// <summary>Asks the <see cref="AgentSessionStore"/> 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="AgentSessionStore"/>,
+    /// including itself or any services it might be wrapping. This is particularly useful for inspecting delegation chains
+    /// to verify that specific store implementations are present.
+    /// </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="AgentSessionStore"/> 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="AgentSessionStore"/>,
+    /// including itself or any services it might be wrapping. This is particularly useful for inspecting delegation chains
+    /// to verify that specific store implementations are present.
+    /// </remarks>
+    public TService? GetService<TService>(object? serviceKey = null)
+        => this.GetService(typeof(TService), serviceKey) is TService service ? service : default;
 }

dotnet/src/Microsoft.Agents.AI.Hosting/DelegatingAgentSessionStore.cs

@@ -59,4 +59,23 @@ public abstract class DelegatingAgentSessionStore : AgentSessionStore
     /// <inheritdoc/>
     public override ValueTask SaveSessionAsync(AIAgent agent, string conversationId, AgentSession session, CancellationToken cancellationToken = default)
         => this.InnerStore.SaveSessionAsync(agent, conversationId, session, cancellationToken);
+
+    /// <inheritdoc/>
+    /// <remarks>
+    /// This implementation first checks if this instance satisfies the service request.
+    /// If not, it chains the request to the inner store, allowing services to be retrieved
+    /// from any store in the delegation chain.
+    /// </remarks>
+    public override object? GetService(Type serviceType, object? serviceKey = null)
+    {
+        // First, check if this instance satisfies the request
+        object? service = base.GetService(serviceType, serviceKey);
+        if (service is not null)
+        {
+            return service;
+        }
+
+        // Chain to the inner store
+        return this.InnerStore.GetService(serviceType, serviceKey);
+    }
 }

dotnet/tests/Microsoft.Agents.AI.Hosting.UnitTests/DelegatingAgentSessionStoreTests.cs

@@ -191,6 +191,182 @@ public class DelegatingAgentSessionStoreTests
 
     #endregion
 
+    #region GetService Tests
+
+    /// <summary>
+    /// Verify that GetService returns itself when requesting the exact type.
+    /// </summary>
+    [Fact]
+    public void GetServiceReturnsItselfForExactType()
+    {
+        // Act
+        var result = this._delegatingStore.GetService(typeof(TestDelegatingAgentSessionStore));
+
+        // Assert
+        Assert.Same(this._delegatingStore, result);
+    }
+
+    /// <summary>
+    /// Verify that GetService returns itself when requesting a base type.
+    /// </summary>
+    [Fact]
+    public void GetServiceReturnsItselfForBaseType()
+    {
+        // Act
+        var result = this._delegatingStore.GetService(typeof(DelegatingAgentSessionStore));
+
+        // Assert
+        Assert.Same(this._delegatingStore, result);
+    }
+
+    /// <summary>
+    /// Verify that GetService returns itself when requesting AgentSessionStore.
+    /// </summary>
+    [Fact]
+    public void GetServiceReturnsItselfForAgentSessionStoreType()
+    {
+        // Act
+        var result = this._delegatingStore.GetService(typeof(AgentSessionStore));
+
+        // Assert
+        Assert.Same(this._delegatingStore, result);
+    }
+
+    /// <summary>
+    /// Verify that GetService chains to inner store when type is not satisfied by outer store.
+    /// </summary>
+    [Fact]
+    public void GetServiceChainsToInnerStore()
+    {
+        // Arrange
+        var innerStore = new ConcreteAgentSessionStore();
+        var delegatingStore = new TestDelegatingAgentSessionStore(innerStore);
+
+        // Act
+        var result = delegatingStore.GetService(typeof(ConcreteAgentSessionStore));
+
+        // Assert
+        Assert.Same(innerStore, result);
+    }
+
+    /// <summary>
+    /// Verify that GetService chains through multiple delegation layers.
+    /// </summary>
+    [Fact]
+    public void GetServiceChainsThoughMultipleDelegationLayers()
+    {
+        // Arrange - create a three-layer chain: outer -> middle -> inner
+        var innerStore = new ConcreteAgentSessionStore();
+        var middleStore = new AnotherDelegatingAgentSessionStore(innerStore);
+        var outerStore = new TestDelegatingAgentSessionStore(middleStore);
+
+        // Act - request the innermost store type
+        var result = outerStore.GetService(typeof(ConcreteAgentSessionStore));
+
+        // Assert
+        Assert.Same(innerStore, result);
+    }
+
+    /// <summary>
+    /// Verify that GetService can find a store in the middle of the delegation chain.
+    /// </summary>
+    [Fact]
+    public void GetServiceFindsMiddleStoreInChain()
+    {
+        // Arrange - create a three-layer chain: outer -> middle -> inner
+        var innerStore = new ConcreteAgentSessionStore();
+        var middleStore = new AnotherDelegatingAgentSessionStore(innerStore);
+        var outerStore = new TestDelegatingAgentSessionStore(middleStore);
+
+        // Act - request the middle store type
+        var result = outerStore.GetService(typeof(AnotherDelegatingAgentSessionStore));
+
+        // Assert
+        Assert.Same(middleStore, result);
+    }
+
+    /// <summary>
+    /// Verify that GetService returns null when the requested type is not found in the chain.
+    /// </summary>
+    [Fact]
+    public void GetServiceReturnsNullWhenTypeNotFound()
+    {
+        // Arrange
+        var innerStore = new ConcreteAgentSessionStore();
+        var delegatingStore = new TestDelegatingAgentSessionStore(innerStore);
+
+        // Act
+        var result = delegatingStore.GetService(typeof(string));
+
+        // Assert
+        Assert.Null(result);
+    }
+
+    /// <summary>
+    /// Verify that GetService returns null when a service key is provided but not matched.
+    /// </summary>
+    [Fact]
+    public void GetServiceReturnsNullWhenServiceKeyProvided()
+    {
+        // Act
+        var result = this._delegatingStore.GetService(typeof(TestDelegatingAgentSessionStore), "some-key");
+
+        // Assert
+        Assert.Null(result);
+    }
+
+    /// <summary>
+    /// Verify that GetService throws ArgumentNullException when serviceType is null.
+    /// </summary>
+    [Fact]
+    public void GetServiceThrowsWhenServiceTypeIsNull() =>
+        Assert.Throws<ArgumentNullException>("serviceType", () => this._delegatingStore.GetService(null!));
+
+    /// <summary>
+    /// Verify that GetService generic method works correctly.
+    /// </summary>
+    [Fact]
+    public void GetServiceGenericReturnsItself()
+    {
+        // Act
+        var result = this._delegatingStore.GetService<TestDelegatingAgentSessionStore>();
+
+        // Assert
+        Assert.Same(this._delegatingStore, result);
+    }
+
+    /// <summary>
+    /// Verify that GetService generic method chains to inner store.
+    /// </summary>
+    [Fact]
+    public void GetServiceGenericChainsToInnerStore()
+    {
+        // Arrange
+        var innerStore = new ConcreteAgentSessionStore();
+        var delegatingStore = new TestDelegatingAgentSessionStore(innerStore);
+
+        // Act
+        var result = delegatingStore.GetService<ConcreteAgentSessionStore>();
+
+        // Assert
+        Assert.Same(innerStore, result);
+    }
+
+    /// <summary>
+    /// Verify that GetService generic method returns null when type not found.
+    /// </summary>
+    [Fact]
+    public void GetServiceGenericReturnsNullWhenTypeNotFound()
+    {
+        // Act
+        var result = this._delegatingStore.GetService<string>();
+
+        // Assert
+        Assert.Null(result);
+    }
+
+    #endregion
+
     #region Test Implementation
 
     /// <summary>
@@ -201,6 +377,23 @@ public class DelegatingAgentSessionStoreTests
         public new AgentSessionStore InnerStore => base.InnerStore;
     }
 
+    /// <summary>
+    /// Another delegating store implementation for testing multi-layer chains.
+    /// </summary>
+    private sealed class AnotherDelegatingAgentSessionStore(AgentSessionStore innerStore) : DelegatingAgentSessionStore(innerStore);
+
+    /// <summary>
+    /// Concrete (non-delegating) session store for testing GetService chaining.
+    /// </summary>
+    private sealed class ConcreteAgentSessionStore : AgentSessionStore
+    {
+        public override ValueTask<AgentSession> GetSessionAsync(AIAgent agent, string conversationId, CancellationToken cancellationToken = default)
+            => new(new TestAgentSession());
+
+        public override ValueTask SaveSessionAsync(AIAgent agent, string conversationId, AgentSession session, CancellationToken cancellationToken = default)
+            => ValueTask.CompletedTask;
+    }
+
     private sealed class TestAgentSession : AgentSession;
 
     #endregion

dotnet/tests/Microsoft.Agents.AI.Hosting.UnitTests/IsolationKeyScopedAgentSessionStoreTests.cs

@@ -364,6 +364,45 @@ public class IsolationKeyScopedAgentSessionStoreTests
 
     #endregion
 
+    #region GetService Tests
+
+    /// <summary>
+    /// Verify that GetService can retrieve IsolationKeyScopedAgentSessionStore from a delegation chain.
+    /// </summary>
+    [Fact]
+    public void GetServiceReturnsIsolationKeyScopedAgentSessionStore()
+    {
+        // Arrange
+        var provider = new TestSessionIsolationKeyProvider(TestIsolationKey);
+        var store = new IsolationKeyScopedAgentSessionStore(this._innerStoreMock.Object, provider);
+
+        // Act
+        var result = store.GetService<IsolationKeyScopedAgentSessionStore>();
+
+        // Assert
+        Assert.Same(store, result);
+    }
+
+    /// <summary>
+    /// Verify that GetService chains through to find inner store types.
+    /// </summary>
+    [Fact]
+    public void GetServiceChainsToInnerStore()
+    {
+        // Arrange
+        var concreteInnerStore = new ConcreteAgentSessionStore();
+        var provider = new TestSessionIsolationKeyProvider(TestIsolationKey);
+        var store = new IsolationKeyScopedAgentSessionStore(concreteInnerStore, provider);
+
+        // Act
+        var result = store.GetService<ConcreteAgentSessionStore>();
+
+        // Assert
+        Assert.Same(concreteInnerStore, result);
+    }
+
+    #endregion
+
     #region Helper Classes
 
     /// <summary>
@@ -386,5 +425,17 @@ public class IsolationKeyScopedAgentSessionStoreTests
 
     private sealed class TestAgentSession : AgentSession;
 
+    /// <summary>
+    /// Concrete (non-delegating) session store for testing GetService chaining.
+    /// </summary>
+    private sealed class ConcreteAgentSessionStore : AgentSessionStore
+    {
+        public override ValueTask<AgentSession> GetSessionAsync(AIAgent agent, string conversationId, CancellationToken cancellationToken = default)
+            => new(new TestAgentSession());
+
+        public override ValueTask SaveSessionAsync(AIAgent agent, string conversationId, AgentSession session, CancellationToken cancellationToken = default)
+            => ValueTask.CompletedTask;
+    }
+
     #endregion
 }