.NET: Adding Delegating AIAgent (#708)

* Delegating AIAgent impl

* Warning fix

* Fix duplicate

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

@@ -0,0 +1,70 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Shared.Diagnostics;
+
+namespace Microsoft.Extensions.AI.Agents;
+
+/// <summary>
+/// Provides an optional base class for an <see cref="AIAgent"/> that passes through calls to another instance.
+/// </summary>
+/// <remarks>
+/// This is recommended as a base type when building agents that can be chained around an underlying <see cref="AIAgent"/>.
+/// The default implementation simply passes each call to the inner agent instance.
+/// </remarks>
+public class DelegatingAIAgent : AIAgent
+{
+    /// <summary>
+    /// Initializes a new instance of the <see cref="DelegatingAIAgent"/> class.
+    /// </summary>
+    /// <param name="innerAgent">The wrapped agent instance.</param>
+    protected DelegatingAIAgent(AIAgent innerAgent)
+    {
+        this.InnerAgent = Throw.IfNull(innerAgent);
+    }
+
+    /// <summary>Gets the inner <see cref="AIAgent" />.</summary>
+    protected AIAgent InnerAgent { get; }
+
+    /// <inheritdoc />
+    public override string Id => this.InnerAgent.Id;
+
+    /// <inheritdoc />
+    public override string? Name => this.InnerAgent.Name;
+
+    /// <inheritdoc />
+    public override string? Description => this.InnerAgent.Description;
+
+    /// <inheritdoc />
+    public override object? GetService(Type serviceType, object? serviceKey = null)
+    {
+        _ = Throw.IfNull(serviceType);
+
+        // If the key is non-null, we don't know what it means so pass through to the inner service.
+        return
+            serviceKey is null && serviceType.IsInstanceOfType(this) ? this :
+            this.InnerAgent.GetService(serviceType, serviceKey);
+    }
+
+    /// <inheritdoc />
+    public override AgentThread GetNewThread() => this.InnerAgent.GetNewThread();
+
+    /// <inheritdoc />
+    public override Task<AgentRunResponse> RunAsync(
+        IReadOnlyCollection<ChatMessage> messages,
+        AgentThread? thread = null,
+        AgentRunOptions? options = null,
+        CancellationToken cancellationToken = default)
+        => this.InnerAgent.RunAsync(messages, thread, options, cancellationToken);
+
+    /// <inheritdoc />
+    public override IAsyncEnumerable<AgentRunResponseUpdate> RunStreamingAsync(
+        IReadOnlyCollection<ChatMessage> messages,
+        AgentThread? thread = null,
+        AgentRunOptions? options = null,
+        CancellationToken cancellationToken = default)
+        => this.InnerAgent.RunStreamingAsync(messages, thread, options, cancellationToken);
+}

dotnet/src/Microsoft.Extensions.AI.Agents/OpenTelemetryAgent.cs

@@ -24,13 +24,12 @@ namespace Microsoft.Extensions.AI.Agents;
 /// This class provides telemetry instrumentation for agent operations including activities, metrics, and logging.
 /// The telemetry output follows OpenTelemetry semantic conventions in <see href="https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-agent-spans/"/> and is subject to change as the conventions evolve.
 /// </remarks>
-public sealed partial class OpenTelemetryAgent : AIAgent, IDisposable
+public sealed partial class OpenTelemetryAgent : DelegatingAIAgent, IDisposable
 {
     private const LogLevel EventLogLevel = LogLevel.Information;
     private JsonSerializerOptions _jsonSerializerOptions;
     private readonly OpenTelemetryChatClient? _openTelemetryChatClient;
     private readonly string? _system;
-    private readonly AIAgent _innerAgent;
     private readonly ActivitySource _activitySource;
     private readonly Meter _meter;
     private readonly Histogram<double> _operationDurationHistogram;
@@ -44,9 +43,8 @@ public sealed partial class OpenTelemetryAgent : AIAgent, IDisposable
     /// <param name="logger">The <see cref="ILogger"/> to use for emitting events.</param>
     /// <param name="sourceName">An optional source name that will be used on the telemetry data.</param>
     public OpenTelemetryAgent(AIAgent innerAgent, ILogger? logger = null, string? sourceName = null)
+        : base(innerAgent)
     {
-        this._innerAgent = Throw.IfNull(innerAgent);
-
         string name = string.IsNullOrEmpty(sourceName) ? OpenTelemetryConsts.DefaultSourceName : sourceName!;
         this._activitySource = new(name);
         this._meter = new(name);
@@ -54,7 +52,7 @@ public sealed partial class OpenTelemetryAgent : AIAgent, IDisposable
         this._system = this.GetService<AIAgentMetadata>()?.ProviderName ?? OpenTelemetryConsts.GenAI.SystemNameValues.MicrosoftExtensionsAIAgents;
 
         // Attempt to get the open telemetry chat client if the inner agent is a ChatClientAgent.
-        this._openTelemetryChatClient = (innerAgent as ChatClientAgent)?.ChatClient.GetService<OpenTelemetryChatClient>();
+        this._openTelemetryChatClient = (this.InnerAgent as ChatClientAgent)?.ChatClient.GetService<OpenTelemetryChatClient>();
 
         // Inherit by default the EnableSensitiveData setting from the TelemetryChatClient if available.
         this.EnableSensitiveData = this._openTelemetryChatClient?.EnableSensitiveData ?? false;
@@ -112,21 +110,16 @@ public sealed partial class OpenTelemetryAgent : AIAgent, IDisposable
 
     /// <inheritdoc/>
     public override object? GetService(Type serviceType, object? serviceKey = null)
-        => base.GetService(serviceType, serviceKey)
-            ?? (serviceType == typeof(ActivitySource) ? this._activitySource
-            : this._innerAgent.GetService(serviceType, serviceKey));
+    {
+        // Handle ActivitySource requests directly - always return our own ActivitySource
+        if (serviceType == typeof(ActivitySource))
+        {
+            return this._activitySource;
+        }
 
-    /// <inheritdoc/>
-    public override string Id => this._innerAgent.Id;
-
-    /// <inheritdoc/>
-    public override string? Name => this._innerAgent.Name;
-
-    /// <inheritdoc/>
-    public override string? Description => this._innerAgent.Description;
-
-    /// <inheritdoc/>
-    public override AgentThread GetNewThread() => this._innerAgent.GetNewThread();
+        // For other service types, use the base delegation logic
+        return base.GetService(serviceType, serviceKey);
+    }
 
     /// <inheritdoc/>
     public override async Task<AgentRunResponse> RunAsync(
@@ -146,7 +139,7 @@ public sealed partial class OpenTelemetryAgent : AIAgent, IDisposable
         Exception? error = null;
         try
         {
-            response = await this._innerAgent.RunAsync(messages, thread, options, cancellationToken).ConfigureAwait(false);
+            response = await base.RunAsync(messages, thread, options, cancellationToken).ConfigureAwait(false);
             return response;
         }
         catch (Exception ex)
@@ -175,7 +168,7 @@ public sealed partial class OpenTelemetryAgent : AIAgent, IDisposable
         IAsyncEnumerable<AgentRunResponseUpdate> updates;
         try
         {
-            updates = this._innerAgent.RunStreamingAsync(messages, thread, options, cancellationToken);
+            updates = base.RunStreamingAsync(messages, thread, options, cancellationToken);
         }
         catch (Exception ex)
         {
@@ -224,7 +217,7 @@ public sealed partial class OpenTelemetryAgent : AIAgent, IDisposable
     private Activity? CreateAndConfigureActivity(string operationName, IReadOnlyCollection<ChatMessage> messages, AgentThread? thread)
     {
         // Get the GenAI system name for telemetry
-        var chatClientAgent = this._innerAgent as ChatClientAgent;
+        var chatClientAgent = this.InnerAgent as ChatClientAgent;
         Activity? activity = null;
         if (this._activitySource.HasListeners())
         {

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

@@ -0,0 +1,308 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
+using Moq;
+
+namespace Microsoft.Extensions.AI.Agents.Abstractions.UnitTests;
+
+/// <summary>
+/// Unit tests for the <see cref="DelegatingAIAgent"/> class.
+/// </summary>
+public class DelegatingAIAgentTests
+{
+    private readonly Mock<AIAgent> _innerAgentMock;
+    private readonly TestDelegatingAIAgent _delegatingAgent;
+    private readonly AgentRunResponse _testResponse;
+    private readonly List<AgentRunResponseUpdate> _testStreamingResponses;
+    private readonly AgentThread _testThread;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="DelegatingAIAgentTests"/> class.
+    /// </summary>
+    public DelegatingAIAgentTests()
+    {
+        this._innerAgentMock = new Mock<AIAgent>();
+        this._testResponse = new AgentRunResponse(new ChatMessage(ChatRole.Assistant, "Test response"));
+        this._testStreamingResponses = [new AgentRunResponseUpdate(ChatRole.Assistant, "Test streaming response")];
+        this._testThread = new AgentThread();
+
+        // Setup inner agent mock
+        this._innerAgentMock.Setup(x => x.Id).Returns("test-agent-id");
+        this._innerAgentMock.Setup(x => x.Name).Returns("Test Agent");
+        this._innerAgentMock.Setup(x => x.Description).Returns("Test Description");
+        this._innerAgentMock.Setup(x => x.GetNewThread()).Returns(this._testThread);
+
+        this._innerAgentMock
+            .Setup(x => x.RunAsync(
+                It.IsAny<IReadOnlyCollection<ChatMessage>>(),
+                It.IsAny<AgentThread?>(),
+                It.IsAny<AgentRunOptions?>(),
+                It.IsAny<CancellationToken>()))
+            .ReturnsAsync(this._testResponse);
+
+        this._innerAgentMock
+            .Setup(x => x.RunStreamingAsync(
+                It.IsAny<IReadOnlyCollection<ChatMessage>>(),
+                It.IsAny<AgentThread?>(),
+                It.IsAny<AgentRunOptions?>(),
+                It.IsAny<CancellationToken>()))
+            .Returns(ToAsyncEnumerableAsync(this._testStreamingResponses));
+
+        this._delegatingAgent = new TestDelegatingAIAgent(this._innerAgentMock.Object);
+    }
+
+    #region Constructor Tests
+
+    /// <summary>
+    /// Verify that constructor throws ArgumentNullException when innerAgent is null.
+    /// </summary>
+    [Fact]
+    public void RequiresInnerAgent()
+    {
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>("innerAgent", () => new TestDelegatingAIAgent(null!));
+    }
+
+    /// <summary>
+    /// Verify that constructor sets the inner agent correctly.
+    /// </summary>
+    [Fact]
+    public void Constructor_WithValidInnerAgent_SetsInnerAgent()
+    {
+        // Act
+        var delegatingAgent = new TestDelegatingAIAgent(this._innerAgentMock.Object);
+
+        // Assert
+        Assert.Same(this._innerAgentMock.Object, delegatingAgent.InnerAgent);
+    }
+
+    #endregion
+
+    #region Property Delegation Tests
+
+    /// <summary>
+    /// Verify that Id property delegates to inner agent.
+    /// </summary>
+    [Fact]
+    public void Id_DelegatesToInnerAgent()
+    {
+        // Act
+        var id = this._delegatingAgent.Id;
+
+        // Assert
+        Assert.Equal("test-agent-id", id);
+        this._innerAgentMock.Verify(x => x.Id, Times.Once);
+    }
+
+    /// <summary>
+    /// Verify that Name property delegates to inner agent.
+    /// </summary>
+    [Fact]
+    public void Name_DelegatesToInnerAgent()
+    {
+        // Act
+        var name = this._delegatingAgent.Name;
+
+        // Assert
+        Assert.Equal("Test Agent", name);
+        this._innerAgentMock.Verify(x => x.Name, Times.Once);
+    }
+
+    /// <summary>
+    /// Verify that Description property delegates to inner agent.
+    /// </summary>
+    [Fact]
+    public void Description_DelegatesToInnerAgent()
+    {
+        // Act
+        var description = this._delegatingAgent.Description;
+
+        // Assert
+        Assert.Equal("Test Description", description);
+        this._innerAgentMock.Verify(x => x.Description, Times.Once);
+    }
+
+    #endregion
+
+    #region Method Delegation Tests
+
+    /// <summary>
+    /// Verify that GetNewThread delegates to inner agent.
+    /// </summary>
+    [Fact]
+    public void GetNewThread_DelegatesToInnerAgent()
+    {
+        // Act
+        var thread = this._delegatingAgent.GetNewThread();
+
+        // Assert
+        Assert.Same(this._testThread, thread);
+        this._innerAgentMock.Verify(x => x.GetNewThread(), Times.Once);
+    }
+
+    /// <summary>
+    /// Verify that RunAsync delegates to inner agent with correct parameters.
+    /// </summary>
+    [Fact]
+    public async Task RunAsyncDefaultsToInnerAgentAsync()
+    {
+        // Arrange
+        var expectedMessages = new[] { new ChatMessage(ChatRole.User, "Test message") };
+        var expectedThread = new AgentThread();
+        var expectedOptions = new AgentRunOptions();
+        var expectedCancellationToken = new CancellationToken();
+        var expectedResult = new TaskCompletionSource<AgentRunResponse>();
+        var expectedResponse = new AgentRunResponse();
+
+        var innerAgentMock = new Mock<AIAgent>();
+        innerAgentMock
+            .Setup(x => x.RunAsync(expectedMessages, expectedThread, expectedOptions, expectedCancellationToken))
+            .Returns(expectedResult.Task);
+
+        var delegatingAgent = new TestDelegatingAIAgent(innerAgentMock.Object);
+
+        // Act
+        var resultTask = delegatingAgent.RunAsync(expectedMessages, expectedThread, expectedOptions, expectedCancellationToken);
+
+        // Assert
+        Assert.False(resultTask.IsCompleted);
+        expectedResult.SetResult(expectedResponse);
+        Assert.True(resultTask.IsCompleted);
+        Assert.Same(expectedResponse, await resultTask);
+    }
+
+    /// <summary>
+    /// Verify that RunStreamingAsync delegates to inner agent with correct parameters.
+    /// </summary>
+    [Fact]
+    public async Task RunStreamingAsyncDefaultsToInnerAgentAsync()
+    {
+        // Arrange
+        var expectedMessages = new[] { new ChatMessage(ChatRole.User, "Test message") };
+        var expectedThread = new AgentThread();
+        var expectedOptions = new AgentRunOptions();
+        var expectedCancellationToken = new CancellationToken();
+        AgentRunResponseUpdate[] expectedResults =
+        [
+            new(ChatRole.Assistant, "Message 1"),
+            new(ChatRole.Assistant, "Message 2")
+        ];
+
+        var innerAgentMock = new Mock<AIAgent>();
+        innerAgentMock
+            .Setup(x => x.RunStreamingAsync(expectedMessages, expectedThread, expectedOptions, expectedCancellationToken))
+            .Returns(ToAsyncEnumerableAsync(expectedResults));
+
+        var delegatingAgent = new TestDelegatingAIAgent(innerAgentMock.Object);
+
+        // Act
+        var resultAsyncEnumerable = delegatingAgent.RunStreamingAsync(expectedMessages, expectedThread, expectedOptions, expectedCancellationToken);
+
+        // Assert
+        var enumerator = resultAsyncEnumerable.GetAsyncEnumerator();
+        Assert.True(await enumerator.MoveNextAsync());
+        Assert.Same(expectedResults[0], enumerator.Current);
+        Assert.True(await enumerator.MoveNextAsync());
+        Assert.Same(expectedResults[1], enumerator.Current);
+        Assert.False(await enumerator.MoveNextAsync());
+    }
+
+    #endregion
+
+    #region GetService Tests
+
+    /// <summary>
+    /// Verify that GetService throws ArgumentNullException when serviceType is null.
+    /// </summary>
+    [Fact]
+    public void GetServiceThrowsForNullType()
+    {
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>("serviceType", () => this._delegatingAgent.GetService(null!));
+    }
+
+    /// <summary>
+    /// Verify that GetService returns the delegating agent itself when requesting compatible type and key is null.
+    /// </summary>
+    [Fact]
+    public void GetServiceReturnsSelfIfCompatibleWithRequestAndKeyIsNull()
+    {
+        // Act
+        var agent = this._delegatingAgent.GetService<DelegatingAIAgent>();
+
+        // Assert
+        Assert.Same(this._delegatingAgent, agent);
+    }
+
+    /// <summary>
+    /// Verify that GetService delegates to inner agent when service key is not null.
+    /// </summary>
+    [Fact]
+    public void GetServiceDelegatesToInnerIfKeyIsNotNull()
+    {
+        // Arrange
+        var expectedKey = new object();
+        var expectedResult = new Mock<AIAgent>().Object;
+        var innerAgentMock = new Mock<AIAgent>();
+        innerAgentMock.Setup(x => x.GetService(typeof(AIAgent), expectedKey)).Returns(expectedResult);
+        var delegatingAgent = new TestDelegatingAIAgent(innerAgentMock.Object);
+
+        // Act
+        var agent = delegatingAgent.GetService<AIAgent>(expectedKey);
+
+        // Assert
+        Assert.Same(expectedResult, agent);
+    }
+
+    /// <summary>
+    /// Verify that GetService delegates to inner agent when not compatible with request.
+    /// </summary>
+    [Fact]
+    public void GetServiceDelegatesToInnerIfNotCompatibleWithRequest()
+    {
+        // Arrange
+        var expectedResult = TimeZoneInfo.Local;
+        var expectedKey = new object();
+        var innerAgentMock = new Mock<AIAgent>();
+        innerAgentMock
+            .Setup(x => x.GetService(typeof(TimeZoneInfo), expectedKey))
+            .Returns(expectedResult);
+        var delegatingAgent = new TestDelegatingAIAgent(innerAgentMock.Object);
+
+        // Act
+        var tzi = delegatingAgent.GetService<TimeZoneInfo>(expectedKey);
+
+        // Assert
+        Assert.Same(expectedResult, tzi);
+    }
+
+    #endregion
+
+    #region Helper Methods
+
+    private static async IAsyncEnumerable<T> ToAsyncEnumerableAsync<T>(IEnumerable<T> values)
+    {
+        await Task.Yield();
+        foreach (var value in values)
+        {
+            yield return value;
+        }
+    }
+
+    #endregion
+
+    #region Test Implementation
+
+    /// <summary>
+    /// Test implementation of DelegatingAIAgent for testing purposes.
+    /// </summary>
+    private sealed class TestDelegatingAIAgent(AIAgent innerAgent) : DelegatingAIAgent(innerAgent)
+    {
+        public new AIAgent InnerAgent => base.InnerAgent;
+    }
+
+    #endregion
+}